#!/usr/bin/env python3 """ pg_plan_cache Agent — Hugging Face Spaces (Gradio UI) Provides a web interface for managing and exploring the pg_plan_cache PostgreSQL extension. Works in two modes: Demo mode — query normalization, hashing, architecture docs (no backend needed) Live mode — full cache stats, Redis inspection, diagnostics (needs PG + Redis) """ import os import gradio as gr from normalizer import normalize_query, compute_query_hash # --------------------------------------------------------------------------- # Optional live DB connection (only when credentials are provided) # --------------------------------------------------------------------------- db = None def _connect_db(): """Try to connect using env vars / HF Secrets. Returns DatabaseManager or None.""" global db if db is not None: return db pg_host = os.getenv("PG_HOST") redis_host = os.getenv("REDIS_HOST") if not pg_host and not redis_host: return None try: from config import load_config from db import DatabaseManager config = load_config() db = DatabaseManager(config) db.pg_query("SELECT 1") return db except Exception: return None def _is_live(): return _connect_db() is not None # --------------------------------------------------------------------------- # Tab 1: Query Normalizer & Hasher (always works — no backend needed) # --------------------------------------------------------------------------- def normalize_tab(query: str): """Normalize a SQL query and compute its SHA-256 cache key.""" if not query.strip(): return "", "", "" normalized = normalize_query(query) qhash = compute_query_hash(normalized) # Check cache status if live status = "N/A (demo mode — no Redis connected)" conn = _connect_db() if conn: try: cached = conn.redis_get(f"plan:{qhash}") status = "CACHED" if cached else "NOT CACHED" except Exception as e: status = f"Error: {e}" return normalized, qhash, status # --------------------------------------------------------------------------- # Tab 2: Cache Statistics (needs live PG) # --------------------------------------------------------------------------- def stats_tab(): conn = _connect_db() if not conn: return "Not connected to PostgreSQL.\nSet PG_HOST, PG_USER, etc. in Space Secrets to enable." try: import tools return tools.get_cache_stats(conn) except Exception as e: return f"Error: {e}" # --------------------------------------------------------------------------- # Tab 3: Cache Efficiency Analysis (needs live PG + Redis) # --------------------------------------------------------------------------- def analyze_tab(): conn = _connect_db() if not conn: return "Not connected.\nSet PG and Redis credentials in Space Secrets to enable." try: import tools return tools.analyze_cache_efficiency(conn) except Exception as e: return f"Error: {e}" # --------------------------------------------------------------------------- # Tab 4: Full Diagnostics (needs live PG + Redis) # --------------------------------------------------------------------------- def diagnose_tab(): conn = _connect_db() if not conn: return "Not connected.\nSet PG and Redis credentials in Space Secrets to enable." try: import tools return tools.diagnose(conn) except Exception as e: return f"Error: {e}" # --------------------------------------------------------------------------- # Tab 5: Redis Health (needs Redis) # --------------------------------------------------------------------------- def redis_health_tab(): conn = _connect_db() if not conn: return "Not connected.\nSet REDIS_HOST, REDIS_PORT, etc. in Space Secrets to enable." try: import tools return tools.check_redis_health(conn) except Exception as e: return f"Error: {e}" # --------------------------------------------------------------------------- # Tab 6: List Cached Plans (needs Redis) # --------------------------------------------------------------------------- def list_plans_tab(limit: int): conn = _connect_db() if not conn: return "Not connected.\nSet PG and Redis credentials in Space Secrets to enable." try: import tools return tools.list_cached_plans(conn, int(limit)) except Exception as e: return f"Error: {e}" # --------------------------------------------------------------------------- # Tab 7: Inspect Plan (needs Redis) # --------------------------------------------------------------------------- def inspect_plan_tab(query_hash: str): conn = _connect_db() if not conn: return "Not connected.\nSet Redis credentials in Space Secrets to enable." if not query_hash.strip(): return "Enter a query hash." try: import tools return tools.get_cached_plan_detail(conn, query_hash.strip()) except Exception as e: return f"Error: {e}" # --------------------------------------------------------------------------- # Tab 8: Table Dependencies (needs Redis) # --------------------------------------------------------------------------- def deps_tab(table_name: str, query_hash: str): conn = _connect_db() if not conn: return "Not connected.\nSet Redis credentials in Space Secrets to enable." try: import tools return tools.get_table_dependencies( conn, table_name.strip() or None, query_hash.strip() or None, ) except Exception as e: return f"Error: {e}" # --------------------------------------------------------------------------- # Tab 9: Extension Config (needs PG) # --------------------------------------------------------------------------- def config_tab(): conn = _connect_db() if not conn: return "Not connected.\nSet PG credentials in Space Secrets to enable." try: import tools return tools.get_extension_config(conn) except Exception as e: return f"Error: {e}" # --------------------------------------------------------------------------- # Tab 10: Invalidate (needs PG) # --------------------------------------------------------------------------- def invalidate_tab(table_name: str): conn = _connect_db() if not conn: return "Not connected.\nSet PG credentials in Space Secrets to enable." try: import tools if table_name.strip(): return tools.invalidate_table_plans(conn, table_name.strip()) else: return tools.invalidate_all_plans(conn) except Exception as e: return f"Error: {e}" # --------------------------------------------------------------------------- # Tab 11: Run SQL (needs PG) # --------------------------------------------------------------------------- def sql_tab(sql: str): conn = _connect_db() if not conn: return "Not connected.\nSet PG credentials in Space Secrets to enable." if not sql.strip(): return "Enter a SELECT query." try: import tools return tools.run_sql_query(conn, sql.strip()) except Exception as e: return f"Error: {e}" # --------------------------------------------------------------------------- # Architecture reference (always works) # --------------------------------------------------------------------------- ARCHITECTURE_MD = """ ## pg_plan_cache — Architecture ### How It Works ``` Query arrives at PostgreSQL │ ▼ ┌─────────────────────────┐ │ planner_hook │ ← pgpc_planner_hook() │ intercepts planning │ └────────┬────────────────┘ │ ▼ ┌─────────────────────────┐ │ normalize_query() │ Replace literals with $N │ compute_query_hash() │ SHA-256 of normalized query └────────┬────────────────┘ │ ▼ ┌─────────────────────────┐ ┌──────────────┐ │ Redis GET plan: │────▶│ Cache HIT? │ └────────┬────────────────┘ └──────┬───────┘ │ │ Cache MISS Return cached │ PlannedStmt ▼ ┌─────────────────────────┐ │ standard_planner() │ Run PostgreSQL planner └────────┬────────────────┘ │ ▼ ┌─────────────────────────┐ │ serialize + Redis SET │ Store plan with TTL │ extract_table_names() │ Track dependencies │ redis_set_dependencies()│ Bidirectional mapping └─────────────────────────┘ ``` ### Key Components | Component | File | Purpose | |-----------|------|---------| | Planner Hook | `pg_plan_cache.c` | Intercepts queries, checks cache, stores plans | | Query Normalizer | `query_normalize.c` | Literals → `$N`, collapse whitespace, lowercase | | Redis Client | `redis_client.c` | Connection pool, GET/SET/DEL, dependency tracking | | Schema Invalidation | `pg_plan_cache.c` | Relcache callback bumps version on DDL | ### Redis Key Patterns | Pattern | Purpose | |---------|---------| | `plan:` | Serialized PlannedStmt + metadata | | `deps:table:` | SET of query hashes depending on this table | | `qdeps:` | SET of table names this query depends on | ### GUC Parameters | Parameter | Default | Restart? | |-----------|---------|----------| | `pg_plan_cache.redis_host` | 127.0.0.1 | No (SIGHUP) | | `pg_plan_cache.redis_port` | 6379 | No (SIGHUP) | | `pg_plan_cache.ttl` | 3600 | No (SIGHUP) | | `pg_plan_cache.enabled` | true | No (SUSET) | | `pg_plan_cache.redis_timeout_ms` | 500 | No (SIGHUP) | | `pg_plan_cache.redis_pool_size` | 4 | Yes (POSTMASTER) | ### Serialization Format ``` schema_version|total_cost|unix_timestamp|nodeToString(PlannedStmt) ``` Plans with a stale `schema_version` are discarded on read. ### Normalization Examples | Input | Normalized | |-------|-----------| | `SELECT * FROM users WHERE id = 42` | `select * from users where id = $1` | | `SELECT name FROM t WHERE x = 'hello'` | `select name from t where x = $1` | | `SELECT /* comment */ a FROM b` | `select a from b` | """ # --------------------------------------------------------------------------- # Build Gradio app # --------------------------------------------------------------------------- def build_app(): with gr.Blocks( title="pg_plan_cache Agent", theme=gr.themes.Soft(), ) as app: gr.Markdown("# pg_plan_cache Agent") gr.Markdown( "Manage, monitor, and diagnose the **pg_plan_cache** PostgreSQL extension. " "The **Normalizer** and **Architecture** tabs work without any backend. " "All other tabs require PostgreSQL and Redis credentials set as environment variables." ) live_status = "Connected" if _is_live() else "Demo mode (no backend)" gr.Markdown(f"**Status:** {live_status}") with gr.Tab("Normalizer"): gr.Markdown("Normalize a SQL query and compute its SHA-256 cache key.") query_input = gr.Textbox( label="SQL Query", placeholder="SELECT * FROM users WHERE id = 42 AND name = 'Alice'", lines=3, ) norm_btn = gr.Button("Normalize & Hash", variant="primary") norm_output = gr.Textbox(label="Normalized Query", lines=2) hash_output = gr.Textbox(label="Cache Key (SHA-256)") status_output = gr.Textbox(label="Cache Status") norm_btn.click( fn=normalize_tab, inputs=[query_input], outputs=[norm_output, hash_output, status_output], ) with gr.Tab("Cache Stats"): gr.Markdown("Live cache hit/miss/error statistics from `pg_plan_cache_stats()`.") stats_btn = gr.Button("Refresh Stats", variant="primary") stats_output = gr.Textbox(label="Statistics", lines=12) stats_btn.click(fn=stats_tab, outputs=[stats_output]) with gr.Tab("Analyze"): gr.Markdown("Cache efficiency analysis with tuning recommendations.") analyze_btn = gr.Button("Run Analysis", variant="primary") analyze_output = gr.Textbox(label="Analysis", lines=20) analyze_btn.click(fn=analyze_tab, outputs=[analyze_output]) with gr.Tab("Diagnose"): gr.Markdown("Full diagnostic: PostgreSQL, Redis, stats, and configuration.") diag_btn = gr.Button("Run Diagnostics", variant="primary") diag_output = gr.Textbox(label="Diagnostic Report", lines=25) diag_btn.click(fn=diagnose_tab, outputs=[diag_output]) with gr.Tab("Redis Health"): gr.Markdown("Redis connectivity, memory, and keyspace info.") redis_btn = gr.Button("Check Redis", variant="primary") redis_output = gr.Textbox(label="Redis Health", lines=12) redis_btn.click(fn=redis_health_tab, outputs=[redis_output]) with gr.Tab("Cached Plans"): gr.Markdown("List cached plan keys stored in Redis.") plan_limit = gr.Slider(minimum=5, maximum=200, value=50, step=5, label="Limit") plans_btn = gr.Button("List Plans", variant="primary") plans_output = gr.Textbox(label="Plans", lines=15) plans_btn.click(fn=list_plans_tab, inputs=[plan_limit], outputs=[plans_output]) with gr.Tab("Inspect Plan"): gr.Markdown("View details of a specific cached plan by hash.") hash_input = gr.Textbox(label="Query Hash (64-char hex)", placeholder="abc123...") inspect_btn = gr.Button("Inspect", variant="primary") inspect_output = gr.Textbox(label="Plan Details", lines=12) inspect_btn.click(fn=inspect_plan_tab, inputs=[hash_input], outputs=[inspect_output]) with gr.Tab("Dependencies"): gr.Markdown("Explore bidirectional dependency mappings between queries and tables.") dep_table = gr.Textbox(label="Table Name", placeholder="public.users") dep_hash = gr.Textbox(label="Query Hash (optional)", placeholder="abc123...") deps_btn = gr.Button("Look Up", variant="primary") deps_output = gr.Textbox(label="Dependencies", lines=12) deps_btn.click(fn=deps_tab, inputs=[dep_table, dep_hash], outputs=[deps_output]) with gr.Tab("Configuration"): gr.Markdown("View current GUC parameter values for pg_plan_cache.") cfg_btn = gr.Button("Show Config", variant="primary") cfg_output = gr.Textbox(label="Configuration", lines=8) cfg_btn.click(fn=config_tab, outputs=[cfg_output]) with gr.Tab("Invalidate"): gr.Markdown("Invalidate cached plans. Leave table name empty to invalidate ALL plans.") inv_table = gr.Textbox(label="Table Name (optional)", placeholder="public.users") inv_btn = gr.Button("Invalidate", variant="stop") inv_output = gr.Textbox(label="Result") inv_btn.click(fn=invalidate_tab, inputs=[inv_table], outputs=[inv_output]) with gr.Tab("SQL Query"): gr.Markdown("Execute a read-only SELECT query against PostgreSQL.") sql_input = gr.Textbox( label="SQL (SELECT only)", placeholder="SELECT * FROM pg_stat_activity LIMIT 5", lines=3, ) sql_btn = gr.Button("Execute", variant="primary") sql_output = gr.Textbox(label="Results", lines=15) sql_btn.click(fn=sql_tab, inputs=[sql_input], outputs=[sql_output]) with gr.Tab("Architecture"): gr.Markdown(ARCHITECTURE_MD) return app # --------------------------------------------------------------------------- # Entry point # --------------------------------------------------------------------------- if __name__ == "__main__": app = build_app() app.launch()