mcp-code-indexer 1.9.1__py3-none-any.whl → 2.0.1__py3-none-any.whl

mcp_code_indexer/database/database.py

@@ -13,12 +13,20 @@ from datetime import datetime, timedelta
 from pathlib import Path
 from typing import List, Optional, Dict, Any, Tuple, AsyncIterator
 
+import asyncio
+import random
 import aiosqlite
 
 from mcp_code_indexer.database.models import (
     Project, FileDescription, MergeConflict, SearchResult,
     CodebaseSizeInfo, ProjectOverview, WordFrequencyResult, WordFrequencyTerm
 )
+from mcp_code_indexer.database.retry_handler import (
+    RetryHandler, ConnectionRecoveryManager, create_retry_handler
+)
+from mcp_code_indexer.database.connection_health import (
+    ConnectionHealthMonitor, DatabaseMetricsCollector
+)
 
 logger = logging.getLogger(__name__)
 
@@ -31,19 +39,52 @@ class DatabaseManager:
     and caching with proper transaction management and error handling.
     """
     
-    def __init__(self, db_path: Path, pool_size: int = 5):
+    def __init__(self, 
+                 db_path: Path, 
+                 pool_size: int = 3,
+                 retry_count: int = 5,
+                 timeout: float = 10.0,
+                 enable_wal_mode: bool = True,
+                 health_check_interval: float = 30.0):
         """Initialize database manager with path to SQLite database."""
         self.db_path = db_path
         self.pool_size = pool_size
+        self.retry_count = retry_count
+        self.timeout = timeout
+        self.enable_wal_mode = enable_wal_mode
+        self.health_check_interval = health_check_interval
         self._connection_pool: List[aiosqlite.Connection] = []
         self._pool_lock = None  # Will be initialized in async context
+        self._write_lock = None  # Write serialization lock, initialized in async context
+        
+        # Retry and recovery components - configure with provided settings
+        from .retry_handler import RetryConfig
+        retry_config = RetryConfig(max_attempts=retry_count)
+        self._retry_handler = create_retry_handler(retry_config)
+        self._recovery_manager = None  # Initialized in async context
+        
+        # Health monitoring and metrics
+        self._health_monitor = None  # Initialized in async context
+        self._metrics_collector = DatabaseMetricsCollector()
         
     async def initialize(self) -> None:
         """Initialize database schema and configuration."""
         import asyncio
         
-        # Initialize pool lock
+        # Initialize locks
         self._pool_lock = asyncio.Lock()
+        self._write_lock = asyncio.Lock()
+        
+        # Initialize connection recovery manager
+        self._recovery_manager = ConnectionRecoveryManager(self)
+        
+        # Initialize health monitoring with configured interval
+        self._health_monitor = ConnectionHealthMonitor(
+            self, 
+            check_interval=self.health_check_interval,
+            timeout_seconds=self.timeout
+        )
+        await self._health_monitor.start_monitoring()
         
         # Ensure database directory exists
         self.db_path.parent.mkdir(parents=True, exist_ok=True)
@@ -56,6 +97,9 @@ class DatabaseManager:
             # Enable row factory for easier data access
             db.row_factory = aiosqlite.Row
             
+            # Configure WAL mode and optimizations for concurrent access
+            await self._configure_database_optimizations(db, include_wal_mode=self.enable_wal_mode)
+            
             # Apply each migration
             for migration_file in migration_files:
                 logger.info(f"Applying migration: {migration_file.name}")
@@ -67,6 +111,48 @@ class DatabaseManager:
             
         logger.info(f"Database initialized at {self.db_path} with {len(migration_files)} migrations")
     
+    async def _configure_database_optimizations(self, db: aiosqlite.Connection, include_wal_mode: bool = True) -> None:
+        """
+        Configure SQLite optimizations for concurrent access and performance.
+        
+        Args:
+            db: Database connection to configure
+            include_wal_mode: Whether to set WAL mode (only needed once per database)
+        """
+        optimizations = []
+        
+        # WAL mode is database-level, only set during initialization
+        if include_wal_mode:
+            optimizations.append("PRAGMA journal_mode = WAL")
+            logger.info("Enabling WAL mode for database concurrency")
+        
+        # Connection-level optimizations that can be set per connection
+        optimizations.extend([
+            "PRAGMA synchronous = NORMAL",      # Balance durability/performance  
+            "PRAGMA cache_size = -64000",       # 64MB cache
+            "PRAGMA temp_store = MEMORY",       # Use memory for temp tables
+            "PRAGMA mmap_size = 268435456",     # 256MB memory mapping
+            "PRAGMA busy_timeout = 10000",      # 10 second timeout (reduced from 30s)
+            "PRAGMA optimize"                   # Enable query planner optimizations
+        ])
+        
+        # WAL-specific settings (only if WAL mode is being set)
+        if include_wal_mode:
+            optimizations.append("PRAGMA wal_autocheckpoint = 1000")  # Checkpoint after 1000 pages
+        
+        for pragma in optimizations:
+            try:
+                await db.execute(pragma)
+                logger.debug(f"Applied optimization: {pragma}")
+            except Exception as e:
+                logger.warning(f"Failed to apply optimization '{pragma}': {e}")
+        
+        await db.commit()
+        if include_wal_mode:
+            logger.info("Database optimizations configured for concurrent access with WAL mode")
+        else:
+            logger.debug("Connection optimizations applied")
+    
     @asynccontextmanager
     async def get_connection(self) -> AsyncIterator[aiosqlite.Connection]:
         """Get a database connection from pool or create new one."""
@@ -83,11 +169,8 @@ class DatabaseManager:
             conn = await aiosqlite.connect(self.db_path)
             conn.row_factory = aiosqlite.Row
             
-            # Apply performance settings to new connections
-            await conn.execute("PRAGMA busy_timeout = 30000")  # 30 second timeout
-            await conn.execute("PRAGMA synchronous = NORMAL")   # Balanced durability/performance
-            await conn.execute("PRAGMA cache_size = -64000")    # 64MB cache
-            await conn.execute("PRAGMA temp_store = MEMORY")    # Use memory for temp tables
+            # Apply connection-level optimizations (WAL mode already set during initialization)
+            await self._configure_database_optimizations(conn, include_wal_mode=False)
         
         try:
             yield conn
@@ -104,18 +187,259 @@ class DatabaseManager:
                 await conn.close()
     
     async def close_pool(self) -> None:
-        """Close all connections in the pool."""
+        """Close all connections in the pool and stop monitoring."""
+        # Stop health monitoring
+        if self._health_monitor:
+            await self._health_monitor.stop_monitoring()
+        
+        # Close connections
         if self._pool_lock:
             async with self._pool_lock:
                 for conn in self._connection_pool:
                     await conn.close()
                 self._connection_pool.clear()
     
+    @asynccontextmanager
+    async def get_write_connection(self) -> AsyncIterator[aiosqlite.Connection]:
+        """
+        Get a database connection with write serialization.
+        
+        This ensures only one write operation occurs at a time across the entire
+        application, preventing database locking issues in multi-client scenarios.
+        """
+        if self._write_lock is None:
+            raise RuntimeError("DatabaseManager not initialized - call initialize() first")
+        
+        async with self._write_lock:
+            async with self.get_connection() as conn:
+                yield conn
+    
+    @asynccontextmanager
+    async def get_write_connection_with_retry(self, operation_name: str = "write_operation") -> AsyncIterator[aiosqlite.Connection]:
+        """
+        Get a database connection with write serialization and automatic retry logic.
+        
+        This combines write serialization with retry handling for maximum resilience
+        against database locking issues.
+        
+        Args:
+            operation_name: Name of the operation for logging and monitoring
+        """
+        if self._write_lock is None or self._retry_handler is None:
+            raise RuntimeError("DatabaseManager not initialized - call initialize() first")
+        
+        async with self._retry_handler.with_retry(operation_name):
+            try:
+                async with self._write_lock:
+                    async with self.get_connection() as conn:
+                        yield conn
+                        
+                # Reset failure count on success
+                if self._recovery_manager:
+                    self._recovery_manager.reset_failure_count()
+                    
+            except Exception as e:
+                # Handle persistent failures
+                if self._recovery_manager:
+                    await self._recovery_manager.handle_persistent_failure(operation_name, e)
+                raise
+    
+    def get_database_stats(self) -> Dict[str, Any]:
+        """
+        Get database performance and reliability statistics.
+        
+        Returns:
+            Dictionary with retry stats, recovery stats, health status, and metrics
+        """
+        stats = {
+            "connection_pool": {
+                "configured_size": self.pool_size,
+                "current_size": len(self._connection_pool)
+            }
+        }
+        
+        if self._retry_handler:
+            stats["retry_stats"] = self._retry_handler.get_retry_stats()
+        
+        if self._recovery_manager:
+            stats["recovery_stats"] = self._recovery_manager.get_recovery_stats()
+        
+        if self._health_monitor:
+            stats["health_status"] = self._health_monitor.get_health_status()
+        
+        if self._metrics_collector:
+            stats["operation_metrics"] = self._metrics_collector.get_operation_metrics()
+            stats["locking_frequency"] = self._metrics_collector.get_locking_frequency()
+        
+        return stats
+    
+    async def check_health(self) -> Dict[str, Any]:
+        """
+        Perform an immediate health check and return detailed status.
+        
+        Returns:
+            Dictionary with health check result and current metrics
+        """
+        if not self._health_monitor:
+            return {"error": "Health monitoring not initialized"}
+        
+        # Perform immediate health check
+        health_result = await self._health_monitor.check_health()
+        
+        return {
+            "health_check": {
+                "is_healthy": health_result.is_healthy,
+                "response_time_ms": health_result.response_time_ms,
+                "error_message": health_result.error_message,
+                "timestamp": health_result.timestamp.isoformat()
+            },
+            "overall_status": self._health_monitor.get_health_status(),
+            "recent_history": self._health_monitor.get_recent_history()
+        }
+    
+    @asynccontextmanager
+    async def get_immediate_transaction(
+        self, 
+        operation_name: str = "immediate_transaction",
+        timeout_seconds: float = 10.0
+    ) -> AsyncIterator[aiosqlite.Connection]:
+        """
+        Get a database connection with BEGIN IMMEDIATE transaction and timeout.
+        
+        This ensures write locks are acquired immediately, preventing lock escalation
+        failures that can occur with DEFERRED transactions.
+        
+        Args:
+            operation_name: Name of the operation for monitoring
+            timeout_seconds: Transaction timeout in seconds
+        """
+        async with self.get_write_connection_with_retry(operation_name) as conn:
+            try:
+                # Start immediate transaction with timeout
+                async with asyncio.timeout(timeout_seconds):
+                    await conn.execute("BEGIN IMMEDIATE")
+                    yield conn
+                    await conn.commit()
+            except asyncio.TimeoutError:
+                logger.warning(
+                    f"Transaction timeout after {timeout_seconds}s for {operation_name}",
+                    extra={
+                        "structured_data": {
+                            "transaction_timeout": {
+                                "operation": operation_name,
+                                "timeout_seconds": timeout_seconds
+                            }
+                        }
+                    }
+                )
+                await conn.rollback()
+                raise
+            except Exception as e:
+                logger.error(f"Transaction failed for {operation_name}: {e}")
+                await conn.rollback()
+                raise
+    
+    async def execute_transaction_with_retry(
+        self,
+        operation_func,
+        operation_name: str = "transaction_operation",
+        max_retries: int = 3,
+        timeout_seconds: float = 10.0
+    ) -> Any:
+        """
+        Execute a database operation within a transaction with automatic retry.
+        
+        Args:
+            operation_func: Async function that takes a connection and performs the operation
+            operation_name: Name of the operation for logging
+            max_retries: Maximum retry attempts
+            timeout_seconds: Transaction timeout in seconds
+            
+        Returns:
+            Result from operation_func
+            
+        Example:
+            async def my_operation(conn):
+                await conn.execute("INSERT INTO ...", (...))
+                return "success"
+            
+            result = await db.execute_transaction_with_retry(my_operation, "insert_data")
+        """
+        last_error = None
+        
+        for attempt in range(1, max_retries + 1):
+            try:
+                async with self.get_immediate_transaction(operation_name, timeout_seconds) as conn:
+                    result = await operation_func(conn)
+                    
+                # Record successful operation metrics
+                if self._metrics_collector:
+                    self._metrics_collector.record_operation(
+                        operation_name, 
+                        timeout_seconds * 1000,  # Convert to ms
+                        True,
+                        len(self._connection_pool)
+                    )
+                
+                return result
+                
+            except (aiosqlite.OperationalError, asyncio.TimeoutError) as e:
+                last_error = e
+                
+                # Record locking event for metrics
+                if self._metrics_collector and "locked" in str(e).lower():
+                    self._metrics_collector.record_locking_event(operation_name, str(e))
+                
+                if attempt < max_retries:
+                    # Exponential backoff with jitter
+                    delay = 0.1 * (2 ** (attempt - 1))
+                    jitter = delay * 0.1 * (2 * random.random() - 1)  # ±10% jitter
+                    wait_time = max(0.05, delay + jitter)
+                    
+                    logger.warning(
+                        f"Transaction attempt {attempt} failed for {operation_name}, retrying in {wait_time:.2f}s: {e}",
+                        extra={
+                            "structured_data": {
+                                "transaction_retry": {
+                                    "operation": operation_name,
+                                    "attempt": attempt,
+                                    "delay_seconds": wait_time,
+                                    "error": str(e)
+                                }
+                            }
+                        }
+                    )
+                    await asyncio.sleep(wait_time)
+                else:
+                    # Record failed operation metrics
+                    if self._metrics_collector:
+                        self._metrics_collector.record_operation(
+                            operation_name,
+                            timeout_seconds * 1000,
+                            False,
+                            len(self._connection_pool)
+                        )
+                    
+                    logger.error(
+                        f"Transaction failed after {max_retries} attempts for {operation_name}: {e}",
+                        extra={
+                            "structured_data": {
+                                "transaction_failure": {
+                                    "operation": operation_name,
+                                    "max_retries": max_retries,
+                                    "final_error": str(e)
+                                }
+                            }
+                        }
+                    )
+        
+        raise last_error
+    
     # Project operations
     
     async def create_project(self, project: Project) -> None:
         """Create a new project record."""
-        async with self.get_connection() as db:
+        async with self.get_write_connection_with_retry("create_project") as db:
             await db.execute(
                 """
                 INSERT INTO projects (id, name, remote_origin, upstream_origin, aliases, created, last_accessed)
@@ -302,7 +626,7 @@ class DatabaseManager:
     
     async def update_project_access_time(self, project_id: str) -> None:
         """Update the last accessed time for a project."""
-        async with self.get_connection() as db:
+        async with self.get_write_connection_with_retry("update_project_access_time") as db:
             await db.execute(
                 "UPDATE projects SET last_accessed = ? WHERE id = ?",
                 (datetime.utcnow(), project_id)
@@ -311,7 +635,7 @@ class DatabaseManager:
     
     async def update_project(self, project: Project) -> None:
         """Update an existing project record."""
-        async with self.get_connection() as db:
+        async with self.get_write_connection_with_retry("update_project") as db:
             await db.execute(
                 """
                 UPDATE projects 
@@ -373,7 +697,7 @@ class DatabaseManager:
     
     async def create_file_description(self, file_desc: FileDescription) -> None:
         """Create or update a file description."""
-        async with self.get_connection() as db:
+        async with self.get_write_connection_with_retry("create_file_description") as db:
             await db.execute(
                 """
                 INSERT OR REPLACE INTO file_descriptions 
@@ -456,11 +780,11 @@ class DatabaseManager:
             ]
     
     async def batch_create_file_descriptions(self, file_descriptions: List[FileDescription]) -> None:
-        """Batch create multiple file descriptions efficiently."""
+        """Batch create multiple file descriptions efficiently with optimized transactions."""
         if not file_descriptions:
             return
-            
-        async with self.get_connection() as db:
+        
+        async def batch_operation(conn: aiosqlite.Connection) -> None:
             data = [
                 (
                     fd.project_id,
@@ -475,7 +799,7 @@ class DatabaseManager:
                 for fd in file_descriptions
             ]
             
-            await db.executemany(
+            await conn.executemany(
                 """
                 INSERT OR REPLACE INTO file_descriptions 
                 (project_id, branch, file_path, description, file_hash, last_modified, version, source_project_id)
@@ -483,8 +807,13 @@ class DatabaseManager:
                 """,
                 data
             )
-            await db.commit()
             logger.debug(f"Batch created {len(file_descriptions)} file descriptions")
+        
+        await self.execute_transaction_with_retry(
+            batch_operation,
+            f"batch_create_file_descriptions_{len(file_descriptions)}_files",
+            timeout_seconds=30.0  # Longer timeout for batch operations
+        )
     
     # Search operations
     
@@ -552,7 +881,7 @@ class DatabaseManager:
         """Cache token count with TTL."""
         expires = datetime.utcnow() + timedelta(hours=ttl_hours)
         
-        async with self.get_connection() as db:
+        async with self.get_write_connection() as db:
             await db.execute(
                 """
                 INSERT OR REPLACE INTO token_cache (cache_key, token_count, expires)
@@ -564,7 +893,7 @@ class DatabaseManager:
     
     async def cleanup_expired_cache(self) -> None:
         """Remove expired cache entries."""
-        async with self.get_connection() as db:
+        async with self.get_write_connection() as db:
             await db.execute(
                 "DELETE FROM token_cache WHERE expires < ?",
                 (datetime.utcnow(),)
@@ -663,7 +992,7 @@ class DatabaseManager:
     
     async def create_project_overview(self, overview: ProjectOverview) -> None:
         """Create or update a project overview."""
-        async with self.get_connection() as db:
+        async with self.get_write_connection() as db:
             await db.execute(
                 """
                 INSERT OR REPLACE INTO project_overviews 
@@ -716,9 +1045,9 @@ class DatabaseManager:
         """
         removed_files = []
         
-        async with self.get_connection() as db:
+        async def cleanup_operation(conn: aiosqlite.Connection) -> List[str]:
             # Get all file descriptions for this project/branch
-            cursor = await db.execute(
+            cursor = await conn.execute(
                 "SELECT file_path FROM file_descriptions WHERE project_id = ? AND branch = ?",
                 (project_id, branch)
             )
@@ -733,16 +1062,22 @@ class DatabaseManager:
                 
                 if not full_path.exists():
                     to_remove.append(file_path)
-                    removed_files.append(file_path)
             
             # Remove descriptions for missing files
             if to_remove:
-                await db.executemany(
+                await conn.executemany(
                     "DELETE FROM file_descriptions WHERE project_id = ? AND branch = ? AND file_path = ?",
                     [(project_id, branch, path) for path in to_remove]
                 )
-                await db.commit()
                 logger.info(f"Cleaned up {len(to_remove)} missing files from {project_id}/{branch}")
+            
+            return to_remove
+        
+        removed_files = await self.execute_transaction_with_retry(
+            cleanup_operation,
+            f"cleanup_missing_files_{project_id}_{branch}",
+            timeout_seconds=60.0  # Longer timeout for file system operations
+        )
         
         return removed_files
     
@@ -820,7 +1155,7 @@ class DatabaseManager:
         Returns:
             Number of projects removed
         """
-        async with self.get_connection() as db:
+        async with self.get_write_connection() as db:
             # Find projects with no descriptions and no overview
             cursor = await db.execute("""
                 SELECT p.id, p.name