refactor: enhance Redis connection management and backup functionality

atyrode · atyrode · commit e69cb4e6b744 · 2025-05-02T20:00:01.000Z
- Introduced a Redis connection pool to optimize Redis client management, improving performance and resource utilization.
- Updated session management functions to utilize the new Redis connection pool, ensuring efficient access to session data.
- Implemented a new method in BackupService to retrieve backups for a user's first pad using a join operation, addressing the N+1 query problem.
- Enhanced pad router to create backups conditionally based on time intervals, improving backup efficiency and management.
- Streamlined user router to utilize the new Redis client retrieval method, promoting better code organization and maintainability.
diff --git a/src/backend/config.py b/src/backend/config.py
@@ -3,6 +3,7 @@
 import time
 import httpx
 import redis
+from redis import ConnectionPool, Redis
 import jwt
 from jwt.jwks_client import PyJWKClient
 from typing import Optional, Dict, Any, Tuple
@@ -16,6 +17,11 @@
 ASSETS_DIR = os.getenv("ASSETS_DIR")
 FRONTEND_URL = os.getenv('FRONTEND_URL')
 
+MAX_BACKUPS_PER_USER = 10  # Maximum number of backups to keep per user
+MIN_INTERVAL_MINUTES = 5  # Minimum interval in minutes between backups
+DEFAULT_PAD_NAME = "Untitled"  # Default name for new pads
+DEFAULT_TEMPLATE_NAME = "default" # Template name to use when a user doesn't have a pad
+
 # ===== PostHog Configuration =====
 POSTHOG_API_KEY = os.getenv("VITE_PUBLIC_POSTHOG_KEY")
 POSTHOG_HOST = os.getenv("VITE_PUBLIC_POSTHOG_HOST")
@@ -32,15 +38,26 @@
 REDIS_PASSWORD = os.getenv('REDIS_PASSWORD', None)
 REDIS_PORT = int(os.getenv('REDIS_PORT', 6379))
 
-# Redis connection
-redis_client = redis.Redis(
+# Create a Redis connection pool
+redis_pool = ConnectionPool(
     host=REDIS_HOST,
     password=REDIS_PASSWORD,
     port=REDIS_PORT,
     db=0,
-    decode_responses=True
+    decode_responses=True,
+    max_connections=10,  # Adjust based on your application's needs
+    socket_timeout=5.0,
+    socket_connect_timeout=1.0,
+    health_check_interval=30
 )
 
+# Create a Redis client that uses the connection pool
+redis_client = Redis(connection_pool=redis_pool)
+
+def get_redis_client():
+    """Get a Redis client from the connection pool"""
+    return Redis(connection_pool=redis_pool)
+
 # ===== Coder API Configuration =====
 CODER_API_KEY = os.getenv("CODER_API_KEY")
 CODER_URL = os.getenv("CODER_URL")
@@ -54,22 +71,25 @@
 # Session management functions
 def get_session(session_id: str) -> Optional[Dict[str, Any]]:
     """Get session data from Redis"""
-    session_data = redis_client.get(f"session:{session_id}")
+    client = get_redis_client()
+    session_data = client.get(f"session:{session_id}")
     if session_data:
         return json.loads(session_data)
     return None
 
 def set_session(session_id: str, data: Dict[str, Any], expiry: int) -> None:
     """Store session data in Redis with expiry in seconds"""
-    redis_client.setex(
+    client = get_redis_client()
+    client.setex(
         f"session:{session_id}", 
         expiry,
         json.dumps(data)
     )
 
 def delete_session(session_id: str) -> None:
     """Delete session data from Redis"""
-    redis_client.delete(f"session:{session_id}")
+    client = get_redis_client()
+    client.delete(f"session:{session_id}")
 
 def get_auth_url() -> str:
     """Generate the authentication URL for Keycloak login"""
@@ -161,4 +181,4 @@ def get_jwks_client():
     if _jwks_client is None:
         jwks_url = f"{OIDC_SERVER_URL}/realms/{OIDC_REALM}/protocol/openid-connect/certs"
         _jwks_client = PyJWKClient(jwks_url)
-    return _jwks_client
+    return _jwks_client
diff --git a/src/backend/database/repository/backup_repository.py b/src/backend/database/repository/backup_repository.py
@@ -8,9 +8,9 @@
 
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.future import select
-from sqlalchemy import delete, func
+from sqlalchemy import delete, func, join
 
-from ..models import BackupModel
+from ..models import BackupModel, PadModel
 
 class BackupRepository:
     """Repository for backup-related database operations"""
@@ -89,3 +89,28 @@ async def count_by_source(self, source_id: UUID) -> int:
         stmt = select(func.count()).select_from(BackupModel).where(BackupModel.source_id == source_id)
         result = await self.session.execute(stmt)
         return result.scalar()
+        
+    async def get_backups_by_user(self, user_id: UUID, limit: int = 10) -> List[BackupModel]:
+        """
+        Get backups for a user's first pad directly using a join operation.
+        This eliminates the N+1 query problem by fetching the pad and its backups in a single query.
+        
+        Args:
+            user_id: The user ID to get backups for
+            limit: Maximum number of backups to return
+            
+        Returns:
+            List of backup models
+        """
+        # Create a join between PadModel and BackupModel
+        stmt = select(BackupModel).join(
+            PadModel, 
+            BackupModel.source_id == PadModel.id
+        ).where(
+            PadModel.owner_id == user_id
+        ).order_by(
+            BackupModel.created_at.desc()
+        ).limit(limit)
+        
+        result = await self.session.execute(stmt)
+        return result.scalars().all()
diff --git a/src/backend/database/service/backup_service.py b/src/backend/database/service/backup_service.py
@@ -8,7 +8,7 @@
 
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from ..repository import BackupRepository, PadRepository
+from ..repository import BackupRepository, PadRepository, UserRepository
 
 class BackupService:
     """Service for backup-related business logic"""
@@ -101,3 +101,76 @@ async def manage_backups(self, source_id: UUID, max_backups: int = 10) -> int:
             return await self.repository.delete_older_than(source_id, max_backups)
         
         return 0  # No backups deleted
+        
+    async def get_backups_by_user(self, user_id: UUID, limit: int = 10) -> List[Dict[str, Any]]:
+        """
+        Get backups for a user's first pad directly using a join operation.
+        This eliminates the N+1 query problem by fetching the pad and its backups in a single query.
+        
+        Args:
+            user_id: The user ID to get backups for
+            limit: Maximum number of backups to return
+            
+        Returns:
+            List of backup dictionaries
+        """
+        # Check if user exists
+        user_repository = UserRepository(self.session)
+        user = await user_repository.get_by_id(user_id)
+        if not user:
+            raise ValueError(f"User with ID '{user_id}' does not exist")
+        
+        # Get backups directly with a single query
+        backups = await self.repository.get_backups_by_user(user_id, limit)
+        return [backup.to_dict() for backup in backups]
+        
+    async def create_backup_if_needed(self, source_id: UUID, data: Dict[str, Any], 
+                                     min_interval_minutes: int = 5, 
+                                     max_backups: int = 10) -> Optional[Dict[str, Any]]:
+        """
+        Create a backup only if needed:
+        - If there are no existing backups
+        - If the latest backup is older than the specified interval
+        
+        Args:
+            source_id: The ID of the source pad
+            data: The data to backup
+            min_interval_minutes: Minimum time between backups in minutes
+            max_backups: Maximum number of backups to keep
+            
+        Returns:
+            The created backup dict if a backup was created, None otherwise
+        """
+        # Check if source pad exists
+        source_pad = await self.pad_repository.get_by_id(source_id)
+        if not source_pad:
+            raise ValueError(f"Pad with ID '{source_id}' does not exist")
+        
+        # Get the latest backup
+        latest_backup = await self.repository.get_latest_by_source(source_id)
+        
+        # Calculate the current time
+        current_time = datetime.now()
+        
+        # Determine if we need to create a backup
+        create_backup = False
+        
+        if not latest_backup:
+            # No backups exist yet, so create one
+            create_backup = True
+        else:
+            # Check if the latest backup is older than the minimum interval
+            backup_age = current_time - latest_backup.created_at
+            if backup_age.total_seconds() > (min_interval_minutes * 60):
+                create_backup = True
+        
+        # Create a backup if needed
+        if create_backup:
+            backup = await self.repository.create(source_id, data)
+            
+            # Manage backups (clean up old ones)
+            await self.manage_backups(source_id, max_backups)
+            
+            return backup.to_dict()
+        
+        return None
diff --git a/src/backend/main.py b/src/backend/main.py
@@ -10,7 +10,7 @@
 from fastapi.staticfiles import StaticFiles
 
 from database import init_db
-from config import STATIC_DIR, ASSETS_DIR, POSTHOG_API_KEY, POSTHOG_HOST
+from config import STATIC_DIR, ASSETS_DIR, POSTHOG_API_KEY, POSTHOG_HOST, redis_client, redis_pool
 from dependencies import UserSession, optional_auth
 from routers.auth_router import auth_router
 from routers.user_router import user_router
@@ -76,14 +76,29 @@ async def load_templates():
 
 @asynccontextmanager
 async def lifespan(_: FastAPI):
+    # Initialize database
     await init_db()
     print("Database connection established successfully")
     
+    # Check Redis connection
+    try:
+        redis_client.ping()
+        print("Redis connection established successfully")
+    except Exception as e:
+        print(f"Warning: Redis connection failed: {str(e)}")
+    
     # Load all templates from the templates directory
     await load_templates()
     print("Templates loaded successfully")
     
     yield
+    
+    # Clean up connections when shutting down
+    try:
+        redis_pool.disconnect()
+        print("Redis connections closed")
+    except Exception as e:
+        print(f"Error closing Redis connections: {str(e)}")
 
 app = FastAPI(lifespan=lifespan)
 
diff --git a/src/backend/routers/pad_router.py b/src/backend/routers/pad_router.py
@@ -6,12 +6,7 @@
 from dependencies import UserSession, require_auth
 from database import get_pad_service, get_backup_service, get_template_pad_service
 from database.service import PadService, BackupService, TemplatePadService
-
-# Constants
-MAX_BACKUPS_PER_USER = 10  # Maximum number of backups to keep per user
-DEFAULT_PAD_NAME = "Untitled"  # Default name for new pads
-DEFAULT_TEMPLATE_NAME = "default" # Template name to use when a user doesn't have a pad
-
+from config import MAX_BACKUPS_PER_USER, MIN_INTERVAL_MINUTES, DEFAULT_PAD_NAME, DEFAULT_TEMPLATE_NAME
 pad_router = APIRouter()
 
 
@@ -39,11 +34,13 @@ async def save_canvas(
             pad = user_pads[0]  # Use the first pad (assuming one pad per user for now)
             await pad_service.update_pad_data(pad["id"], data)
             
-        # Create a backup
-        await backup_service.create_backup(pad["id"], data)
-        
-        # Manage backups (keep only the most recent ones)
-        await backup_service.manage_backups(pad["id"], MAX_BACKUPS_PER_USER)
+        # Create a backup only if needed (if none exist or latest is > 5 min old)
+        await backup_service.create_backup_if_needed(
+            source_id=pad["id"], 
+            data=data,
+            min_interval_minutes=MIN_INTERVAL_MINUTES,
+            max_backups=MAX_BACKUPS_PER_USER
+        )
         
         return {"status": "success"}
     except Exception as e:
@@ -54,7 +51,8 @@ async def save_canvas(
 async def get_canvas(
     user: UserSession = Depends(require_auth),
     pad_service: PadService = Depends(get_pad_service),
-    template_pad_service: TemplatePadService = Depends(get_template_pad_service)
+    template_pad_service: TemplatePadService = Depends(get_template_pad_service),
+    backup_service: BackupService = Depends(get_backup_service)
 ):
     """Get canvas data for the authenticated user"""
     try:
@@ -63,7 +61,14 @@ async def get_canvas(
         
         if not user_pads:
             # Return default canvas if user doesn't have a pad
-            return await create_pad_from_template(DEFAULT_TEMPLATE_NAME, DEFAULT_PAD_NAME, user, pad_service, template_pad_service)
+            return await create_pad_from_template(
+                name=DEFAULT_TEMPLATE_NAME, 
+                display_name=DEFAULT_PAD_NAME, 
+                user=user, 
+                pad_service=pad_service, 
+                template_pad_service=template_pad_service,
+                backup_service=backup_service
+            )
         
         # Return the first pad's data (assuming one pad per user for now)
         return user_pads[0]["data"]
@@ -77,7 +82,8 @@ async def create_pad_from_template(
     display_name: str = DEFAULT_PAD_NAME,
     user: UserSession = Depends(require_auth),
     pad_service: PadService = Depends(get_pad_service),
-    template_pad_service: TemplatePadService = Depends(get_template_pad_service)
+    template_pad_service: TemplatePadService = Depends(get_template_pad_service),
+    backup_service: BackupService = Depends(get_backup_service)
 ):
     """Create a new pad from a template"""
 
@@ -94,6 +100,14 @@ async def create_pad_from_template(
             data=template["data"]
         )
         
+        # Create an initial backup for the new pad
+        await backup_service.create_backup_if_needed(
+            source_id=pad["id"],
+            data=template["data"],
+            min_interval_minutes=0,  # Always create initial backup
+            max_backups=MAX_BACKUPS_PER_USER
+        )
+        
         return pad
     except ValueError as e:
         raise HTTPException(status_code=400, detail=str(e))
@@ -105,7 +119,6 @@ async def create_pad_from_template(
 async def get_recent_canvas_backups(
     limit: int = MAX_BACKUPS_PER_USER, 
     user: UserSession = Depends(require_auth),
-    pad_service: PadService = Depends(get_pad_service),
     backup_service: BackupService = Depends(get_backup_service)
 ):
     """Get the most recent canvas backups for the authenticated user"""
@@ -114,18 +127,12 @@ async def get_recent_canvas_backups(
         limit = MAX_BACKUPS_PER_USER
     
     try:
-        # Get user's pads
-        user_pads = await pad_service.get_pads_by_owner(user.id)
-        
-        # Get the first pad's ID (assuming one pad per user for now)
-        pad_id = UUID(user_pads[0]["id"])
-        
-        # Get backups for the pad
-        backups_data = await backup_service.get_backups_by_source(pad_id)
+        # Get backups directly with a single query
+        backups_data = await backup_service.get_backups_by_user(user.id, limit)
         
         # Format backups to match the expected response format
         backups = []
-        for backup in backups_data[:limit]:
+        for backup in backups_data:
             backups.append({
                 "id": backup["id"],
                 "timestamp": backup["created_at"],
diff --git a/src/backend/routers/user_router.py b/src/backend/routers/user_router.py
@@ -4,7 +4,7 @@
 import posthog
 from fastapi import APIRouter, Depends, HTTPException
 
-from config import redis_client, FRONTEND_URL
+from config import get_redis_client, FRONTEND_URL
 from database import get_user_service
 from database.service import UserService
 from dependencies import UserSession, require_admin, require_auth
@@ -92,7 +92,8 @@ async def get_user_count(
     _: bool = Depends(require_admin),
 ):
     """Get the number of active sessions (admin only)"""
-    session_count = len(redis_client.keys("session:*"))
+    client = get_redis_client()
+    session_count = len(client.keys("session:*"))
     return {"active_sessions": session_count }
 
 
