Implement authentication request tracking and update metrics

- Added a new function to track authentication requests using Prometheus metrics in auth.py.
- Updated login, registration, and email validation endpoints to call the tracking function.
- Refactored Prometheus middleware to handle error responses and record metrics accurately.
- Modified Grafana dashboard configuration to reflect changes in authentication metrics tracking.

Author: PhiRho

app/api/auth.py

@@ -36,6 +36,38 @@
 from app.services.dynamodb import DynamoDBService
 from app.services.ses import SESService
 from app.api.teams import register_team
+from app.middleware.prometheus import auth_requests_total
+
+def track_auth_request(request: Request, identifier: Optional[str] = None) -> None:
+    """
+    Track authentication requests using Prometheus metrics.
+
+    Args:
+        request: The FastAPI request object
+        identifier: Optional identifier (username/email) for the request
+    """
+    # Get the path from the request
+    path = request.url.path
+
+    # If no identifier provided, try to extract from request
+    if not identifier:
+        try:
+            # Try to get from form data
+            form_data = request.form()
+            identifier = form_data.get("username") or form_data.get("email")
+        except:
+            try:
+                # Try to get from JSON body
+                body = request.json()
+                identifier = body.get("username") or body.get("email")
+            except:
+                identifier = "unknown"
+
+    # Increment the counter
+    auth_requests_total.labels(
+        endpoint=path,
+        identifier=identifier or "unknown"
+    ).inc()
 
 router = APIRouter(
     tags=["Authentication"]
@@ -156,6 +188,9 @@ async def login(
             detail="Invalid login data. Please provide username and password in either form data or JSON format."
         )
 
+    # Track the auth request
+    track_auth_request(request, login_data.username)
+
     user = db.query(DBUser).filter(DBUser.email == login_data.username).first()
     if not user or not verify_password(login_data.password, user.hashed_password):
         raise HTTPException(
@@ -243,7 +278,11 @@ async def update_user_me(
     return current_user
 
 @router.post("/register", response_model=User)
-async def register(user: UserCreate, db: Session = Depends(get_db)):
+async def register(
+    request: Request,
+    user: UserCreate,
+    db: Session = Depends(get_db)
+):
     """
     Register a new user account.
 
@@ -252,6 +291,9 @@ async def register(user: UserCreate, db: Session = Depends(get_db)):
 
     After registration, you'll need to login to get an access token.
     """
+    # Track the auth request
+    track_auth_request(request, user.email)
+
     # Check if user with this email exists
     db_user = db.query(DBUser).filter(DBUser.email == user.email).first()
     if db_user:
@@ -309,6 +351,9 @@ async def validate_email(
             detail="Email is required"
         )
 
+    # Track the auth request
+    track_auth_request(request, email)
+
     try:
         email_validator.validate_email(email, check_deliverability=False)
     except EmailNotValidError as e:
@@ -442,6 +487,9 @@ async def sign_in(
             detail="Invalid sign in data. Please provide username and verification code in either form data or JSON format."
         )
 
+    # Track the auth request
+    track_auth_request(request, sign_in_data.username)
+
     # Verify the code using DynamoDB first
     dynamodb_service = DynamoDBService()
     stored_code = dynamodb_service.read_validation_code(sign_in_data.username)

app/main.py

@@ -1,6 +1,4 @@
-from fastapi import FastAPI, Depends, HTTPException
-from fastapi.security import OAuth2PasswordBearer
-from sqlalchemy.orm import Session
+from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.trustedhost import TrustedHostMiddleware
 from starlette.middleware.base import BaseHTTPMiddleware

app/middleware/prometheus.py

@@ -1,7 +1,7 @@
 from prometheus_client import Counter, Histogram, Gauge
 from prometheus_fastapi_instrumentator import Instrumentator, metrics
 from prometheus_fastapi_instrumentator.metrics import Info
-from fastapi import Request
+from fastapi import Request, HTTPException, status
 from starlette.middleware.base import BaseHTTPMiddleware
 from app.core.security import get_current_user_from_auth
 from app.db.database import get_db
@@ -45,25 +45,47 @@
     ["user_id", "method", "endpoint"]
 )
 
+# Auth Metrics
+auth_requests_total = Counter(
+    "auth_requests_total",
+    "Total number of authentication requests",
+    ["endpoint", "identifier"]
+)
+
 class PrometheusMiddleware(BaseHTTPMiddleware):
     async def dispatch(self, request: Request, call_next):
         # Skip metrics for certain paths
         if request.url.path in ["/metrics", "/health", "/docs", "/openapi.json"]:
             return await call_next(request)
 
         start_time = time.time()
+        response = None
+        duration = 0
+        is_error = False
 
-        # Process the request
-        response = await call_next(request)
-
-        # Calculate duration
-        duration = time.time() - start_time
+        try:
+            # Process the request
+            response = await call_next(request)
+            duration = time.time() - start_time
+        except Exception as e:
+            logger.warning(f"Request failed: {e}")
+            # Record the actual duration of the failed request
+            duration = time.time() - start_time
+            # Capture the error response to be raised later
+            is_error = True
+            if not isinstance(e, HTTPException):
+                error_response = HTTPException(
+                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                    detail=str(e)
+                )
+            else:
+                error_response = e
 
         # Record RED metrics
         http_requests_total.labels(
             method=request.method,
             endpoint=request.url.path,
-            status_code=response.status_code
+            status_code=response.status_code if response else 500
         ).inc()
 
         http_request_duration_seconds.labels(
@@ -86,14 +108,19 @@ async def dispatch(self, request: Request, call_next):
                     access_token = parts[1]
 
             if access_token:
-                db = next(get_db())
-                user = await get_current_user_from_auth(
-                    access_token=access_token if access_token else None,
-                    authorization=auth_header if auth_header else None,
-                    db=db
-                )
-                user_id = str(user.id) if user else "anonymous"
-                db.close()
+                try:
+                    db = next(get_db())
+                    user = await get_current_user_from_auth(
+                        access_token=access_token if access_token else None,
+                        authorization=auth_header if auth_header else None,
+                        db=db
+                    )
+                    user_id = str(user.id) if user else "anonymous"
+                except Exception as e:
+                    logger.debug(f"Could not get user for metrics: {str(e)}")
+                    user_id = "anonymous"
+                finally:
+                    db.close()
         except Exception as e:
             logger.debug(f"Could not get user for metrics: {str(e)}")
             user_id = "anonymous"
@@ -105,4 +132,7 @@ async def dispatch(self, request: Request, call_next):
             endpoint=request.url.path
         ).inc()
 
-        return response
+        if is_error:
+            raise error_response
+        else:
+            return response

grafana/provisioning/dashboards/fastapi.json

@@ -621,8 +621,8 @@
             "type": "prometheus",
             "uid": "prometheus"
           },
-          "expr": "rate(auth_attempts_total[5m])",
-          "legendFormat": "{{action}} - {{endpoint}} ({{status}})",
+          "expr": "rate(auth_requests_total[5m])",
+          "legendFormat": "{{endpoint}} - {{identifier}} ({{status_code}})",
           "refId": "A"
         }
       ],
@@ -643,8 +643,8 @@
   },
   "timepicker": {},
   "timezone": "",
-  "title": "FastAPI Dashboard",
-  "uid": "fastapi",
+  "title": "amazee.ai Backend",
+  "uid": "amazeeai-backend",
   "version": 1,
   "weekStart": ""
 }