LangGraph PostgresSaver Connection & Migration Errors

# ✅ One-off setup script — run before first deploy of a new database
from langgraph.checkpoint.postgres import PostgresSaver

DATABASE_URL = "postgresql://user:pass@host:5432/langgraph"

with PostgresSaver.from_conn_string(DATABASE_URL) as checkpointer:
    checkpointer.setup()
    print("Checkpoint tables created:")
    print("  - checkpoints")
    print("  - checkpoint_writes")
    print("  - checkpoint_blobs")


# ✅ Async equivalent — for setup from an async app
import asyncio
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver

async def init():
    async with AsyncPostgresSaver.from_conn_string(DATABASE_URL) as saver:
        await saver.setup()

asyncio.run(init())


# ✅ Migration-style setup — run once per deploy, no-op if already set up
# migrations/001_langgraph_checkpoints.py
from langgraph.checkpoint.postgres import PostgresSaver
import os, sys

def run_migration():
    dsn = os.environ["DATABASE_URL"]
    if not dsn.startswith("postgresql://"):
        # LangGraph/psycopg 3 requires the postgresql:// scheme
        dsn = dsn.replace("postgres://", "postgresql://", 1)

    with PostgresSaver.from_conn_string(dsn) as saver:
        saver.setup()
    print("LangGraph checkpoint schema is up to date.")

if __name__ == "__main__":
    try:
        run_migration()
    except Exception as e:
        print(f"Migration failed: {e}", file=sys.stderr)
        sys.exit(1)


# ❌ ANTI-PATTERN — calling setup() inside every request
# Slow, causes lock contention, floods the query log
# @app.post("/chat")
# def chat():
#     with PostgresSaver.from_conn_string(DSN) as saver:
#         saver.setup()                    # <-- don't do this per request
#         # ...

# ✅ Async pool for FastAPI — sized for the number of concurrent requests you expect
from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from fastapi import FastAPI
from contextlib import asynccontextmanager


DATABASE_URL = "postgresql://user:pass@host:5432/langgraph"


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Sizing rule of thumb:
    #   max_size = num_workers * expected_concurrent_requests_per_worker
    #   min_size = a warm baseline you always want ready
    pool = AsyncConnectionPool(
        DATABASE_URL,
        min_size=5,
        max_size=40,                     # tune to your load
        timeout=10,                      # fail fast if no conn in 10s
        max_lifetime=3600,               # recycle conns hourly
        max_idle=300,                    # release idle conns after 5 min
        kwargs={"autocommit": True, "prepare_threshold": None},
    )
    await pool.open()

    checkpointer = AsyncPostgresSaver(pool)
    await checkpointer.setup()           # idempotent

    app.state.pool = pool
    app.state.graph = build_graph(checkpointer)
    yield

    await pool.close()


api = FastAPI(lifespan=lifespan)


@api.post("/chat/{thread_id}")
async def chat(thread_id: str, message: str):
    config = {"configurable": {"thread_id": thread_id}}
    return await api.state.graph.ainvoke(
        {"messages": [HumanMessage(content=message)]},
        config,
    )


# ✅ Monitor pool health — expose a /metrics endpoint
@api.get("/pool-stats")
async def pool_stats():
    p = api.state.pool
    return {
        "size": p.get_stats().get("pool_size"),
        "available": p.get_stats().get("pool_available"),
        "waiting": p.get_stats().get("requests_waiting"),
    }


# Sync equivalent — for non-async apps
from psycopg_pool import ConnectionPool
from langgraph.checkpoint.postgres import PostgresSaver

pool = ConnectionPool(DATABASE_URL, min_size=5, max_size=40, timeout=10)
checkpointer = PostgresSaver(pool)
checkpointer.setup()

# ✅ Connect to Postgres through PgBouncer in transaction mode (Supabase, RDS Proxy)
from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver


PGBOUNCER_URL = "postgresql://user:pass@pgbouncer:6543/langgraph"


pool = AsyncConnectionPool(
    PGBOUNCER_URL,
    min_size=2,
    max_size=10,
    timeout=10,
    kwargs={
        "autocommit": True,
        # THE fix — disables psycopg prepared statement cache
        # PgBouncer transaction mode cannot preserve prepared statements
        # across connection reuse.
        "prepare_threshold": None,
    },
)


# ✅ Direct-to-Postgres (no PgBouncer) — leave prepared statements ON for perf
DIRECT_URL = "postgresql://user:pass@postgres:5432/langgraph"

direct_pool = AsyncConnectionPool(
    DIRECT_URL,
    min_size=5,
    max_size=40,
    kwargs={"autocommit": True},         # no prepare_threshold override
)


# ✅ Two-URL pattern (Supabase-style):
#   - Direct URL for migrations (needs prepared statements)
#   - Pooler URL for runtime (needs prepare_threshold=None)
import os

if os.environ.get("PHASE") == "migrate":
    # Migration uses direct URL
    with PostgresSaver.from_conn_string(os.environ["DIRECT_URL"]) as saver:
        saver.setup()
else:
    # Runtime uses pooler URL
    pool = AsyncConnectionPool(
        os.environ["POOLER_URL"],
        max_size=20,
        kwargs={"prepare_threshold": None},
    )
    checkpointer = AsyncPostgresSaver(pool)


# ❌ WRONG — using PgBouncer transaction mode without disabling prepare cache
# Every few requests silently returns stale data or errors with
# "prepared statement 'ps_1' does not exist"
# pool = AsyncConnectionPool(PGBOUNCER_URL, kwargs={"autocommit": True})  # missing prepare_threshold=None