Merge pull request #59 from sensein/user_management

Improve performance and make JWT auth non-blocking

Author: tekrajchhetri

docker-compose.unified.yml

@@ -14,6 +14,7 @@ services:
     networks:
       - brainkb-network
     restart: unless-stopped
+    command: postgres -c max_connections=200 -c shared_buffers=256MB
     healthcheck:
       test: ["CMD-SHELL", "pg_isready -U ${JWT_POSTGRES_DATABASE_USER:-postgres}"]
       interval: 10s
@@ -66,6 +67,8 @@ services:
     volumes:
       # Use named volume for local development, or custom path via OXIGRAPH_DATA_PATH
       # For production with shared storage, set OXIGRAPH_DATA_PATH=/fsx/brainkb-kg-repo in .env
+      # IMPORTANT: If using bind mount (OXIGRAPH_DATA_PATH set), ensure the path exists and has correct permissions
+      # The start_services.sh script will validate and create the path if needed
       - ${OXIGRAPH_DATA_PATH:-oxigraph_data}:/data
       - ${OXIGRAPH_TMP_PATH:-oxigraph_tmp}:/tmp
 

ml_service/core/database.py

@@ -20,7 +20,7 @@
 import asyncpg
 from contextlib import asynccontextmanager
 from typing import AsyncGenerator, Optional
-from fastapi import HTTPException
+from fastapi import HTTPException, status
 
 from core.configuration import load_environment
 
@@ -66,16 +66,24 @@ async def init_db_pool():
         num_workers = int(os.getenv("WEB_CONCURRENCY", "1" if env_state == "development" else "4"))
 
         # Calculate per-worker sizes to avoid exceeding PostgreSQL limits
-        # Increased pool size to handle concurrent JWT token requests and other operations
-        # Target: ~90 total connections across all workers (leaves room for other services)
+        # Increased pool size to handle high concurrent JWT token requests
+        # After optimization: each token request uses 1 connection (was 2)
+        # Target: support 100+ concurrent requests safely
         if env_state == "development":
             min_size = 0  # Lazy init - no connections at startup
             max_size = 10
         else:
             min_size = 2  # Keep a few connections ready for fast JWT validation
-            # With 6 workers: 15 per worker = 90 total (safe, PostgreSQL default max is 100)
-            # Increased from 8 to 15 to handle bursts of concurrent requests
-            max_size = max(10, min(15, 90 // num_workers))
+            # Calculate pool size to support high concurrency
+            # Target: 100+ concurrent requests across all workers
+            # Formula: ensure at least 20-30 connections per worker, but cap at reasonable limit
+            # Allow environment override via DB_POOL_MAX_SIZE
+            if os.getenv("DB_POOL_MAX_SIZE"):
+                max_size = int(os.getenv("DB_POOL_MAX_SIZE"))
+            else:
+                # Default: 25 per worker (supports 100 concurrent with 4 workers, 150 with 6)
+                # PostgreSQL default max_connections is usually 100, but can be increased
+                max_size = 25
 
         print(f"[DB POOL INIT] Workers: {num_workers}, min={min_size}, max={max_size} per worker")
         print(f"[DB POOL INIT] Total potential: {max_size * num_workers} connections")
@@ -280,6 +288,13 @@ async def _insert_logic(connection):
             # Manage our own connection
             async with get_db_connection() as connection:
                 return await _insert_logic(connection)
+    except asyncpg.exceptions.UniqueViolationError as e:
+        # Handle unique constraint violations (e.g., duplicate email) atomically
+        logger.warning(f"Registration attempt with duplicate email {email}: {str(e)}")
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="A user with that email already exists"
+        )
     except Exception as e:
         logger.error(f"Error inserting data for user {email}: {str(e)}")
         raise HTTPException(status_code=400, detail=str(e))
@@ -331,10 +346,12 @@ async def select_scope_id(conn: Optional[asyncpg.Connection] = None) -> Optional
         raise HTTPException(status_code=400, detail=str(e))
 
 
-async def get_scopes_by_user(user_id: int):
+async def get_scopes_by_user(user_id: int, conn: Optional[asyncpg.Connection] = None):
     """
     Get all scope names assigned to a user.
     Returns a list of scope names.
+    If conn is provided, uses that connection (caller manages it).
+    Otherwise, manages its own connection.
     """
     query = f"""
     SELECT s.name
@@ -344,11 +361,19 @@ async def get_scopes_by_user(user_id: int):
     """
 
     try:
-        async with get_db_connection() as conn:
+        if conn is not None:
+            # Use provided connection
             results = await conn.fetch(query, user_id)
             assigned_scopes_to_user = [result["name"] for result in results]
             logger.debug(f"Scopes for user {user_id}: {assigned_scopes_to_user}")
             return assigned_scopes_to_user
+        else:
+            # Manage our own connection
+            async with get_db_connection() as connection:
+                results = await connection.fetch(query, user_id)
+                assigned_scopes_to_user = [result["name"] for result in results]
+                logger.debug(f"Scopes for user {user_id}: {assigned_scopes_to_user}")
+                return assigned_scopes_to_user
     except Exception as e:
         logger.error(f"Error getting scopes for user {user_id}: {str(e)}")
         raise HTTPException(status_code=400, detail=str(e))

ml_service/core/routers/jwt_auth.py

@@ -2,7 +2,7 @@
 
 from fastapi import APIRouter, HTTPException, status, Depends
 
-from core.database import connect_postgres, get_user, insert_data, get_scopes_by_user
+from core.database import get_db_connection, insert_data, get_scopes_by_user
 from core.models.user import UserIn, LoginUserIn
 from core.security import get_password_hash, authenticate_user, create_access_token
 
@@ -12,23 +12,32 @@
 
 
 @router.post("/register", status_code=201)
-async def register(user: UserIn, conn=Depends(connect_postgres)):
-
-    if await get_user(conn=conn, email=user.email):
-        raise HTTPException(
-            status_code=status.HTTP_400_BAD_REQUEST,
-            detail="A user with that email already exists",
+async def register(user: UserIn):
+    """
+    Register a new user. Uses proper connection management to avoid connection leaks.
+    
+    Note: The unique constraint check is handled atomically by the database.
+    This prevents race conditions where two concurrent requests could both pass
+    a pre-check and then both attempt to insert the same email.
+    """
+    async with get_db_connection() as conn:
+        hashed_password = await get_password_hash(user.password)
+        
+        # Let the database enforce uniqueness atomically - no pre-check needed
+        # insert_data will catch UniqueViolationError and return a user-friendly message
+        return await insert_data(
+            conn=conn, fullname=user.full_name, email=user.email, password=hashed_password
         )
-    hashed_password = get_password_hash(user.password)
-
-    return await insert_data(
-        conn=conn, fullname=user.full_name, email=user.email, password=hashed_password
-    )
 
 
 @router.post("/token")
-async def login(user: LoginUserIn, conn=Depends(connect_postgres)):
-    user = await authenticate_user(user.email, user.password, conn)
-    scopes = await get_scopes_by_user(user_id=user["id"])
-    access_token = create_access_token(user["email"], scopes)
-    return {"access_token": access_token, "token_type": "bearer"}
+async def login(user: LoginUserIn):
+    """
+    Authenticate user and return JWT token. Uses proper connection management to avoid connection leaks.
+    Reuses the same database connection for both authentication and scope retrieval to minimize connection pool usage.
+    """
+    async with get_db_connection() as conn:
+        authenticated_user = await authenticate_user(user.email, user.password, conn)
+        scopes = await get_scopes_by_user(user_id=authenticated_user["id"], conn=conn)
+        access_token = create_access_token(authenticated_user["email"], scopes)
+        return {"access_token": access_token, "token_type": "bearer"}

ml_service/core/security.py

@@ -18,6 +18,7 @@
 
 import datetime
 import logging
+import asyncio
 from typing import Annotated, List, Optional, Dict
 
 from fastapi import Depends, HTTPException, status, WebSocket
@@ -56,20 +57,22 @@ def create_access_token(email: str, scopes: List[str]) -> str:
     return encoded_jwt
 
 
-def get_password_hash(password: str) -> str:
-    return pwd_context.hash(password)
+async def get_password_hash(password: str) -> str:
+    """Hash password asynchronously to avoid blocking the event loop."""
+    return await asyncio.to_thread(pwd_context.hash, password)
 
 
-def verify_password(plain_password: str, hashed_password: str) -> bool:
-    return pwd_context.verify(plain_password, hashed_password)
+async def verify_password(plain_password: str, hashed_password: str) -> bool:
+    """Verify password asynchronously to avoid blocking the event loop."""
+    return await asyncio.to_thread(pwd_context.verify, plain_password, hashed_password)
 
 
 async def authenticate_user(email, password, conn):
     logger.debug("Authenticating user", extra={"email": email})
     user = await get_user(conn=conn, email=email)
     if not user:
         raise credentials_exception
-    if not verify_password(password, user["password"]):
+    if not await verify_password(password, user["password"]):
         raise credentials_exception
     return user