mcp-hangar 0.2.0__py3-none-any.whl

mcp_hangar/infrastructure/persistence/database.py

@@ -0,0 +1,333 @@
+"""Database connection management for SQLite persistence.
+
+Provides async-compatible database access with connection pooling,
+migrations, and health checking.
+"""
+
+import asyncio
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from pathlib import Path
+import threading
+from typing import Any, AsyncIterator, Dict, List, Optional, Tuple
+
+import aiosqlite
+
+from ...logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class DatabaseConfig:
+    """Configuration for database connection.
+
+    Attributes:
+        path: Path to SQLite database file. Use ":memory:" for in-memory.
+        timeout: Connection timeout in seconds.
+        isolation_level: SQLite isolation level.
+        check_same_thread: Whether to enforce same-thread access.
+        enable_wal: Enable Write-Ahead Logging for better concurrency.
+        busy_timeout_ms: Timeout for busy handler in milliseconds.
+    """
+
+    path: str = "data/mcp_hangar.db"
+    timeout: float = 30.0
+    isolation_level: Optional[str] = "DEFERRED"
+    check_same_thread: bool = False
+    enable_wal: bool = True
+    busy_timeout_ms: int = 5000
+
+    def __post_init__(self):
+        # Ensure data directory exists for file-based database
+        if self.path != ":memory:":
+            Path(self.path).parent.mkdir(parents=True, exist_ok=True)
+
+
+class Database:
+    """Async SQLite database wrapper with connection pooling.
+
+    Provides:
+    - Async connection management
+    - Automatic migrations
+    - Connection health checking
+    - Thread-safe connection pool
+    """
+
+    def __init__(self, config: Optional[DatabaseConfig] = None):
+        """Initialize database with configuration.
+
+        Args:
+            config: Database configuration. Defaults to file-based DB.
+        """
+        self._config = config or DatabaseConfig()
+        self._lock = asyncio.Lock()
+        self._initialized = False
+        self._migrations_applied = False
+
+    @property
+    def config(self) -> DatabaseConfig:
+        """Get current database configuration."""
+        return self._config
+
+    @asynccontextmanager
+    async def connection(self) -> AsyncIterator[aiosqlite.Connection]:
+        """Get a database connection.
+
+        Yields:
+            Async SQLite connection
+
+        Example:
+            async with db.connection() as conn:
+                await conn.execute("SELECT * FROM providers")
+        """
+        conn = await aiosqlite.connect(
+            self._config.path,
+            timeout=self._config.timeout,
+            isolation_level=self._config.isolation_level,
+        )
+        try:
+            # Configure connection
+            await conn.execute(f"PRAGMA busy_timeout = {self._config.busy_timeout_ms}")
+
+            if self._config.enable_wal and self._config.path != ":memory:":
+                await conn.execute("PRAGMA journal_mode = WAL")
+
+            # Enable foreign keys
+            await conn.execute("PRAGMA foreign_keys = ON")
+
+            conn.row_factory = aiosqlite.Row
+            yield conn
+        finally:
+            await conn.close()
+
+    @asynccontextmanager
+    async def transaction(self) -> AsyncIterator[aiosqlite.Connection]:
+        """Get a database connection within a transaction.
+
+        Automatically commits on success, rolls back on exception.
+
+        Yields:
+            Async SQLite connection
+        """
+        async with self.connection() as conn:
+            try:
+                yield conn
+                await conn.commit()
+            except (aiosqlite.Error, ValueError, TypeError) as e:
+                logger.debug("transaction_rollback", error=str(e))
+                await conn.rollback()
+                raise
+
+    async def initialize(self) -> None:
+        """Initialize database and run migrations.
+
+        Creates tables and applies any pending migrations.
+        Safe to call multiple times - idempotent.
+        """
+        async with self._lock:
+            if self._initialized:
+                return
+
+            await self._apply_migrations()
+            self._initialized = True
+            logger.info(f"Database initialized: {self._config.path}")
+
+    async def _apply_migrations(self) -> None:
+        """Apply database migrations."""
+        async with self.connection() as conn:
+            # Create migrations tracking table
+            await conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS _migrations (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    name TEXT NOT NULL UNIQUE,
+                    applied_at TEXT NOT NULL DEFAULT (datetime('now'))
+                )
+            """
+            )
+
+            # Get applied migrations
+            cursor = await conn.execute("SELECT name FROM _migrations")
+            applied = {row[0] for row in await cursor.fetchall()}
+
+            # Apply pending migrations in order
+            for migration_name, migration_sql in MIGRATIONS:
+                if migration_name not in applied:
+                    logger.info(f"Applying migration: {migration_name}")
+                    await conn.executescript(migration_sql)
+                    await conn.execute(
+                        "INSERT INTO _migrations (name) VALUES (?)",
+                        (migration_name,),
+                    )
+
+            await conn.commit()
+            self._migrations_applied = True
+
+    async def health_check(self) -> Dict[str, Any]:
+        """Check database health.
+
+        Returns:
+            Dictionary with health status and metrics
+        """
+        try:
+            async with self.connection() as conn:
+                # Basic connectivity check
+                cursor = await conn.execute("SELECT 1")
+                await cursor.fetchone()
+
+                # Get database stats
+                cursor = await conn.execute("SELECT COUNT(*) FROM provider_configs")
+                provider_count = (await cursor.fetchone())[0]
+
+                cursor = await conn.execute("SELECT COUNT(*) FROM audit_log")
+                audit_count = (await cursor.fetchone())[0]
+
+                return {
+                    "status": "healthy",
+                    "database_path": self._config.path,
+                    "initialized": self._initialized,
+                    "migrations_applied": self._migrations_applied,
+                    "provider_count": provider_count,
+                    "audit_entries": audit_count,
+                }
+        except Exception as e:
+            logger.error(f"Database health check failed: {e}")
+            return {
+                "status": "unhealthy",
+                "error": str(e),
+                "database_path": self._config.path,
+            }
+
+    async def close(self) -> None:
+        """Close database resources.
+
+        For SQLite with aiosqlite, connections are managed per-operation,
+        but this method ensures clean shutdown.
+        """
+        self._initialized = False
+        logger.info("Database closed")
+
+
+# Database migrations - applied in order
+MIGRATIONS: List[Tuple[str, str]] = [
+    (
+        "001_initial_schema",
+        """
+        -- Provider configurations table
+        CREATE TABLE IF NOT EXISTS provider_configs (
+            provider_id TEXT PRIMARY KEY,
+            mode TEXT NOT NULL,
+            config_json TEXT NOT NULL,
+            enabled INTEGER NOT NULL DEFAULT 1,
+            version INTEGER NOT NULL DEFAULT 1,
+            created_at TEXT NOT NULL DEFAULT (datetime('now')),
+            updated_at TEXT NOT NULL DEFAULT (datetime('now'))
+        );
+
+        -- Index for listing enabled providers
+        CREATE INDEX IF NOT EXISTS idx_provider_configs_enabled
+            ON provider_configs(enabled);
+
+        -- Audit log table
+        CREATE TABLE IF NOT EXISTS audit_log (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            entity_id TEXT NOT NULL,
+            entity_type TEXT NOT NULL,
+            action TEXT NOT NULL,
+            actor TEXT NOT NULL,
+            timestamp TEXT NOT NULL,
+            old_state_json TEXT,
+            new_state_json TEXT,
+            metadata_json TEXT,
+            correlation_id TEXT,
+            created_at TEXT NOT NULL DEFAULT (datetime('now'))
+        );
+
+        -- Indexes for audit log queries
+        CREATE INDEX IF NOT EXISTS idx_audit_log_entity
+            ON audit_log(entity_id, entity_type);
+        CREATE INDEX IF NOT EXISTS idx_audit_log_timestamp
+            ON audit_log(timestamp DESC);
+        CREATE INDEX IF NOT EXISTS idx_audit_log_correlation
+            ON audit_log(correlation_id) WHERE correlation_id IS NOT NULL;
+        CREATE INDEX IF NOT EXISTS idx_audit_log_action
+            ON audit_log(action, timestamp DESC);
+        """,
+    ),
+    (
+        "002_add_provider_metadata",
+        """
+        -- Add metadata column to provider_configs
+        ALTER TABLE provider_configs ADD COLUMN metadata_json TEXT;
+
+        -- Add last_started_at for recovery prioritization
+        ALTER TABLE provider_configs ADD COLUMN last_started_at TEXT;
+
+        -- Add failure_count for recovery decisions
+        ALTER TABLE provider_configs ADD COLUMN consecutive_failures INTEGER DEFAULT 0;
+        """,
+    ),
+    (
+        "003_add_audit_indexes",
+        """
+        -- Composite index for time-range queries with filters
+        CREATE INDEX IF NOT EXISTS idx_audit_log_time_entity
+            ON audit_log(timestamp DESC, entity_type);
+
+        -- Index for actor-based queries (who did what)
+        CREATE INDEX IF NOT EXISTS idx_audit_log_actor
+            ON audit_log(actor, timestamp DESC);
+        """,
+    ),
+]
+
+
+# Singleton database instance
+_database: Optional[Database] = None
+_database_lock = threading.Lock()
+
+
+def get_database(config: Optional[DatabaseConfig] = None) -> Database:
+    """Get or create the global database instance.
+
+    Args:
+        config: Optional configuration. Only used when creating new instance.
+
+    Returns:
+        Database instance
+    """
+    global _database
+    with _database_lock:
+        if _database is None:
+            _database = Database(config)
+        return _database
+
+
+def set_database(database: Database) -> None:
+    """Set the global database instance.
+
+    Useful for testing with custom configurations.
+
+    Args:
+        database: Database instance to use
+    """
+    global _database
+    with _database_lock:
+        _database = database
+
+
+async def initialize_database(config: Optional[DatabaseConfig] = None) -> Database:
+    """Initialize and return the database.
+
+    Convenience function that gets the database and initializes it.
+
+    Args:
+        config: Optional database configuration
+
+    Returns:
+        Initialized database instance
+    """
+    db = get_database(config)
+    await db.initialize()
+    return db

mcp_hangar/infrastructure/persistence/database_common.py

@@ -0,0 +1,330 @@
+"""Common database utilities for SQLite and PostgreSQL.
+
+Provides shared connection management, schema migrations, and utilities
+that can be reused across different stores (auth, events, knowledge base).
+"""
+
+from contextlib import contextmanager
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+import sqlite3
+import threading
+from typing import Any, Generator, Protocol
+
+import structlog
+
+logger = structlog.get_logger(__name__)
+
+
+@dataclass
+class SQLiteConfig:
+    """SQLite database configuration.
+
+    Attributes:
+        path: Path to database file. Use ":memory:" for in-memory.
+        enable_wal: Enable Write-Ahead Logging for better concurrency.
+        busy_timeout_ms: Timeout for busy handler in milliseconds.
+        foreign_keys: Enable foreign key constraints.
+    """
+
+    path: str = ":memory:"
+    enable_wal: bool = True
+    busy_timeout_ms: int = 5000
+    foreign_keys: bool = True
+
+    def __post_init__(self):
+        if self.path != ":memory:":
+            Path(self.path).parent.mkdir(parents=True, exist_ok=True)
+
+
+@dataclass
+class PostgresConfig:
+    """PostgreSQL database configuration.
+
+    Attributes:
+        host: Database host.
+        port: Database port.
+        database: Database name.
+        user: Database user.
+        password: Database password.
+        min_connections: Minimum pool connections.
+        max_connections: Maximum pool connections.
+    """
+
+    host: str = "localhost"
+    port: int = 5432
+    database: str = "mcp_hangar"
+    user: str = "mcp_hangar"
+    password: str = ""
+    min_connections: int = 2
+    max_connections: int = 10
+
+
+class IConnectionFactory(Protocol):
+    """Protocol for database connection factories."""
+
+    def get_connection(self) -> Any:
+        """Get a database connection."""
+        ...
+
+    def close(self) -> None:
+        """Close all connections."""
+        ...
+
+
+class SQLiteConnectionFactory:
+    """Thread-safe SQLite connection factory.
+
+    Uses thread-local storage to provide one connection per thread.
+    For in-memory databases, uses a single persistent connection.
+    """
+
+    def __init__(self, config: SQLiteConfig):
+        self._config = config
+        self._local = threading.local()
+        self._lock = threading.Lock()
+
+        # For in-memory database, keep a persistent connection
+        self._persistent_conn: sqlite3.Connection | None = None
+        if config.path == ":memory:":
+            self._persistent_conn = self._create_connection()
+
+    def _create_connection(self) -> sqlite3.Connection:
+        """Create a new database connection with proper settings."""
+        conn = sqlite3.connect(
+            self._config.path,
+            check_same_thread=False,
+            timeout=self._config.busy_timeout_ms / 1000,
+        )
+        conn.row_factory = sqlite3.Row
+
+        if self._config.foreign_keys:
+            conn.execute("PRAGMA foreign_keys = ON")
+
+        if self._config.enable_wal and self._config.path != ":memory:":
+            conn.execute("PRAGMA journal_mode = WAL")
+
+        conn.execute(f"PRAGMA busy_timeout = {self._config.busy_timeout_ms}")
+
+        return conn
+
+    @contextmanager
+    def get_connection(self) -> Generator[sqlite3.Connection, None, None]:
+        """Get a database connection for the current thread.
+
+        Yields:
+            sqlite3.Connection configured for use.
+        """
+        if self._persistent_conn is not None:
+            yield self._persistent_conn
+            return
+
+        if not hasattr(self._local, "connection") or self._local.connection is None:
+            self._local.connection = self._create_connection()
+
+        yield self._local.connection
+
+    def close(self) -> None:
+        """Close all connections."""
+        if self._persistent_conn:
+            try:
+                self._persistent_conn.commit()
+            except Exception:
+                pass
+            self._persistent_conn.close()
+            self._persistent_conn = None
+
+        if hasattr(self._local, "connection") and self._local.connection:
+            try:
+                self._local.connection.commit()
+                self._local.connection.execute("PRAGMA wal_checkpoint(TRUNCATE)")
+            except Exception:
+                pass
+            self._local.connection.close()
+            self._local.connection = None
+
+
+class PostgresConnectionFactory:
+    """PostgreSQL connection factory with connection pooling.
+
+    Uses psycopg2 ThreadedConnectionPool for thread-safe connections.
+    """
+
+    def __init__(self, config: PostgresConfig):
+        self._config = config
+        self._pool = None
+
+    def _ensure_pool(self):
+        """Lazily create connection pool."""
+        if self._pool is None:
+            try:
+                import psycopg2  # noqa: F401
+                from psycopg2 import pool
+            except ImportError:
+                raise ImportError("psycopg2 is required for PostgreSQL. " "Install with: pip install psycopg2-binary")
+
+            self._pool = pool.ThreadedConnectionPool(
+                minconn=self._config.min_connections,
+                maxconn=self._config.max_connections,
+                host=self._config.host,
+                port=self._config.port,
+                database=self._config.database,
+                user=self._config.user,
+                password=self._config.password,
+            )
+
+    @contextmanager
+    def get_connection(self) -> Generator[Any, None, None]:
+        """Get a database connection from the pool.
+
+        Yields:
+            psycopg2 connection.
+        """
+        self._ensure_pool()
+        conn = self._pool.getconn()
+        try:
+            yield conn
+        finally:
+            self._pool.putconn(conn)
+
+    def close(self) -> None:
+        """Close the connection pool."""
+        if self._pool:
+            self._pool.closeall()
+            self._pool = None
+
+
+class MigrationRunner:
+    """Runs database migrations in order.
+
+    Supports both SQLite and PostgreSQL via connection factory.
+    """
+
+    def __init__(
+        self,
+        connection_factory: IConnectionFactory,
+        migrations: list[dict],
+        table_name: str = "schema_migrations",
+    ):
+        """Initialize migration runner.
+
+        Args:
+            connection_factory: Factory for database connections.
+            migrations: List of migration dicts with 'version', 'name', 'sql'.
+            table_name: Name of migrations tracking table.
+        """
+        self._conn_factory = connection_factory
+        self._migrations = sorted(migrations, key=lambda m: m["version"])
+        self._table_name = table_name
+
+    def run(self) -> int:
+        """Run pending migrations.
+
+        Returns:
+            Number of migrations applied.
+        """
+        with self._conn_factory.get_connection() as conn:
+            # Create migrations table
+            self._ensure_migrations_table(conn)
+
+            # Get current version
+            current_version = self._get_current_version(conn)
+
+            # Apply pending migrations
+            applied = 0
+            for migration in self._migrations:
+                if migration["version"] > current_version:
+                    self._apply_migration(conn, migration)
+                    applied += 1
+
+            return applied
+
+    def _ensure_migrations_table(self, conn) -> None:
+        """Create migrations tracking table if not exists."""
+        cursor = conn.cursor()
+        cursor.execute(
+            f"""
+            CREATE TABLE IF NOT EXISTS {self._table_name} (
+                version INTEGER PRIMARY KEY,
+                name TEXT NOT NULL,
+                applied_at TEXT NOT NULL
+            )
+        """
+        )
+        conn.commit()
+
+    def _get_current_version(self, conn) -> int:
+        """Get the current schema version."""
+        cursor = conn.cursor()
+        cursor.execute(f"SELECT MAX(version) FROM {self._table_name}")
+        row = cursor.fetchone()
+        return row[0] if row and row[0] else 0
+
+    def _apply_migration(self, conn, migration: dict) -> None:
+        """Apply a single migration."""
+        cursor = conn.cursor()
+
+        logger.info(
+            "applying_migration",
+            version=migration["version"],
+            name=migration["name"],
+        )
+
+        # Execute migration SQL
+        if hasattr(conn, "executescript"):
+            # SQLite
+            conn.executescript(migration["sql"])
+        else:
+            # PostgreSQL
+            cursor.execute(migration["sql"])
+
+        # Record migration
+        cursor.execute(
+            (
+                f"INSERT INTO {self._table_name} (version, name, applied_at) VALUES (?, ?, ?)"
+                if hasattr(conn, "executescript")
+                else f"INSERT INTO {self._table_name} (version, name, applied_at) VALUES (%s, %s, %s)"
+            ),
+            (migration["version"], migration["name"], datetime.now(timezone.utc).isoformat()),
+        )
+
+        conn.commit()
+
+        logger.info(
+            "migration_applied",
+            version=migration["version"],
+            name=migration["name"],
+        )
+
+
+def create_connection_factory(
+    driver: str,
+    sqlite_config: SQLiteConfig | None = None,
+    postgres_config: PostgresConfig | None = None,
+) -> IConnectionFactory:
+    """Create appropriate connection factory based on driver.
+
+    Args:
+        driver: Database driver ("sqlite" or "postgresql").
+        sqlite_config: SQLite configuration (if driver is sqlite).
+        postgres_config: PostgreSQL configuration (if driver is postgresql).
+
+    Returns:
+        Connection factory for the specified driver.
+
+    Raises:
+        ValueError: If unknown driver or missing config.
+    """
+    if driver == "sqlite":
+        if sqlite_config is None:
+            sqlite_config = SQLiteConfig()
+        return SQLiteConnectionFactory(sqlite_config)
+
+    elif driver in ("postgresql", "postgres"):
+        if postgres_config is None:
+            postgres_config = PostgresConfig()
+        return PostgresConnectionFactory(postgres_config)
+
+    else:
+        raise ValueError(f"Unknown database driver: {driver}")