mcp-hangar 0.2.0__py3-none-any.whl

mcp_hangar/infrastructure/knowledge_base/postgres.py

@@ -0,0 +1,545 @@
+"""PostgreSQL implementation of IKnowledgeBase."""
+
+from datetime import datetime, timedelta, timezone
+import hashlib
+import json
+from typing import Any, Optional
+
+from ...logging_config import get_logger
+from .contracts import AuditEntry, IKnowledgeBase, KnowledgeBaseConfig, MetricEntry, ProviderStateEntry
+
+logger = get_logger(__name__)
+
+# SQL Migrations for PostgreSQL
+MIGRATIONS = [
+    {
+        "version": 1,
+        "name": "initial_schema",
+        "sql": """
+CREATE TABLE IF NOT EXISTS schema_migrations (
+    version INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    applied_at TIMESTAMPTZ DEFAULT NOW()
+);
+
+CREATE TABLE IF NOT EXISTS tool_cache (
+    id SERIAL PRIMARY KEY,
+    provider TEXT NOT NULL,
+    tool TEXT NOT NULL,
+    arguments_hash TEXT NOT NULL,
+    result JSONB NOT NULL,
+    created_at TIMESTAMPTZ DEFAULT NOW(),
+    expires_at TIMESTAMPTZ NOT NULL,
+    UNIQUE(provider, tool, arguments_hash)
+);
+
+CREATE INDEX IF NOT EXISTS idx_tool_cache_lookup
+    ON tool_cache(provider, tool, arguments_hash);
+CREATE INDEX IF NOT EXISTS idx_tool_cache_expires
+    ON tool_cache(expires_at);
+
+CREATE TABLE IF NOT EXISTS audit_log (
+    id SERIAL PRIMARY KEY,
+    timestamp TIMESTAMPTZ DEFAULT NOW(),
+    event_type TEXT NOT NULL,
+    provider TEXT,
+    tool TEXT,
+    arguments JSONB,
+    result_summary TEXT,
+    duration_ms INTEGER,
+    success BOOLEAN NOT NULL,
+    error_message TEXT,
+    correlation_id TEXT
+);
+
+CREATE INDEX IF NOT EXISTS idx_audit_log_timestamp
+    ON audit_log(timestamp DESC);
+CREATE INDEX IF NOT EXISTS idx_audit_log_provider
+    ON audit_log(provider, timestamp DESC);
+
+CREATE TABLE IF NOT EXISTS provider_state_history (
+    id SERIAL PRIMARY KEY,
+    provider_id TEXT NOT NULL,
+    old_state TEXT,
+    new_state TEXT NOT NULL,
+    timestamp TIMESTAMPTZ DEFAULT NOW(),
+    reason TEXT
+);
+
+CREATE INDEX IF NOT EXISTS idx_provider_state_provider
+    ON provider_state_history(provider_id, timestamp DESC);
+
+CREATE TABLE IF NOT EXISTS provider_metrics (
+    id SERIAL PRIMARY KEY,
+    provider_id TEXT NOT NULL,
+    timestamp TIMESTAMPTZ DEFAULT NOW(),
+    metric_name TEXT NOT NULL,
+    metric_value DOUBLE PRECISION NOT NULL,
+    labels JSONB DEFAULT '{}'
+);
+
+CREATE INDEX IF NOT EXISTS idx_provider_metrics_lookup
+    ON provider_metrics(provider_id, metric_name, timestamp DESC);
+""",
+    },
+    {
+        "version": 2,
+        "name": "cleanup_function",
+        "sql": """
+CREATE OR REPLACE FUNCTION cleanup_expired_cache() RETURNS INTEGER AS $$
+DECLARE
+    deleted_count INTEGER;
+BEGIN
+    DELETE FROM tool_cache WHERE expires_at < NOW();
+    GET DIAGNOSTICS deleted_count = ROW_COUNT;
+    RETURN deleted_count;
+END;
+$$ LANGUAGE plpgsql;
+""",
+    },
+]
+
+
+class PostgresKnowledgeBase(IKnowledgeBase):
+    """PostgreSQL implementation of knowledge base."""
+
+    def __init__(self, config: KnowledgeBaseConfig):
+        self._config = config
+        self._pool = None
+        self._initialized = False
+
+    async def initialize(self) -> bool:
+        """Initialize connection pool and run migrations."""
+        try:
+            import asyncpg
+
+            self._pool = await asyncpg.create_pool(
+                self._config.dsn,
+                min_size=1,
+                max_size=self._config.pool_size,
+                command_timeout=10,
+            )
+
+            # Run migrations
+            await self._run_migrations()
+
+            self._initialized = True
+            logger.info(
+                "postgres_kb_initialized",
+                pool_size=self._config.pool_size,
+            )
+            return True
+
+        except ImportError:
+            logger.error("asyncpg_not_installed", hint="pip install asyncpg")
+            return False
+        except Exception as e:
+            logger.error("postgres_kb_init_failed", error=str(e))
+            return False
+
+    async def _run_migrations(self) -> None:
+        """Run pending migrations."""
+        async with self._pool.acquire() as conn:
+            # Get current version
+            try:
+                version = await conn.fetchval("SELECT MAX(version) FROM schema_migrations")
+            except OSError:
+                # Table doesn't exist yet or connection failed - this is first run
+                version = 0
+
+            current_version = version or 0
+
+            for migration in MIGRATIONS:
+                if migration["version"] <= current_version:
+                    continue
+
+                logger.info(
+                    "applying_migration",
+                    version=migration["version"],
+                    name=migration["name"],
+                )
+
+                await conn.execute(migration["sql"])
+                await conn.execute(
+                    "INSERT INTO schema_migrations (version, name) VALUES ($1, $2)",
+                    migration["version"],
+                    migration["name"],
+                )
+
+            final_version = await conn.fetchval("SELECT MAX(version) FROM schema_migrations")
+            logger.info("postgres_kb_schema_ready", version=final_version)
+
+    async def close(self) -> None:
+        """Close connection pool."""
+        if self._pool:
+            await self._pool.close()
+            self._pool = None
+            self._initialized = False
+            logger.info("postgres_kb_closed")
+
+    async def is_healthy(self) -> bool:
+        """Check if database is reachable."""
+        if not self._pool:
+            return False
+        try:
+            async with self._pool.acquire() as conn:
+                await conn.fetchval("SELECT 1")
+            return True
+        except (OSError, ConnectionError, TimeoutError) as e:
+            logger.debug("postgres_health_check_failed", error=str(e))
+            return False
+
+    def _hash_arguments(self, arguments: dict) -> str:
+        """Create hash of arguments for cache key."""
+        serialized = json.dumps(arguments, sort_keys=True, default=str)
+        return hashlib.sha256(serialized.encode()).hexdigest()[:32]
+
+    # === Cache Operations ===
+
+    async def cache_get(self, provider: str, tool: str, arguments: dict) -> Optional[dict]:
+        if not self._pool:
+            return None
+
+        args_hash = self._hash_arguments(arguments)
+
+        try:
+            async with self._pool.acquire() as conn:
+                row = await conn.fetchrow(
+                    """
+                    SELECT result FROM tool_cache
+                    WHERE provider = $1 AND tool = $2 AND arguments_hash = $3
+                      AND expires_at > NOW()
+                    """,
+                    provider,
+                    tool,
+                    args_hash,
+                )
+                if row:
+                    logger.debug("cache_hit", provider=provider, tool=tool)
+                    return json.loads(row["result"])
+                return None
+        except Exception as e:
+            logger.warning("cache_get_failed", error=str(e))
+            return None
+
+    async def cache_set(
+        self,
+        provider: str,
+        tool: str,
+        arguments: dict,
+        result: Any,
+        ttl_s: Optional[int] = None,
+    ) -> bool:
+        if not self._pool:
+            return False
+
+        args_hash = self._hash_arguments(arguments)
+        ttl = ttl_s or self._config.cache_ttl_s
+        expires_at = datetime.now(timezone.utc) + timedelta(seconds=ttl)
+
+        try:
+            async with self._pool.acquire() as conn:
+                await conn.execute(
+                    """
+                    INSERT INTO tool_cache (provider, tool, arguments_hash, result, expires_at)
+                    VALUES ($1, $2, $3, $4, $5)
+                    ON CONFLICT (provider, tool, arguments_hash)
+                    DO UPDATE SET result = $4, expires_at = $5, created_at = NOW()
+                    """,
+                    provider,
+                    tool,
+                    args_hash,
+                    json.dumps(result, default=str),
+                    expires_at,
+                )
+                return True
+        except Exception as e:
+            logger.warning("cache_set_failed", error=str(e))
+            return False
+
+    async def cache_invalidate(self, provider: Optional[str] = None, tool: Optional[str] = None) -> int:
+        if not self._pool:
+            return 0
+
+        try:
+            async with self._pool.acquire() as conn:
+                if provider and tool:
+                    result = await conn.execute(
+                        "DELETE FROM tool_cache WHERE provider = $1 AND tool = $2",
+                        provider,
+                        tool,
+                    )
+                elif provider:
+                    result = await conn.execute("DELETE FROM tool_cache WHERE provider = $1", provider)
+                else:
+                    result = await conn.execute("DELETE FROM tool_cache")
+                # Parse "DELETE N" to get count
+                return int(result.split()[-1]) if result else 0
+        except Exception as e:
+            logger.warning("cache_invalidate_failed", error=str(e))
+            return 0
+
+    async def cache_cleanup(self) -> int:
+        if not self._pool:
+            return 0
+
+        try:
+            async with self._pool.acquire() as conn:
+                result = await conn.fetchval("SELECT cleanup_expired_cache()")
+                logger.info("cache_cleanup", deleted=result)
+                return result or 0
+        except Exception as e:
+            logger.warning("cache_cleanup_failed", error=str(e))
+            return 0
+
+    # === Audit Operations ===
+
+    async def audit_log(self, entry: AuditEntry) -> bool:
+        if not self._pool:
+            return False
+
+        try:
+            async with self._pool.acquire() as conn:
+                await conn.execute(
+                    """
+                    INSERT INTO audit_log
+                        (event_type, provider, tool, arguments, result_summary,
+                         duration_ms, success, error_message, correlation_id)
+                    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
+                    """,
+                    entry.event_type,
+                    entry.provider,
+                    entry.tool,
+                    json.dumps(entry.arguments, default=str) if entry.arguments else None,
+                    entry.result_summary,
+                    entry.duration_ms,
+                    entry.success,
+                    entry.error_message,
+                    entry.correlation_id,
+                )
+                return True
+        except Exception as e:
+            logger.warning("audit_log_failed", error=str(e))
+            return False
+
+    async def audit_query(
+        self,
+        provider: Optional[str] = None,
+        tool: Optional[str] = None,
+        success: Optional[bool] = None,
+        since: Optional[datetime] = None,
+        limit: int = 100,
+    ) -> list[AuditEntry]:
+        if not self._pool:
+            return []
+
+        try:
+            conditions = []
+            params = []
+            param_idx = 1
+
+            if provider:
+                conditions.append(f"provider = ${param_idx}")
+                params.append(provider)
+                param_idx += 1
+            if tool:
+                conditions.append(f"tool = ${param_idx}")
+                params.append(tool)
+                param_idx += 1
+            if success is not None:
+                conditions.append(f"success = ${param_idx}")
+                params.append(success)
+                param_idx += 1
+            if since:
+                conditions.append(f"timestamp >= ${param_idx}")
+                params.append(since)
+                param_idx += 1
+
+            where_clause = " AND ".join(conditions) if conditions else "1=1"
+            params.append(limit)
+
+            async with self._pool.acquire() as conn:
+                rows = await conn.fetch(
+                    f"""
+                    SELECT event_type, provider, tool, arguments, result_summary,
+                           duration_ms, success, error_message, correlation_id, timestamp
+                    FROM audit_log
+                    WHERE {where_clause}
+                    ORDER BY timestamp DESC
+                    LIMIT ${param_idx}
+                    """,
+                    *params,
+                )
+
+                return [
+                    AuditEntry(
+                        event_type=row["event_type"],
+                        provider=row["provider"],
+                        tool=row["tool"],
+                        arguments=json.loads(row["arguments"]) if row["arguments"] else None,
+                        result_summary=row["result_summary"],
+                        duration_ms=row["duration_ms"],
+                        success=row["success"],
+                        error_message=row["error_message"],
+                        correlation_id=row["correlation_id"],
+                        timestamp=row["timestamp"],
+                    )
+                    for row in rows
+                ]
+        except Exception as e:
+            logger.warning("audit_query_failed", error=str(e))
+            return []
+
+    async def audit_stats(self, hours: int = 24) -> dict:
+        if not self._pool:
+            return {}
+
+        try:
+            async with self._pool.acquire() as conn:
+                row = await conn.fetchrow(
+                    f"""
+                    SELECT
+                        COUNT(*) as total,
+                        COUNT(*) FILTER (WHERE success) as success_count,
+                        COUNT(*) FILTER (WHERE NOT success) as error_count,
+                        COUNT(DISTINCT provider) as providers,
+                        COUNT(DISTINCT tool) as tools,
+                        AVG(duration_ms) FILTER (WHERE duration_ms IS NOT NULL) as avg_duration_ms
+                    FROM audit_log
+                    WHERE timestamp > NOW() - INTERVAL '{hours} hours'
+                    """
+                )
+                return dict(row) if row else {}
+        except Exception as e:
+            logger.warning("audit_stats_failed", error=str(e))
+            return {}
+
+    # === Provider State Operations ===
+
+    async def record_state_change(self, entry: ProviderStateEntry) -> bool:
+        if not self._pool:
+            return False
+
+        try:
+            async with self._pool.acquire() as conn:
+                await conn.execute(
+                    """
+                    INSERT INTO provider_state_history
+                        (provider_id, old_state, new_state, reason)
+                    VALUES ($1, $2, $3, $4)
+                    """,
+                    entry.provider_id,
+                    entry.old_state,
+                    entry.new_state,
+                    entry.reason,
+                )
+                return True
+        except Exception as e:
+            logger.warning("record_state_failed", error=str(e))
+            return False
+
+    async def get_state_history(self, provider_id: str, limit: int = 100) -> list[ProviderStateEntry]:
+        if not self._pool:
+            return []
+
+        try:
+            async with self._pool.acquire() as conn:
+                rows = await conn.fetch(
+                    """
+                    SELECT provider_id, old_state, new_state, reason, timestamp
+                    FROM provider_state_history
+                    WHERE provider_id = $1
+                    ORDER BY timestamp DESC
+                    LIMIT $2
+                    """,
+                    provider_id,
+                    limit,
+                )
+                return [
+                    ProviderStateEntry(
+                        provider_id=row["provider_id"],
+                        old_state=row["old_state"],
+                        new_state=row["new_state"],
+                        reason=row["reason"],
+                        timestamp=row["timestamp"],
+                    )
+                    for row in rows
+                ]
+        except Exception as e:
+            logger.warning("get_state_history_failed", error=str(e))
+            return []
+
+    # === Metrics Operations ===
+
+    async def record_metric(self, entry: MetricEntry) -> bool:
+        if not self._pool:
+            return False
+
+        try:
+            async with self._pool.acquire() as conn:
+                await conn.execute(
+                    """
+                    INSERT INTO provider_metrics
+                        (provider_id, metric_name, metric_value, labels)
+                    VALUES ($1, $2, $3, $4)
+                    """,
+                    entry.provider_id,
+                    entry.metric_name,
+                    entry.metric_value,
+                    json.dumps(entry.labels or {}),
+                )
+                return True
+        except Exception as e:
+            logger.warning("record_metric_failed", error=str(e))
+            return False
+
+    async def get_metrics(
+        self,
+        provider_id: str,
+        metric_name: Optional[str] = None,
+        since: Optional[datetime] = None,
+        limit: int = 100,
+    ) -> list[MetricEntry]:
+        if not self._pool:
+            return []
+
+        try:
+            conditions = ["provider_id = $1"]
+            params = [provider_id]
+            param_idx = 2
+
+            if metric_name:
+                conditions.append(f"metric_name = ${param_idx}")
+                params.append(metric_name)
+                param_idx += 1
+            if since:
+                conditions.append(f"timestamp >= ${param_idx}")
+                params.append(since)
+                param_idx += 1
+
+            params.append(limit)
+
+            async with self._pool.acquire() as conn:
+                rows = await conn.fetch(
+                    f"""
+                    SELECT provider_id, metric_name, metric_value, labels, timestamp
+                    FROM provider_metrics
+                    WHERE {" AND ".join(conditions)}
+                    ORDER BY timestamp DESC
+                    LIMIT ${param_idx}
+                    """,
+                    *params,
+                )
+                return [
+                    MetricEntry(
+                        provider_id=row["provider_id"],
+                        metric_name=row["metric_name"],
+                        metric_value=row["metric_value"],
+                        labels=json.loads(row["labels"]) if row["labels"] else None,
+                        timestamp=row["timestamp"],
+                    )
+                    for row in rows
+                ]
+        except Exception as e:
+            logger.warning("get_metrics_failed", error=str(e))
+            return []