postgres-aiops 0.1.0__py3-none-any.whl

postgres_aiops/ops/_util.py

@@ -0,0 +1,102 @@
+"""Shared helpers for the PostgreSQL ops modules.
+
+Two jobs:
+
+  * ``s`` — pass catalog/query text through the governance ``sanitize`` (bounded
+    length, prompt-injection defense) before it reaches an agent.
+  * ``qualify`` / ``quote_ident`` — the ONLY sanctioned way to place an
+    identifier (schema/table/index/column) into a statement that cannot be
+    parameterised (DDL, ``VACUUM``, ``ANALYZE``, ``REINDEX``). Every part is
+    validated against strict identifier rules and then double-quoted, so a value
+    that is not a plain identifier is rejected rather than interpolated.
+
+Values (pids, thresholds, limits, setting values) are ALWAYS passed as query
+parameters — never string-formatted into SQL.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+from postgres_aiops.governance import sanitize
+
+# A PostgreSQL unquoted identifier component: letter/underscore then
+# letters/digits/underscore/dollar. We deliberately reject everything else
+# (spaces, quotes, semicolons, operators) so interpolation cannot inject SQL.
+_IDENT_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_$]*$")
+
+# Whitelist of ORDER-BY columns callers may choose from (maps a friendly name to
+# a real column). Used so an ordering choice is never taken from raw user text.
+STATEMENT_ORDER_COLUMNS = {
+    "total_time": "total_exec_time",
+    "mean_time": "mean_exec_time",
+    "calls": "calls",
+    "rows": "rows",
+    "io": "shared_blks_read",
+}
+
+
+def s(value: Any, limit: int = 200) -> str:
+    """Sanitize an arbitrary value to a bounded, injection-safe string."""
+    return sanitize(str(value if value is not None else ""), limit)
+
+
+def quote_ident(part: str) -> str:
+    """Validate a single identifier component and return it double-quoted.
+
+    Raises ``ValueError`` for anything that is not a plain identifier — this is
+    the boundary that makes identifier interpolation safe.
+    """
+    if not isinstance(part, str) or not _IDENT_RE.match(part):
+        raise ValueError(
+            f"Invalid SQL identifier {part!r}: only letters, digits, underscore "
+            f"and '$' are allowed (must start with a letter/underscore)."
+        )
+    return '"' + part + '"'
+
+
+def qualify(name: str) -> str:
+    """Validate + quote a possibly schema-qualified name (``schema.table``).
+
+    ``qualify('public.orders')`` → ``"public"."orders"``; ``qualify('orders')``
+    → ``"orders"``. Each component is validated by :func:`quote_ident`.
+    """
+    if not isinstance(name, str) or not name.strip():
+        raise ValueError("Empty identifier is not allowed.")
+    parts = name.split(".")
+    if len(parts) > 2:
+        raise ValueError(f"Too many name parts in {name!r} (expected schema.table).")
+    return ".".join(quote_ident(p) for p in parts)
+
+
+def quote_literal(value: str) -> str:
+    """Quote a string as a SQL literal (single-quoted, doubling embedded quotes).
+
+    Only used for ``ALTER SYSTEM SET`` where the value cannot be parameterised.
+    """
+    return "'" + str(value).replace("'", "''") + "'"
+
+
+def order_column(choice: str) -> str:
+    """Map a caller's order-by choice to a real column via the whitelist."""
+    col = STATEMENT_ORDER_COLUMNS.get(choice)
+    if col is None:
+        allowed = ", ".join(sorted(STATEMENT_ORDER_COLUMNS))
+        raise ValueError(f"Unknown order_by '{choice}'. Allowed: {allowed}.")
+    return col
+
+
+def human_bytes(n: Any) -> str:
+    """Render a byte count as a human string (e.g. 1536 -> '1.5 kB')."""
+    try:
+        size = float(n)
+    except (TypeError, ValueError):
+        return "0 bytes"
+    for unit in ("bytes", "kB", "MB", "GB", "TB"):
+        if abs(size) < 1024.0 or unit == "TB":
+            if unit == "bytes":
+                return f"{int(size)} {unit}"
+            return f"{size:.1f} {unit}"
+        size /= 1024.0
+    return f"{size:.1f} TB"

postgres_aiops/ops/activity.py

@@ -0,0 +1,193 @@
+"""Activity reads: sessions, long-running queries, locks and blocking pairs.
+
+Every query surfaced to an agent is truncated in SQL (``left(query, N)``) and
+sanitised again via ``s`` — a running statement can contain arbitrary text.
+Thresholds (seconds) are passed as bound parameters, never string-formatted.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from postgres_aiops.ops._util import s
+
+_ACTIVITY_SQL = """
+SELECT pid,
+       usename AS username,
+       datname AS database,
+       client_addr::text AS client_addr,
+       application_name,
+       state,
+       wait_event_type,
+       wait_event,
+       backend_type,
+       xact_start,
+       query_start,
+       state_change,
+       EXTRACT(EPOCH FROM (now() - query_start))::int AS query_age_seconds,
+       EXTRACT(EPOCH FROM (now() - xact_start))::int AS xact_age_seconds,
+       left(query, 500) AS query
+FROM pg_stat_activity
+WHERE pid <> pg_backend_pid()
+  AND (%(state)s IS NULL OR state = %(state)s)
+  AND (%(include_idle)s OR state IS DISTINCT FROM 'idle')
+ORDER BY query_start NULLS LAST
+"""
+
+_LONG_RUNNING_SQL = """
+SELECT pid, usename AS username, datname AS database, state,
+       wait_event_type, wait_event,
+       EXTRACT(EPOCH FROM (now() - query_start))::int AS duration_seconds,
+       (now() - query_start) AS duration,
+       left(query, 500) AS query
+FROM pg_stat_activity
+WHERE state <> 'idle'
+  AND query_start IS NOT NULL
+  AND (now() - query_start) >= make_interval(secs => %(min_seconds)s)
+  AND pid <> pg_backend_pid()
+ORDER BY query_start ASC
+"""
+
+_LOCKS_SQL = """
+SELECT a.pid, a.usename AS username, a.datname AS database,
+       l.locktype, l.mode, l.granted,
+       COALESCE(c.relname, l.locktype) AS object,
+       a.wait_event_type, a.wait_event,
+       left(a.query, 300) AS query
+FROM pg_locks l
+JOIN pg_stat_activity a ON a.pid = l.pid
+LEFT JOIN pg_class c ON c.oid = l.relation
+WHERE a.pid <> pg_backend_pid()
+ORDER BY l.granted, a.pid
+"""
+
+# One row per (blocked backend -> a backend blocking it). pg_blocking_pids gives
+# the exact wait-for edges; unnesting produces the graph the flagship RCA walks.
+_BLOCKING_PAIRS_SQL = """
+SELECT blocked.pid AS blocked_pid,
+       blocked.usename AS blocked_user,
+       blocked.datname AS database,
+       blocked.wait_event_type,
+       blocked.wait_event,
+       left(blocked.query, 300) AS blocked_query,
+       blocking.pid AS blocking_pid,
+       blocking.usename AS blocking_user,
+       blocking.state AS blocking_state,
+       left(blocking.query, 300) AS blocking_query
+FROM pg_stat_activity blocked
+JOIN LATERAL unnest(pg_blocking_pids(blocked.pid)) AS bp(pid) ON true
+JOIN pg_stat_activity blocking ON blocking.pid = bp.pid
+WHERE cardinality(pg_blocking_pids(blocked.pid)) > 0
+"""
+
+
+def _activity_row(r: dict) -> dict:
+    return {
+        "pid": r.get("pid"),
+        "username": s(r.get("username"), 128),
+        "database": s(r.get("database"), 128),
+        "clientAddr": s(r.get("client_addr"), 64),
+        "applicationName": s(r.get("application_name"), 128),
+        "state": s(r.get("state"), 64),
+        "waitEventType": s(r.get("wait_event_type"), 64),
+        "waitEvent": s(r.get("wait_event"), 64),
+        "backendType": s(r.get("backend_type"), 64),
+        "queryAgeSeconds": r.get("query_age_seconds"),
+        "xactAgeSeconds": r.get("xact_age_seconds"),
+        "query": s(r.get("query"), 500),
+    }
+
+
+def list_activity(
+    conn: Any, state: str | None = None, include_idle: bool = True
+) -> dict:
+    """[READ] Current sessions from pg_stat_activity, with per-state counts.
+
+    Flags idle-in-transaction backends (an open transaction holding resources).
+    """
+    rows = conn.query(_ACTIVITY_SQL, {"state": state, "include_idle": include_idle})
+    sessions = [_activity_row(r) for r in rows]
+    by_state: dict[str, int] = {}
+    for r in sessions:
+        key = r["state"] or "unknown"
+        by_state[key] = by_state.get(key, 0) + 1
+    idle_in_txn = [
+        r for r in sessions if r["state"] == "idle in transaction"
+    ]
+    return {
+        "total": len(sessions),
+        "byState": dict(sorted(by_state.items(), key=lambda kv: kv[1], reverse=True)),
+        "idleInTransactionCount": len(idle_in_txn),
+        "idleInTransaction": idle_in_txn,
+        "sessions": sessions,
+    }
+
+
+def long_running_queries(conn: Any, min_seconds: int = 60) -> dict:
+    """[READ] Active queries running at least ``min_seconds``, oldest first."""
+    rows = conn.query(_LONG_RUNNING_SQL, {"min_seconds": int(min_seconds)})
+    queries = [
+        {
+            "pid": r.get("pid"),
+            "username": s(r.get("username"), 128),
+            "database": s(r.get("database"), 128),
+            "state": s(r.get("state"), 64),
+            "durationSeconds": r.get("duration_seconds"),
+            "duration": s(r.get("duration"), 64),
+            "waitEventType": s(r.get("wait_event_type"), 64),
+            "waitEvent": s(r.get("wait_event"), 64),
+            "query": s(r.get("query"), 500),
+        }
+        for r in rows
+    ]
+    return {
+        "thresholdSeconds": int(min_seconds),
+        "count": len(queries),
+        "queries": queries,
+    }
+
+
+def list_locks(conn: Any) -> dict:
+    """[READ] Held/awaited locks joined to their owning backend and object."""
+    rows = conn.query(_LOCKS_SQL)
+    locks = [
+        {
+            "pid": r.get("pid"),
+            "username": s(r.get("username"), 128),
+            "database": s(r.get("database"), 128),
+            "lockType": s(r.get("locktype"), 64),
+            "mode": s(r.get("mode"), 64),
+            "granted": bool(r.get("granted")),
+            "object": s(r.get("object"), 128),
+            "waitEventType": s(r.get("wait_event_type"), 64),
+            "query": s(r.get("query"), 300),
+        }
+        for r in rows
+    ]
+    waiting = [lock for lock in locks if not lock["granted"]]
+    return {
+        "total": len(locks),
+        "waitingCount": len(waiting),
+        "waiting": waiting,
+        "locks": locks,
+    }
+
+
+def blocking_pairs(conn: Any) -> list[dict]:
+    """[READ] Wait-for edges (blocked pid -> blocking pid) from pg_blocking_pids."""
+    rows = conn.query(_BLOCKING_PAIRS_SQL)
+    return [
+        {
+            "blockedPid": r.get("blocked_pid"),
+            "blockedUser": s(r.get("blocked_user"), 128),
+            "database": s(r.get("database"), 128),
+            "waitEventType": s(r.get("wait_event_type"), 64),
+            "waitEvent": s(r.get("wait_event"), 64),
+            "blockedQuery": s(r.get("blocked_query"), 300),
+            "blockingPid": r.get("blocking_pid"),
+            "blockingUser": s(r.get("blocking_user"), 128),
+            "blockingState": s(r.get("blocking_state"), 64),
+            "blockingQuery": s(r.get("blocking_query"), 300),
+        }
+        for r in rows
+    ]

postgres_aiops/ops/analysis.py

@@ -0,0 +1,263 @@
+"""Flagship signature analyses over PostgreSQL telemetry (pure analysis).
+
+The differentiators — transparent heuristics, every flag reported with its number
+so a DBA can see *why* something was ranked, never a black-box verdict:
+
+  1. ``slow_query_rca`` — take the worst pg_stat_statements entry (+ an optional
+     EXPLAIN plan) and map it to a likely cause + concrete action.
+  2. ``bloat_and_vacuum_analysis`` — combine dead-tuple ratio, table size and
+     autovacuum recency into a per-table recommendation.
+  3. ``blocking_lock_chain_rca`` — build the wait-for tree from blocking pairs,
+     name the root blocker, and give the action.
+
+All three are pure functions (no I/O): pass them telemetry (from the reads in the
+other ops modules, or injected) and they return the analysis.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+MAX_ROWS = 100
+
+
+# ── 1. slow query RCA ───────────────────────────────────────────────────────
+# Thresholds that flip a signal on (each reported with its measured number).
+_SLOW_MEAN_MS = 100.0
+_LOW_CACHE_HIT_PCT = 95.0
+_HIGH_CALLS = 100_000
+
+
+def _plan_node_types(plan: Any, found: set[str]) -> None:
+    """Recursively collect ``Node Type`` values from an EXPLAIN JSON plan."""
+    if isinstance(plan, dict):
+        node = plan.get("Node Type")
+        if isinstance(node, str):
+            found.add(node)
+        for value in plan.values():
+            _plan_node_types(value, found)
+    elif isinstance(plan, list):
+        for item in plan:
+            _plan_node_types(item, found)
+
+
+def _slow_findings(worst: dict, node_types: set[str]) -> list[dict]:
+    """Build the list of cited findings (cause + action) for the worst statement."""
+    findings: list[dict] = []
+    mean = float(worst.get("meanExecTimeMs") or 0)
+    cache = worst.get("cacheHitRatioPct")
+    temp_written = worst.get("tempBlksWritten") or 0
+    calls = worst.get("calls") or 0
+
+    if "Seq Scan" in node_types and mean >= _SLOW_MEAN_MS:
+        findings.append({
+            "signal": "sequential scan on a slow statement",
+            "detail": f"plan contains Seq Scan; mean {mean}ms >= {_SLOW_MEAN_MS}ms",
+            "cause": "A hot query scans a table with no usable index.",
+            "action": "Add an index on the filter/join columns; confirm with EXPLAIN.",
+        })
+    if isinstance(cache, (int, float)) and cache < _LOW_CACHE_HIT_PCT:
+        findings.append({
+            "signal": "low shared-buffer cache hit ratio",
+            "detail": f"cacheHitRatioPct {cache}% < {_LOW_CACHE_HIT_PCT}%",
+            "cause": "Reads miss the buffer cache and hit disk repeatedly.",
+            "action": "Add a covering index to cut blocks read, or raise shared_buffers.",
+        })
+    if temp_written and int(temp_written) > 0:
+        findings.append({
+            "signal": "temp blocks written (spill to disk)",
+            "detail": f"tempBlksWritten={temp_written} — sorts/hashes spilled",
+            "cause": "work_mem is too small for this query's sort/hash.",
+            "action": "Raise work_mem for the session/role, or reduce the sorted set.",
+        })
+    if calls and int(calls) >= _HIGH_CALLS:
+        findings.append({
+            "signal": "very high call count",
+            "detail": f"calls={calls} >= {_HIGH_CALLS}",
+            "cause": "A cheap statement is executed enormously often (possible N+1).",
+            "action": "Batch/cache at the application, or use a prepared/set-based query.",
+        })
+    if not findings:
+        findings.append({
+            "signal": "no dominant signal",
+            "detail": f"mean {mean}ms, calls {calls}",
+            "cause": "The statement is costly but shows no single clear driver.",
+            "action": "EXPLAIN (ANALYZE, BUFFERS) it and inspect the most expensive node.",
+        })
+    return findings
+
+
+def slow_query_rca(statements: list[dict], explain: dict | None = None) -> dict:
+    """[READ] RCA for the worst pg_stat_statements entry (+ optional EXPLAIN plan).
+
+    Picks the statement with the greatest total execution time, then maps its
+    numbers (mean time, cache-hit ratio, temp spill, call count) — and any
+    EXPLAIN plan node types supplied — to cited causes and concrete actions.
+    """
+    ranked = sorted(
+        (s for s in (statements or []) if isinstance(s, dict)),
+        key=lambda x: float(x.get("totalExecTimeMs") or 0),
+        reverse=True,
+    )
+    if not ranked:
+        return {"evaluated": 0, "worst": None, "findings": [], "note": "No statements supplied."}
+
+    worst = ranked[0]
+    node_types: set[str] = set()
+    if explain:
+        _plan_node_types(explain.get("plan", explain), node_types)
+    findings = _slow_findings(worst, node_types)
+    return {
+        "evaluated": len(ranked),
+        "worst": {
+            "queryId": worst.get("queryId"),
+            "query": worst.get("query"),
+            "calls": worst.get("calls"),
+            "totalExecTimeMs": worst.get("totalExecTimeMs"),
+            "meanExecTimeMs": worst.get("meanExecTimeMs"),
+            "cacheHitRatioPct": worst.get("cacheHitRatioPct"),
+            "tempBlksWritten": worst.get("tempBlksWritten"),
+        },
+        "planNodeTypes": sorted(node_types),
+        "findings": findings,
+        "note": (
+            "Advisory read-only heuristic over pg_stat_statements; every finding "
+            "cites the measured number. Worst = greatest total execution time."
+        ),
+    }
+
+
+# ── 2. bloat & vacuum analysis ──────────────────────────────────────────────
+_DEAD_PCT_WARN = 20.0
+_DEAD_TUP_WARN = 1000
+_MOD_ANALYZE_WARN = 50_000
+
+
+def _vacuum_recommendation(row: dict) -> dict | None:
+    """Return a cited recommendation for one table, or None if it looks healthy."""
+    dead_pct = float(row.get("deadPct") or 0)
+    dead = int(row.get("deadTuples") or 0)
+    last_autovac = row.get("lastAutovacuum")
+    reasons: list[str] = []
+    action = None
+    if dead_pct >= _DEAD_PCT_WARN and dead >= _DEAD_TUP_WARN:
+        reasons.append(f"deadPct {dead_pct}% >= {_DEAD_PCT_WARN}% ({dead} dead tuples)")
+        action = "Run VACUUM (ANALYZE); if disk must be reclaimed, VACUUM FULL/pg_repack."
+    if not last_autovac and dead >= _DEAD_TUP_WARN:
+        reasons.append(f"never autovacuumed but has {dead} dead tuples")
+        action = action or "Lower autovacuum_vacuum_scale_factor for this table."
+    if not reasons:
+        return None
+    return {
+        "schema": row.get("schema"),
+        "table": row.get("table"),
+        "deadPct": dead_pct,
+        "deadTuples": dead,
+        "sizePretty": row.get("sizePretty"),
+        "lastAutovacuum": last_autovac,
+        "reasons": reasons,
+        "action": action,
+    }
+
+
+def bloat_and_vacuum_analysis(tables: list[dict]) -> dict:
+    """[READ] Rank tables needing vacuum from dead-tuple ratio + autovacuum recency.
+
+    Pure analysis over table-bloat rows ({schema, table, deadPct, deadTuples,
+    sizePretty, lastAutovacuum}). Each recommendation cites the numbers that
+    triggered it; healthy tables are omitted.
+    """
+    recs = [rec for rec in (_vacuum_recommendation(r) for r in (tables or [])) if rec]
+    recs.sort(key=lambda r: r["deadPct"], reverse=True)
+    return {
+        "tablesEvaluated": len(tables or []),
+        "needsAttentionCount": len(recs),
+        "thresholds": {"deadPct": _DEAD_PCT_WARN, "deadTuples": _DEAD_TUP_WARN},
+        "recommendations": recs[:MAX_ROWS],
+        "note": (
+            "Advisory read-only heuristic: flags deadPct >= "
+            f"{_DEAD_PCT_WARN}% with >= {_DEAD_TUP_WARN} dead tuples, or a table "
+            "never autovacuumed that has accumulated dead tuples."
+        ),
+    }
+
+
+# ── 3. blocking lock chain RCA ──────────────────────────────────────────────
+
+
+def _descendants(root: int, children: dict[int, list[int]]) -> set[int]:
+    """All pids transitively blocked by ``root`` (BFS, cycle-safe)."""
+    seen: set[int] = set()
+    stack = list(children.get(root, []))
+    while stack:
+        pid = stack.pop()
+        if pid in seen:
+            continue
+        seen.add(pid)
+        stack.extend(children.get(pid, []))
+    return seen
+
+
+def blocking_lock_chain_rca(pairs: list[dict]) -> dict:
+    """[READ] Build the wait-for tree from blocking pairs and name the root blocker.
+
+    Pure analysis over pairs ({blockedPid, blockingPid, blockedQuery,
+    blockingQuery, ...}). A root blocker blocks others but is itself blocked by
+    none; the worst root is the one with the most transitively-blocked backends.
+    A cycle (everyone blocked) is reported as a likely deadlock.
+    """
+    edges = [
+        p for p in (pairs or [])
+        if isinstance(p, dict) and p.get("blockedPid") and p.get("blockingPid")
+    ]
+    if not edges:
+        return {"blockedBackends": 0, "roots": [], "note": "No blocking detected."}
+
+    children: dict[int, list[int]] = {}
+    blocking_query: dict[int, str] = {}
+    blocked_pids: set[int] = set()
+    blocking_pids: set[int] = set()
+    for e in edges:
+        b, g = e["blockedPid"], e["blockingPid"]
+        children.setdefault(g, []).append(b)
+        blocked_pids.add(b)
+        blocking_pids.add(g)
+        blocking_query.setdefault(g, e.get("blockingQuery") or "")
+
+    root_pids = [pid for pid in blocking_pids if pid not in blocked_pids]
+    if not root_pids:
+        return {
+            "blockedBackends": len(blocked_pids),
+            "roots": [],
+            "deadlockSuspected": True,
+            "note": (
+                "Every blocker is itself blocked — a cycle / possible deadlock. "
+                "Inspect the transactions and terminate one to break the cycle."
+            ),
+        }
+
+    roots = []
+    for pid in root_pids:
+        blocked = _descendants(pid, children)
+        roots.append({
+            "rootPid": pid,
+            "blockedCount": len(blocked),
+            "blockedPids": sorted(blocked),
+            "rootQuery": blocking_query.get(pid, ""),
+            "action": (
+                f"Backend {pid} is the head of the chain — investigate its "
+                "transaction; cancel_query(pid) to stop its statement, or "
+                "terminate_backend(pid) to end the session and release the locks."
+            ),
+        })
+    roots.sort(key=lambda r: r["blockedCount"], reverse=True)
+    return {
+        "blockedBackends": len(blocked_pids),
+        "rootCount": len(roots),
+        "worstRootPid": roots[0]["rootPid"],
+        "roots": roots[:MAX_ROWS],
+        "note": (
+            "Advisory read-only heuristic: a root blocker holds locks others wait "
+            "on but waits on nobody; worst root blocks the most backends."
+        ),
+    }