Fix auth module

shared/monitoring/metrics.py

@@ -1,12 +1,14 @@
+# shared/monitoring/metrics.py
 """
 Metrics collection for microservices
 """
 
 import time
 import logging
-from typing import Dict, Any
+from typing import Dict, Any, List # Added List import
 from prometheus_client import Counter, Histogram, Gauge, start_http_server
 from functools import wraps
+from prometheus_client import generate_latest # Moved this import here for consistency
 
 logger = logging.getLogger(__name__)
 
@@ -43,11 +45,14 @@ FORECASTS_GENERATED = Counter(
 
 class MetricsCollector:
     """Metrics collector for microservices"""
-    
+
     def __init__(self, service_name: str):
         self.service_name = service_name
         self.start_time = time.time()
-    
+        # Initialize dictionaries to hold custom counters and histograms
+        self._counters: Dict[str, Counter] = {}
+        self._histograms: Dict[str, Histogram] = {}
+
     def start_metrics_server(self, port: int = 8080):
         """Start Prometheus metrics server"""
         try:
@@ -55,7 +60,7 @@ class MetricsCollector:
             logger.info(f"Metrics server started on port {port}")
         except Exception as e:
             logger.error(f"Failed to start metrics server: {e}")
-    
+
     def record_request(self, method: str, endpoint: str, status_code: int, duration: float):
         """Record HTTP request metrics"""
         REQUEST_COUNT.labels(
@@ -64,49 +69,110 @@ class MetricsCollector:
             status_code=status_code,
             service=self.service_name
         ).inc()
-        
+
         REQUEST_DURATION.labels(
             method=method,
             endpoint=endpoint,
             service=self.service_name
         ).observe(duration)
-    
+
     def record_training_job(self, status: str):
         """Record training job metrics"""
         TRAINING_JOBS.labels(
             status=status,
             service=self.service_name
         ).inc()
-    
+
     def record_forecast_generated(self):
         """Record forecast generation metrics"""
         FORECASTS_GENERATED.labels(
             service=self.service_name
         ).inc()
-    
+
     def set_active_connections(self, count: int):
         """Set active database connections"""
         ACTIVE_CONNECTIONS.labels(
             service=self.service_name
         ).set(count)
 
+    def register_counter(self, name: str, documentation: str, labels: List[str] = None):
+        """Register a custom Counter metric."""
+        if name not in self._counters:
+            if labels is None:
+                labels = ['service']
+            elif 'service' not in labels:
+                labels.append('service')
+            # Pass labelnames as a keyword argument
+            self._counters[name] = Counter(name, documentation, labelnames=labels)
+            logger.info(f"Registered counter: {name}")
+        else:
+            logger.warning(f"Counter '{name}' already registered.")
+        return self._counters[name] # Return the counter for direct use if needed
+
+    def increment_counter(self, name: str, value: int = 1, labels: Dict[str, str] = None):
+        """Increment a custom Counter metric."""
+        if name not in self._counters:
+            logger.error(f"Counter '{name}' not registered. Cannot increment.")
+            return
+
+        # Ensure the 'service' label is always present
+        if labels is None:
+            labels = {'service': self.service_name}
+        elif 'service' not in labels:
+            labels['service'] = self.service_name
+
+        self._counters[name].labels(**labels).inc(value)
+
+    def register_histogram(self, name: str, documentation: str, labels: List[str] = None, buckets: tuple = Histogram.DEFAULT_BUCKETS):
+        """Register a custom Histogram metric."""
+        if name not in self._histograms:
+            if labels is None:
+                labels = ['service']
+            elif 'service' not in labels:
+                labels.append('service')
+            # Pass labelnames and buckets as keyword arguments
+            self._histograms[name] = Histogram(name, documentation, labelnames=labels, buckets=buckets)
+            logger.info(f"Registered histogram: {name}")
+        else:
+            logger.warning(f"Histogram '{name}' already registered.")
+        return self._histograms[name] # Return the histogram for direct use if needed
+
+    def observe_histogram(self, name: str, value: float, labels: Dict[str, str] = None):
+        """Observe a custom Histogram metric."""
+        if name not in self._histograms:
+            logger.error(f"Histogram '{name}' not registered. Cannot observe.")
+            return
+
+        if labels is None:
+            labels = {'service': self.service_name}
+        elif 'service' not in labels:
+            labels['service'] = self.service_name
+
+        self._histograms[name].labels(**labels).observe(value)
+
+    def get_metrics(self) -> str:
+        """Return Prometheus metrics in exposition format."""
+        return generate_latest().decode('utf-8')
+
+
 def metrics_middleware(metrics_collector: MetricsCollector):
     """Middleware to collect metrics"""
-    
-    def middleware(request, call_next):
+
+    async def middleware(request, call_next):
         start_time = time.time()
-        
-        response = call_next(request)
-        
+
+        response = await call_next(request)
+
         duration = time.time() - start_time
-        
+
+        # Use the specific record_request for HTTP requests
         metrics_collector.record_request(
             method=request.method,
             endpoint=request.url.path,
             status_code=response.status_code,
             duration=duration
         )
-        
+
         return response
-    
+
     return middleware