AI 替查询算命：LLM 报这行 SQL 要慢 10 倍，你敢跑吗？

# Bespoke-Card 三件套的最小可运行骨架
# pip install langgraph typing-extensions

from typing import TypedDict
from langgraph.graph import StateGraph, END


# ===== 状态：贯穿整条工作流的"记忆" =====
class EstimatorState(TypedDict):
    workload: list             # 原始 SQL 查询列表（从 pg_stat_statements 抓）
    planner_output: dict       # planner 拆解结果：要写哪几个函数
    code: str                  # coder 写的 Python estimation 函数源码
    q_error: float             # validator 算出来的 q-error（越接近 1 越好）
    iteration: int             # 反馈循环计数，超过 3 就停


# ===== 节点：planner / coder / validator =====
def planner(state: EstimatorState) -> dict:
    """拆解 workload：分析 SQL 模式，输出要写哪几个 estimation 函数。"""
    # 真实实现里会用 LLM 读 SQL，识别高频 JOIN 模式和高基数列
    return {
        "planner_output": {
            "functions": [
                {"name": "estimate_orders_by_status",
                 "table": "orders",
                 "cols": ["status", "user_id"]},
            ]
        }
    }


def coder(state: EstimatorState) -> dict:
    """根据 planner 输出，写 Python estimation 函数（白盒、可读）。"""
    # 真实实现里会调 LLM 生成函数体，并要求走沙箱 + 单元测试
    return {"code": "def estimate_orders_by_status(...): ..."}


def validator(state: EstimatorState) -> dict:
    """在 EXPLAIN ANALYZE 上跑回归，算 q-error，决定要不要回环。"""
    # 真实实现：拿真实行数 vs estimator 输出行数，算 max(est/act, act/est)
    return {"q_error": 3.2, "iteration": state["iteration"] + 1}


# ===== 条件边：决定回环还是 END =====
def should_continue(state: EstimatorState) -> str:
    if state["q_error"] < 2.0 or state["iteration"] >= 3:
        return END       # 收敛了，退出
    return "coder"       # 没收敛，让 coder 再改一版


# ===== 组装状态机 =====
# 显式声明入口节点（LangGraph 0.2+ 强烈推荐，避免自动推断带来的"幽灵入口"警告）
workflow = StateGraph(EstimatorState)
workflow.set_entry_point("planner")                              # 显式 entry point
workflow.add_node("planner", planner)
workflow.add_node("coder", coder)
workflow.add_node("validator", validator)

workflow.add_edge("planner", "coder")                        # planner → coder
workflow.add_edge("coder", "validator")                      # coder → validator
workflow.add_conditional_edges(
    "validator", should_continue, {"coder": "coder", END: END}
)

app = workflow.compile()
# 首次 invoke 必须给 q_error 一个默认值（否则 should_continue 第一次访问会 KeyError）
# app.invoke({"workload": [...], "iteration": 0, "q_error": float("inf")})  # 跑起来

-- 1. 启用 PL/Python（首次需要）
-- CREATE EXTENSION plpython3u;

-- 2. 把 coder 生成的 Python 函数注册成 PG UDF
CREATE OR REPLACE FUNCTION estimate_orders_by_status(
    p_status text,
    p_user_id bigint
) RETURNS double precision AS $$
    # Bespoke-Card 合成的 estimation 函数
    # 输入：query 中具体的 status / user_id
    # 输出：优化器能拿来算代价的"估计行数"
    if p_status == 'paid' and p_user_id > 1000:
        return 12.0          # 历史看这种组合平均 12 行
    if p_status == 'paid':
        return 1500.0
    return 80000.0
$$ LANGUAGE plpython3u IMMUTABLE PARALLEL SAFE;

-- 3. 让 PG 优化器感知这个谓词的列相关性（PG 14+ 接口）
CREATE STATISTICS orders_status_userid_stats (dependencies)
    ON status, user_id FROM orders;

-- 4. 验证：拿真实 EXPLAIN ANALYZE 行数对比
EXPLAIN ANALYZE
SELECT * FROM orders WHERE status = 'paid' AND user_id = 1234;
-- 优化器现在会调 estimate_orders_by_status('paid', 1234) 拿估计值