75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
|
"""pg_stat_statements deltas and EXPLAIN helpers, per phase-0a-test-suite-spec.md §7."""
|
||
|
|
from contextlib import contextmanager
|
||
|
|
|
||
|
|
import psycopg
|
||
|
|
|
||
|
|
|
||
|
|
class _Counter:
|
||
|
|
"""Populated with `.calls` (and `.total_time_ms`) once the `with` block exits."""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self.calls = None
|
||
|
|
self.total_time_ms = None
|
||
|
|
|
||
|
|
def __repr__(self):
|
||
|
|
return f"Counter(calls={self.calls}, total_time_ms={self.total_time_ms})"
|
||
|
|
|
||
|
|
def __str__(self):
|
||
|
|
return repr(self)
|
||
|
|
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def statement_count(dsn):
|
||
|
|
"""Reset pg_stat_statements, run the caller's block, yield the call-count delta.
|
||
|
|
|
||
|
|
pg_stat_statements is reset globally, so this is only meaningful when
|
||
|
|
nothing else is hitting the database concurrently during the `with` block
|
||
|
|
-- true for the single-request pathology assertions this is built for.
|
||
|
|
"""
|
||
|
|
counter = _Counter()
|
||
|
|
with psycopg.connect(dsn, autocommit=True) as conn:
|
||
|
|
with conn.cursor() as cur:
|
||
|
|
cur.execute("SELECT pg_stat_statements_reset()")
|
||
|
|
try:
|
||
|
|
yield counter
|
||
|
|
finally:
|
||
|
|
with conn.cursor() as cur:
|
||
|
|
cur.execute(
|
||
|
|
"SELECT COALESCE(SUM(calls), 0), COALESCE(SUM(total_exec_time), 0) "
|
||
|
|
"FROM pg_stat_statements"
|
||
|
|
)
|
||
|
|
calls, total_time_ms = cur.fetchone()
|
||
|
|
counter.calls = calls
|
||
|
|
counter.total_time_ms = total_time_ms
|
||
|
|
|
||
|
|
|
||
|
|
def explain(dsn, sql, params=None):
|
||
|
|
"""Run EXPLAIN (ANALYZE, FORMAT JSON) and return the plan dict."""
|
||
|
|
with psycopg.connect(dsn, autocommit=True) as conn:
|
||
|
|
with conn.cursor() as cur:
|
||
|
|
cur.execute(f"EXPLAIN (ANALYZE, FORMAT JSON) {sql}", params or ())
|
||
|
|
(plan,) = cur.fetchone()
|
||
|
|
return plan[0]
|
||
|
|
|
||
|
|
|
||
|
|
def plan_node_types(plan: dict) -> set:
|
||
|
|
"""Flatten an EXPLAIN JSON plan (or its top-level {"Plan": ...} wrapper)
|
||
|
|
into the set of "Node Type" values it contains, at any depth."""
|
||
|
|
types: set = set()
|
||
|
|
root = plan.get("Plan", plan) if isinstance(plan, dict) else plan
|
||
|
|
|
||
|
|
def walk(node):
|
||
|
|
if isinstance(node, dict):
|
||
|
|
if "Node Type" in node:
|
||
|
|
types.add(node["Node Type"])
|
||
|
|
for key, value in node.items():
|
||
|
|
if key == "Plans" and isinstance(value, list):
|
||
|
|
for child in value:
|
||
|
|
walk(child)
|
||
|
|
elif isinstance(node, list):
|
||
|
|
for item in node:
|
||
|
|
walk(item)
|
||
|
|
|
||
|
|
walk(root)
|
||
|
|
return types
|