postgres-aiops 0.1.0__py3-none-any.whl

mcp_server/tools/remediation.py

@@ -0,0 +1,281 @@
+"""PostgreSQL maintenance MCP tools (guarded writes).
+
+The state-changing tools. Every one is wrapped with the governance harness
+(audit + graduated approval tier) and takes a ``dry_run`` preview. Reversible
+writes pass an ``undo=`` callback that turns the fetched before-state into an
+inverse descriptor the harness records; irreversible ones record none.
+
+Risk tiers:
+  * terminate_backend / cancel_query / drop_index = high (destructive / irreversible)
+  * run_vacuum / run_analyze / create_index / reindex / update_setting = medium
+"""
+
+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 remediation as ops
+
+# ── undo descriptors (built from the fetched before-state) ──────────────────
+
+
+def _create_index_undo(params: dict[str, Any], result: Any) -> Optional[dict]:
+    if not isinstance(result, dict) or not result.get("index"):
+        return None
+    return {
+        "tool": "drop_index",
+        "params": {"name": result["index"]},
+        "skill": "postgres-aiops",
+        "note": "Inverse of create_index: drop the index that was just created.",
+    }
+
+
+def _drop_index_undo(params: dict[str, Any], result: Any) -> Optional[dict]:
+    if not isinstance(result, dict):
+        return None
+    indexdef = (result.get("priorState") or {}).get("indexdef")
+    if not indexdef:
+        return None
+    return {
+        "tool": "create_index",
+        "params": {"definition": indexdef},
+        "skill": "postgres-aiops",
+        "note": (
+            "Inverse of drop_index: recreate the index from its captured "
+            "definition (replay this CREATE INDEX statement)."
+        ),
+    }
+
+
+def _update_setting_undo(params: dict[str, Any], result: Any) -> Optional[dict]:
+    if not isinstance(result, dict):
+        return None
+    prior = (result.get("priorState") or {}).get("value")
+    if prior is None or prior == "":
+        return None
+    return {
+        "tool": "update_setting",
+        "params": {"name": params.get("name"), "value": prior},
+        "skill": "postgres-aiops",
+        "note": "Inverse of update_setting: ALTER SYSTEM SET back to the prior value.",
+    }
+
+
+# ── activity control (high; irreversible) ───────────────────────────────────
+
+
+@mcp.tool()
+@governed_tool(risk_level="high")
+@tool_errors("dict")
+def terminate_backend(pid: int, dry_run: bool = False, target: Optional[str] = None) -> dict:
+    """[WRITE][risk=high] Terminate a backend (pg_terminate_backend). No safe inverse.
+
+    Captures the backend's pid + query for the audit trail; a terminate cannot be
+    undone, so no undo is offered. Pass dry_run=True to preview.
+
+    Args:
+        pid: Backend process id (from list_activity).
+        dry_run: If True, preview without terminating.
+        target: Target name from config; omit for the default.
+    """
+    conn = _get_connection(target)
+    if dry_run:
+        return {"dryRun": True, "wouldTerminate": {"pid": pid}}
+    return ops.terminate_backend(conn, pid)
+
+
+@mcp.tool()
+@governed_tool(risk_level="high")
+@tool_errors("dict")
+def cancel_query(pid: int, dry_run: bool = False, target: Optional[str] = None) -> dict:
+    """[WRITE][risk=high] Cancel a backend's running query (pg_cancel_backend). No inverse.
+
+    Captures the backend's pid + query for audit; a cancel has no undo. Pass
+    dry_run=True to preview.
+
+    Args:
+        pid: Backend process id (from list_activity).
+        dry_run: If True, preview without cancelling.
+        target: Target name from config; omit for the default.
+    """
+    conn = _get_connection(target)
+    if dry_run:
+        return {"dryRun": True, "wouldCancel": {"pid": pid}}
+    return ops.cancel_query(conn, pid)
+
+
+# ── vacuum / analyze (medium; irreversible, record prior stats) ─────────────
+
+
+@mcp.tool()
+@governed_tool(risk_level="medium")
+@tool_errors("dict")
+def run_vacuum(
+    table: str,
+    full: bool = False,
+    analyze: bool = False,
+    dry_run: bool = False,
+    target: Optional[str] = None,
+) -> dict:
+    """[WRITE][risk=medium] VACUUM a table (optionally FULL/ANALYZE). Records prior stats.
+
+    No undo (a vacuum has no inverse); the prior dead-tuple/last-vacuum stats are
+    captured for the audit trail. VACUUM FULL takes an exclusive lock. Pass
+    dry_run=True to preview.
+
+    Args:
+        table: Table name (optionally schema-qualified, e.g. public.orders).
+        full: Run VACUUM FULL (rewrites the table, exclusive lock).
+        analyze: Also refresh planner statistics (VACUUM ANALYZE).
+        dry_run: If True, preview without running.
+        target: Target name from config; omit for the default.
+    """
+    conn = _get_connection(target)
+    if dry_run:
+        return {"dryRun": True, "wouldVacuum": {"table": table, "full": full, "analyze": analyze}}
+    return ops.run_vacuum(conn, table, full=full, analyze=analyze)
+
+
+@mcp.tool()
+@governed_tool(risk_level="medium")
+@tool_errors("dict")
+def run_analyze(table: str, dry_run: bool = False, target: Optional[str] = None) -> dict:
+    """[WRITE][risk=medium] ANALYZE a table to refresh planner statistics.
+
+    No undo; captures prior stats for audit. Pass dry_run=True to preview.
+
+    Args:
+        table: Table name (optionally schema-qualified).
+        dry_run: If True, preview without running.
+        target: Target name from config; omit for the default.
+    """
+    conn = _get_connection(target)
+    if dry_run:
+        return {"dryRun": True, "wouldAnalyze": {"table": table}}
+    return ops.run_analyze(conn, table)
+
+
+# ── index create/drop/reindex ───────────────────────────────────────────────
+
+
+@mcp.tool()
+@governed_tool(risk_level="medium", undo=_create_index_undo)
+@tool_errors("dict")
+def create_index(
+    table: str,
+    columns: list[str],
+    name: Optional[str] = None,
+    unique: bool = False,
+    concurrently: bool = False,
+    method: Optional[str] = None,
+    dry_run: bool = False,
+    target: Optional[str] = None,
+) -> dict:
+    """[WRITE][risk=medium] Create an index. Reversible: undo drops the created name.
+
+    Supports CONCURRENTLY (non-blocking build). The created name is returned so
+    the harness records an undo that drops exactly this index. Pass dry_run=True
+    to preview.
+
+    Args:
+        table: Table to index (optionally schema-qualified).
+        columns: Column names to index.
+        name: Index name (auto-generated from table+columns when omitted).
+        unique: Create a UNIQUE index.
+        concurrently: Build with CONCURRENTLY (no table lock).
+        method: Index method — btree/hash/gist/gin/brin/spgist (default btree).
+        dry_run: If True, preview without creating.
+        target: Target name from config; omit for the default.
+    """
+    conn = _get_connection(target)
+    if dry_run:
+        return {"dryRun": True, "wouldCreate": {"table": table, "columns": columns, "name": name}}
+    return ops.create_index(
+        conn, table, columns, name=name, unique=unique,
+        concurrently=concurrently, method=method,
+    )
+
+
+@mcp.tool()
+@governed_tool(risk_level="high", undo=_drop_index_undo)
+@tool_errors("dict")
+def drop_index(
+    name: str,
+    concurrently: bool = False,
+    dry_run: bool = False,
+    target: Optional[str] = None,
+) -> dict:
+    """[WRITE][risk=high] Drop an index. Reversible: captures pg_get_indexdef first.
+
+    Before dropping, the exact index definition is captured so the harness records
+    an undo that recreates it. Pass dry_run=True to preview.
+
+    Args:
+        name: Index name (optionally schema-qualified).
+        concurrently: Drop with CONCURRENTLY (no table lock).
+        dry_run: If True, preview without dropping.
+        target: Target name from config; omit for the default.
+    """
+    conn = _get_connection(target)
+    if dry_run:
+        return {"dryRun": True, "wouldDrop": {"name": name}}
+    return ops.drop_index(conn, name, concurrently=concurrently)
+
+
+@mcp.tool()
+@governed_tool(risk_level="medium")
+@tool_errors("dict")
+def reindex(
+    target_name: str,
+    kind: str = "INDEX",
+    concurrently: bool = False,
+    dry_run: bool = False,
+    target: Optional[str] = None,
+) -> dict:
+    """[WRITE][risk=medium] REINDEX an index/table/schema (rebuild in place, no undo).
+
+    Rebuilds physical index storage; there is no inverse. Pass dry_run=True to
+    preview.
+
+    Args:
+        target_name: The index/table/schema name to reindex.
+        kind: INDEX, TABLE, or SCHEMA (default INDEX).
+        concurrently: Rebuild with CONCURRENTLY.
+        dry_run: If True, preview without rebuilding.
+        target: Target name from config; omit for the default.
+    """
+    conn = _get_connection(target)
+    if dry_run:
+        return {"dryRun": True, "wouldReindex": {"kind": kind, "target": target_name}}
+    return ops.reindex(conn, target_name, kind=kind, concurrently=concurrently)
+
+
+# ── server settings (medium; reversible via ALTER SYSTEM) ───────────────────
+
+
+@mcp.tool()
+@governed_tool(risk_level="medium", undo=_update_setting_undo)
+@tool_errors("dict")
+def update_setting(
+    name: str,
+    value: str,
+    dry_run: bool = False,
+    target: Optional[str] = None,
+) -> dict:
+    """[WRITE][risk=medium] ALTER SYSTEM SET a parameter. Reversible: captures prior value.
+
+    Writes postgresql.auto.conf; most parameters need SELECT pg_reload_conf() (or
+    a restart for postmaster-context settings) to take effect — reported but NOT
+    performed automatically. The prior value is captured so the harness records an
+    undo that sets it back. Pass dry_run=True to preview.
+
+    Args:
+        name: The configuration parameter name (e.g. work_mem).
+        value: The new value (as a string).
+        dry_run: If True, preview without changing.
+        target: Target name from config; omit for the default.
+    """
+    conn = _get_connection(target)
+    if dry_run:
+        return {"dryRun": True, "wouldSet": {"name": name, "value": value}}
+    return ops.update_setting(conn, name, value)

mcp_server/tools/replication.py

@@ -0,0 +1,43 @@
+"""Replication 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 replication as ops
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def replication_status(target: Optional[str] = None) -> dict:
+    """[READ] Connected standbys and their replay lag (pg_stat_replication).
+
+    Args:
+        target: Target name from config; omit for the default.
+    """
+    return ops.replication_status(_get_connection(target))
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def replication_slots(target: Optional[str] = None) -> dict:
+    """[READ] Replication slots, flagging inactive slots that retain WAL.
+
+    Args:
+        target: Target name from config; omit for the default.
+    """
+    return ops.replication_slots(_get_connection(target))
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def wal_status(target: Optional[str] = None) -> dict:
+    """[READ] WAL position, level, size settings and archiver health.
+
+    Args:
+        target: Target name from config; omit for the default.
+    """
+    return ops.wal_status(_get_connection(target))

mcp_server/tools/server.py

@@ -0,0 +1,86 @@
+"""Server-level PostgreSQL MCP tools (read-only): overview + catalog reads."""
+
+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 overview as overview_ops
+from postgres_aiops.ops import server as ops
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def overview(target: Optional[str] = None) -> dict:
+    """[READ] One-shot cluster health snapshot.
+
+    Version + uptime, connections by state, idle-in-transaction count, the
+    longest-running query, the worst dead-tuple table, and standby replay lag —
+    each section captured defensively so one failing probe does not sink the rest.
+
+    Args:
+        target: Target name from config; omit for the default.
+    """
+    return overview_ops.snapshot(_get_connection(target))
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def server_version(target: Optional[str] = None) -> dict:
+    """[READ] Server version, uptime, recovery state, and data directory.
+
+    Args:
+        target: Target name from config; omit for the default.
+    """
+    return ops.server_version(_get_connection(target))
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("list")
+def show_settings(pattern: Optional[str] = None, target: Optional[str] = None) -> list:
+    """[READ] Configuration parameters from pg_settings.
+
+    Args:
+        pattern: Optional case-insensitive substring to filter setting names
+            (e.g. 'work_mem', 'autovacuum').
+        target: Target name from config; omit for the default.
+    """
+    return ops.show_settings(_get_connection(target), pattern)
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("list")
+def list_extensions(target: Optional[str] = None) -> list:
+    """[READ] Installed extensions and whether a newer version is available.
+
+    Args:
+        target: Target name from config; omit for the default.
+    """
+    return ops.list_extensions(_get_connection(target))
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("list")
+def list_databases(target: Optional[str] = None) -> list:
+    """[READ] Databases with owner, encoding and on-disk size (largest first).
+
+    Args:
+        target: Target name from config; omit for the default.
+    """
+    return ops.list_databases(_get_connection(target))
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("list")
+def list_roles(target: Optional[str] = None) -> list:
+    """[READ] Roles and their attributes (superuser/login/replication/…).
+
+    Args:
+        target: Target name from config; omit for the default.
+    """
+    return ops.list_roles(_get_connection(target))

mcp_server/tools/tables.py

@@ -0,0 +1,46 @@
+"""Table-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 tables as ops
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def table_sizes(limit: int = 20, target: Optional[str] = None) -> dict:
+    """[READ] Largest tables by total relation size (table + indexes + TOAST).
+
+    Args:
+        limit: Number of tables to return, largest first (default 20).
+        target: Target name from config; omit for the default.
+    """
+    return ops.table_sizes(_get_connection(target), limit=limit)
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def table_bloat(limit: int = 50, target: Optional[str] = None) -> dict:
+    """[READ] Dead-tuple bloat proxy per table (dead / (live + dead)), worst first.
+
+    Args:
+        limit: Number of tables to inspect (default 50).
+        target: Target name from config; omit for the default.
+    """
+    return ops.table_bloat(_get_connection(target), limit=limit)
+
+
+@mcp.tool()
+@governed_tool(risk_level="low")
+@tool_errors("dict")
+def autovacuum_status(limit: int = 50, target: Optional[str] = None) -> dict:
+    """[READ] Per-table dead tuples, mods-since-analyze, and last (auto)vacuum times.
+
+    Args:
+        limit: Number of tables to inspect (default 50).
+        target: Target name from config; omit for the default.
+    """
+    return ops.autovacuum_status(_get_connection(target), limit=limit)

postgres_aiops/__init__.py

@@ -0,0 +1,9 @@
+"""postgres-aiops — governed PostgreSQL DBA operations for AI agents.
+
+Standalone and self-contained: the governance harness (audit, token budget,
+undo-token recording, graduated risk tiers, prompt-injection sanitize) is
+bundled under ``postgres_aiops.governance`` — this package has no external
+skill-family dependency. Preview: not yet full-coverage.
+"""
+
+__version__ = "0.1.0"

postgres_aiops/cli/__init__.py

@@ -0,0 +1,9 @@
+"""CLI package for postgres-aiops.
+
+Re-exports ``app`` so the pyproject entry point
+``postgres-aiops = "postgres_aiops.cli:app"`` works unchanged.
+"""
+
+from postgres_aiops.cli._root import app
+
+__all__ = ["app"]

postgres_aiops/cli/_common.py

@@ -0,0 +1,78 @@
+"""Shared helpers for postgres-aiops CLI sub-modules."""
+
+from __future__ import annotations
+
+import functools
+from collections.abc import Callable
+from pathlib import Path
+from typing import Annotated, Any
+
+import typer
+from rich.console import Console
+
+console = Console()
+
+# ─── Shared Option types ───────────────────────────────────────────────────
+
+TargetOption = Annotated[
+    str | None, typer.Option("--target", "-t", help="Target name from config")
+]
+DryRunOption = Annotated[
+    bool, typer.Option("--dry-run", help="Print the API call without executing")
+]
+
+
+def _cli_error_types() -> tuple[type[BaseException], ...]:
+    """Exceptions translated to a one-line teaching error instead of a traceback."""
+    from postgres_aiops.connection import PgError
+
+    return (PgError, KeyError, OSError, ValueError)
+
+
+def cli_errors(fn: Callable) -> Callable:
+    """Translate known exceptions into one red line + exit code 1."""
+
+    @functools.wraps(fn)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        try:
+            return fn(*args, **kwargs)
+        except (typer.Exit, typer.Abort):
+            raise
+        except _cli_error_types() as e:
+            message = str(e)
+            if isinstance(e, KeyError):
+                message = f"Missing required key or environment variable: {message}"
+            console.print(f"[red]Error: {message}[/]")
+            raise typer.Exit(1) from e
+
+    return wrapper
+
+
+def get_connection(target: str | None, config_path: Path | None = None):
+    """Return a (conn, config) tuple for the given target."""
+    from postgres_aiops.config import load_config
+    from postgres_aiops.connection import ConnectionManager
+
+    cfg = load_config(config_path)
+    mgr = ConnectionManager(cfg)
+    return mgr.connect(target), cfg
+
+
+def dry_run_print(*, operation: str, api_call: str, parameters: dict | None = None) -> None:
+    """Print a dry-run preview of the API call that would be made."""
+    console.print("\n[bold magenta][DRY-RUN] No changes will be made.[/]")
+    console.print(f"[magenta]  Operation: {operation}[/]")
+    console.print(f"[magenta]  API Call:  {api_call}[/]")
+    for k, v in (parameters or {}).items():
+        console.print(f"[magenta]  Param:     {k} = {v}[/]")
+    console.print("[magenta]  Run without --dry-run to execute.[/]\n")
+
+
+def double_confirm(action: str, resource: str) -> None:
+    """Require two confirmations for a destructive operation."""
+    console.print(f"[bold yellow]⚠️  About to: {action} '{resource}'[/]")
+    typer.confirm(f"Confirm 1/2: {action} '{resource}'?", abort=True)
+    typer.confirm(
+        f"Confirm 2/2: really {action} '{resource}'? This may be irreversible.",
+        abort=True,
+    )

postgres_aiops/cli/_root.py

@@ -0,0 +1,68 @@
+"""Top-level Typer app: assembles sub-apps and top-level commands."""
+
+from __future__ import annotations
+
+import typer
+
+from postgres_aiops.cli._common import cli_errors
+from postgres_aiops.cli.activity import activity_app
+from postgres_aiops.cli.analyze import analyze_app
+from postgres_aiops.cli.doctor import doctor_cmd
+from postgres_aiops.cli.index import index_app
+from postgres_aiops.cli.init import init_cmd
+from postgres_aiops.cli.overview import overview_cmd
+from postgres_aiops.cli.query import query_app
+from postgres_aiops.cli.remediate import remediate_app
+from postgres_aiops.cli.replication import repl_app
+from postgres_aiops.cli.secret import secret_app
+from postgres_aiops.cli.server import server_app
+from postgres_aiops.cli.table import table_app
+
+app = typer.Typer(
+    name="postgres-aiops",
+    help="Governed AI-ops for PostgreSQL DBA operations.",
+    no_args_is_help=True,
+)
+
+app.add_typer(server_app, name="server")
+app.add_typer(activity_app, name="activity")
+app.add_typer(query_app, name="query")
+app.add_typer(index_app, name="index")
+app.add_typer(table_app, name="table")
+app.add_typer(repl_app, name="repl")
+app.add_typer(analyze_app, name="analyze")
+app.add_typer(remediate_app, name="remediate")
+app.add_typer(secret_app, name="secret")
+app.command("init")(init_cmd)
+app.command("overview")(overview_cmd)
+app.command("doctor")(doctor_cmd)
+
+
+@app.command("mcp")
+@cli_errors
+def mcp_cmd() -> None:
+    """Start the MCP server (stdio transport).
+
+    Single-command entry point for MCP clients (does not go through uvx/PyPI
+    resolution at launch):
+        postgres-aiops mcp
+    """
+    import sys
+
+    if sys.version_info < (3, 11):
+        typer.echo(
+            f"ERROR: postgres-aiops requires Python >= 3.11 "
+            f"(got {sys.version_info.major}.{sys.version_info.minor}).\n"
+            f"Fix: uv python install 3.12 && "
+            f"uv tool install --python 3.12 --force postgres-aiops",
+            err=True,
+        )
+        raise typer.Exit(2)
+
+    from mcp_server.server import main as _mcp_main
+
+    _mcp_main()
+
+
+if __name__ == "__main__":
+    app()

postgres_aiops/cli/activity.py

@@ -0,0 +1,52 @@
+"""``postgres-aiops activity`` — sessions, long-running queries, locks."""
+
+from __future__ import annotations
+
+import json
+from typing import Annotated
+
+import typer
+
+from postgres_aiops.cli._common import TargetOption, cli_errors, console, get_connection
+
+activity_app = typer.Typer(
+    name="activity",
+    help="Activity: sessions, long-running queries, locks.",
+    no_args_is_help=True,
+)
+
+
+@activity_app.command("list")
+@cli_errors
+def activity_list(
+    state: Annotated[str | None, typer.Option("--state", help="Filter by state")] = None,
+    target: TargetOption = None,
+) -> None:
+    """List current sessions (pg_stat_activity) with per-state counts."""
+    from postgres_aiops.ops import activity as ops
+
+    conn, _ = get_connection(target)
+    console.print_json(json.dumps(ops.list_activity(conn, state=state)))
+
+
+@activity_app.command("long")
+@cli_errors
+def activity_long(
+    min_seconds: Annotated[int, typer.Option("--min-seconds", help="Minimum age")] = 60,
+    target: TargetOption = None,
+) -> None:
+    """List active queries running at least --min-seconds."""
+    from postgres_aiops.ops import activity as ops
+
+    conn, _ = get_connection(target)
+    console.print_json(json.dumps(ops.long_running_queries(conn, min_seconds=min_seconds)))
+
+
+@activity_app.command("locks")
+@cli_errors
+def activity_locks(target: TargetOption = None) -> None:
+    """List held/awaited locks."""
+    from postgres_aiops.ops import activity as ops
+
+    conn, _ = get_connection(target)
+    console.print_json(json.dumps(ops.list_locks(conn)))