app.py

#!/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:<hash>   │────▶│  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:<sha256>` | Serialized PlannedStmt + metadata |
| `deps:table:<schema.table>` | SET of query hashes depending on this table |
| `qdeps:<sha256>` | 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()