postgres-aiops 0.1.0__py3-none-any.whl

mcp_server/__init__.py

@@ -0,0 +1 @@
+"""MCP server package for postgres-aiops."""

mcp_server/_shared.py

@@ -0,0 +1,101 @@
+"""Shared MCP server primitives: the FastMCP instance, connection helper,
+error sanitisation, and the ``@tool_errors`` decorator.
+
+Tool modules under ``mcp_server/tools/`` import ``mcp`` from here and register
+their ``@mcp.tool()`` functions onto it. ``mcp_server/server.py`` then imports
+those modules and runs the server.
+
+Keep ``Optional[X]`` (never PEP 604 ``X | None``) in any FastMCP-reflected
+tool signature — on older mcp/pydantic the union eval'd to ``types.UnionType``
+crashes FastMCP's ``issubclass`` check.
+"""
+
+import functools
+import logging
+import os
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, Optional
+
+from mcp.server.fastmcp import FastMCP
+
+from postgres_aiops.config import load_config
+from postgres_aiops.connection import ConnectionManager, PgError
+from postgres_aiops.governance import sanitize
+
+logger = logging.getLogger(__name__)
+
+_DOCTOR_HINT = "Run 'postgres-aiops doctor' to verify connectivity and credentials."
+
+
+def _safe_error(exc: Exception, tool: str) -> str:
+    """Return an agent-safe error string; log full detail server-side only."""
+    logger.error("Tool %s failed", tool, exc_info=True)
+    _passthrough = (
+        ValueError,
+        FileNotFoundError,
+        KeyError,
+        PermissionError,
+        TimeoutError,
+        ConnectionError,
+        PgError,
+    )
+    if isinstance(exc, _passthrough):
+        return sanitize(str(exc), 300)
+    return f"{type(exc).__name__}: operation failed."
+
+
+def tool_errors(shape: str = "dict") -> Callable:
+    """Wrap a tool body in the canonical try/except → ``_safe_error`` pattern.
+
+    Place this *between* ``@governed_tool`` and the function so the audit
+    decorator and FastMCP still see the original signature.
+    """
+
+    def decorator(func: Callable) -> Callable:
+        name = func.__name__
+
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            try:
+                return func(*args, **kwargs)
+            except Exception as e:  # noqa: BLE001 — sanitised below
+                msg = _safe_error(e, name)
+                if shape == "list":
+                    return [{"error": msg, "hint": _DOCTOR_HINT}]
+                if shape == "str":
+                    return f"Error: {msg} {_DOCTOR_HINT}"
+                return {"error": msg, "hint": _DOCTOR_HINT}
+
+        return wrapper
+
+    return decorator
+
+
+mcp = FastMCP(
+    "postgres-aiops",
+    instructions=(
+        "Governed PostgreSQL DBA operations (preview): a one-shot cluster "
+        "'overview'; server reads (version/settings/extensions/databases/roles); "
+        "activity (sessions, long-running queries, locks); query stats "
+        "(pg_stat_statements top-N, EXPLAIN); index and table health (unused / "
+        "missing / bloat / autovacuum); replication (lag, slots, WAL); three "
+        "flagship analyses — 'slow_query_rca', 'bloat_and_vacuum_analysis', and "
+        "'blocking_lock_chain_rca'; and guarded writes (terminate/cancel, "
+        "vacuum/analyze, create/drop index, reindex, ALTER SYSTEM). Every tool "
+        "runs through the postgres-aiops governance harness (audit / budget / "
+        "risk-tier / undo). Do NOT use for OT/industrial edge — see industrial-aiops."
+    ),
+)
+
+_conn_mgr: Optional[ConnectionManager] = None
+
+
+def _get_connection(target: Optional[str] = None) -> Any:
+    """Return a PostgreSQL connection, lazily initialising the manager."""
+    global _conn_mgr  # noqa: PLW0603
+    if _conn_mgr is None:
+        config_path_str = os.environ.get("POSTGRES_AIOPS_CONFIG")
+        config_path = Path(config_path_str) if config_path_str else None
+        _conn_mgr = ConnectionManager(load_config(config_path))
+    return _conn_mgr.connect(target)

mcp_server/server.py

@@ -0,0 +1,37 @@
+"""MCP server wrapping postgres-aiops operations (stdio transport).
+
+Thin adapter layer: each ``@mcp.tool()`` function (in ``mcp_server/tools/``)
+delegates to the ``postgres_aiops`` ops package and is wrapped with the
+postgres-aiops ``@governed_tool`` harness (audit / budget / undo / risk-tier).
+
+Standalone, self-governed PostgreSQL DBA operations (preview).
+For PostgreSQL servers/clusters via psycopg 3.
+
+Source: https://github.com/AIops-tools/Postgres-AIops
+License: MIT
+"""
+
+import logging
+
+from mcp_server._shared import _safe_error, mcp, tool_errors
+
+# Importing the tool modules registers every @mcp.tool() onto the shared
+# `mcp` instance. Order does not matter; each module is self-contained.
+from mcp_server.tools import (  # noqa: F401 — side effects
+    activity,
+    analysis,
+    indexes,
+    queries,
+    remediation,
+    replication,
+    server,
+    tables,
+)
+
+__all__ = ["mcp", "main", "_safe_error", "tool_errors"]
+
+
+def main() -> None:
+    """Run the MCP server over stdio."""
+    logging.basicConfig(level=logging.INFO)
+    mcp.run(transport="stdio")

mcp_server/tools/__init__.py

@@ -0,0 +1 @@
+"""MCP tool modules. Importing each registers its @mcp.tool() functions."""

mcp_server/tools/activity.py

@@ -0,0 +1,52 @@
+"""Activity PostgreSQL MCP tools (read-only): sessions, long queries, locks."""
+
+from typing import Optional
+
+from mcp_server._shared import _get_connection, mcp, tool_errors
+from postgres_aiops.governance import governed_tool
+from postgres_aiops.ops import activity as ops
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def list_activity(
+    state: Optional[str] = None,
+    include_idle: bool = True,
+    target: Optional[str] = None,
+) -> dict:
+    """[READ] Current sessions (pg_stat_activity) with per-state counts.
+
+    Flags idle-in-transaction backends (open transactions holding resources).
+
+    Args:
+        state: Optional exact state filter (active, idle, 'idle in transaction').
+        include_idle: Include plain idle backends (default True).
+        target: Target name from config; omit for the default.
+    """
+    return ops.list_activity(_get_connection(target), state=state, include_idle=include_idle)
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def long_running_queries(min_seconds: int = 60, target: Optional[str] = None) -> dict:
+    """[READ] Active queries running at least ``min_seconds``, oldest first.
+
+    Args:
+        min_seconds: Minimum age in seconds (default 60).
+        target: Target name from config; omit for the default.
+    """
+    return ops.long_running_queries(_get_connection(target), min_seconds=min_seconds)
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def list_locks(target: Optional[str] = None) -> dict:
+    """[READ] Held/awaited locks joined to their owning backend and object.
+
+    Args:
+        target: Target name from config; omit for the default.
+    """
+    return ops.list_locks(_get_connection(target))

mcp_server/tools/analysis.py

@@ -0,0 +1,89 @@
+"""Flagship PostgreSQL analysis MCP tools (read-only)."""
+
+from typing import Any, Optional
+
+from mcp_server._shared import _get_connection, mcp, tool_errors
+from postgres_aiops.governance import governed_tool
+from postgres_aiops.ops import activity as activity_ops
+from postgres_aiops.ops import analysis as ops
+from postgres_aiops.ops import queries as query_ops
+from postgres_aiops.ops import tables as table_ops
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def slow_query_rca(
+    statements: Optional[list[dict[str, Any]]] = None,
+    explain_sql: Optional[str] = None,
+    limit: int = 20,
+    target: Optional[str] = None,
+) -> dict:
+    """[READ] RCA for the worst pg_stat_statements entry, with cause + action.
+
+    Picks the statement with the greatest total execution time and maps its
+    numbers (mean time, cache-hit ratio, temp spill, calls) — plus an optional
+    EXPLAIN plan — to cited causes and concrete actions. Pass 'statements' for
+    pure/offline analysis, or omit to pull the top statements live.
+
+    Args:
+        statements: Injected pg_stat_statements rows (as from top_queries); if
+            omitted, the worst statements are pulled live.
+        explain_sql: Optional SQL to EXPLAIN so plan node types feed the RCA.
+        limit: How many statements to pull when not injected (default 20).
+        target: Target name from config; omit for the default.
+    """
+    conn = None
+    if statements is None:
+        conn = _get_connection(target)
+        statements = query_ops.top_queries(conn, order_by="total_time", limit=limit)["statements"]
+    explain = None
+    if explain_sql:
+        conn = conn or _get_connection(target)
+        explain = query_ops.explain_query(conn, explain_sql, analyze=False)
+    return ops.slow_query_rca(statements, explain=explain)
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def bloat_and_vacuum_analysis(
+    tables: Optional[list[dict[str, Any]]] = None,
+    limit: int = 50,
+    target: Optional[str] = None,
+) -> dict:
+    """[READ] Rank tables needing vacuum from dead-tuple ratio + autovacuum recency.
+
+    Pass 'tables' (as from table_bloat) for pure/offline analysis, or omit to
+    pull the worst dead-tuple tables live. Each recommendation cites its numbers.
+
+    Args:
+        tables: Injected table-bloat rows; if omitted, pulled live.
+        limit: How many tables to pull when not injected (default 50).
+        target: Target name from config; omit for the default.
+    """
+    if tables is None:
+        tables = table_ops.table_bloat(_get_connection(target), limit=limit)["tables"]
+    return ops.bloat_and_vacuum_analysis(tables)
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def blocking_lock_chain_rca(
+    pairs: Optional[list[dict[str, Any]]] = None,
+    target: Optional[str] = None,
+) -> dict:
+    """[READ] Build the wait-for tree from blocking pairs and name the root blocker.
+
+    Pass 'pairs' (as from the live blocking-pairs read) for pure/offline analysis,
+    or omit to pull the current blocking graph live.
+
+    Args:
+        pairs: Injected blocking pairs {blockedPid, blockingPid, ...}; if omitted,
+            pulled live from pg_blocking_pids.
+        target: Target name from config; omit for the default.
+    """
+    if pairs is None:
+        pairs = activity_ops.blocking_pairs(_get_connection(target))
+    return ops.blocking_lock_chain_rca(pairs)

mcp_server/tools/indexes.py

@@ -0,0 +1,65 @@
+"""Index-health PostgreSQL MCP tools (read-only)."""
+
+from typing import Optional
+
+from mcp_server._shared import _get_connection, mcp, tool_errors
+from postgres_aiops.governance import governed_tool
+from postgres_aiops.ops import indexes as ops
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def unused_indexes(min_size_bytes: int = 0, target: Optional[str] = None) -> dict:
+    """[READ] Non-unique, non-primary indexes with zero scans (drop candidates).
+
+    Args:
+        min_size_bytes: Only report indexes at least this large (default 0).
+        target: Target name from config; omit for the default.
+    """
+    return ops.unused_indexes(_get_connection(target), min_size_bytes=min_size_bytes)
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def missing_index_hints(
+    min_seq_scan: int = 1000,
+    min_live_tup: int = 10000,
+    target: Optional[str] = None,
+) -> dict:
+    """[READ] Tables with heavy sequential scans and few index scans (index hints).
+
+    Args:
+        min_seq_scan: Minimum cumulative sequential scans to flag (default 1000).
+        min_live_tup: Minimum live tuples for a table to qualify (default 10000).
+        target: Target name from config; omit for the default.
+    """
+    return ops.missing_index_hints(
+        _get_connection(target), min_seq_scan=min_seq_scan, min_live_tup=min_live_tup
+    )
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def index_bloat(limit: int = 50, target: Optional[str] = None) -> dict:
+    """[READ] Coarse index-bloat estimate (all inputs returned for transparency).
+
+    Args:
+        limit: Number of indexes to inspect, largest first (default 50).
+        target: Target name from config; omit for the default.
+    """
+    return ops.index_bloat(_get_connection(target), limit=limit)
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def invalid_indexes(target: Optional[str] = None) -> dict:
+    """[READ] Invalid indexes (failed CONCURRENTLY builds) and duplicate indexes.
+
+    Args:
+        target: Target name from config; omit for the default.
+    """
+    return ops.invalid_indexes(_get_connection(target))

mcp_server/tools/queries.py

@@ -0,0 +1,60 @@
+"""Query-statistics PostgreSQL MCP tools: top-N, EXPLAIN (read) + stats reset (write)."""
+
+from typing import Optional
+
+from mcp_server._shared import _get_connection, mcp, tool_errors
+from postgres_aiops.governance import governed_tool
+from postgres_aiops.ops import queries as ops
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def top_queries(
+    order_by: str = "total_time",
+    limit: int = 20,
+    target: Optional[str] = None,
+) -> dict:
+    """[READ] Top statements from pg_stat_statements by a whitelisted metric.
+
+    Args:
+        order_by: One of total_time, mean_time, calls, rows, io.
+        limit: Number of statements to return (1..200, default 20).
+        target: Target name from config; omit for the default.
+    """
+    return ops.top_queries(_get_connection(target), order_by=order_by, limit=limit)
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def explain_query(sql: str, analyze: bool = False, target: Optional[str] = None) -> dict:
+    """[READ] Return the JSON execution plan for ``sql`` (EXPLAIN).
+
+    analyze=False (default) plans without executing; analyze=True runs the
+    statement to collect real timing — only use it for read-only SQL.
+
+    Args:
+        sql: A single SQL statement to EXPLAIN.
+        analyze: If True, execute the statement to gather real row counts/timing.
+        target: Target name from config; omit for the default.
+    """
+    return ops.explain_query(_get_connection(target), sql, analyze=analyze)
+
+
+@mcp.tool()
+@governed_tool(risk_level="medium")
+@tool_errors("dict")
+def reset_query_stats(dry_run: bool = False, target: Optional[str] = None) -> dict:
+    """[WRITE][risk=medium] Reset pg_stat_statements accumulators (irreversible).
+
+    The counters cannot be restored, so no undo is recorded. Pass dry_run=True
+    to preview.
+
+    Args:
+        dry_run: If True, preview without resetting.
+        target: Target name from config; omit for the default.
+    """
+    if dry_run:
+        return {"dryRun": True, "wouldReset": "pg_stat_statements"}
+    return ops.reset_query_stats(_get_connection(target))