PostgresSaver:生產級關聯型檢查點

源码版本1.2.9

職責

PostgresSaver 是 BaseCheckpointSaver 在 Postgres 上的落地實作,也是 LangGraph 官方推薦的生產 saver。程式碼在 libs/checkpoint-postgres/langgraph/checkpoint/postgres/,同步版入口是 libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py 裡的 PostgresSaver (PostgresSaver:40),非同步版在 aio.py 的 AsyncPostgresSaver (AsyncPostgresSaver:40),共用基類 BasePostgresSaver (BasePostgresSaver:297),後者負責 SQL 模板和遷移指令碼。

跟 SQLite 單表 + 主鍵去重的玩法不一樣,PostgresSaver 把資料拆成三張表:checkpoints(主表,JSONB)、checkpoint_blobs(channel value 二進位)、checkpoint_writes(task 增量寫)(MIGRATIONS:43-91)。這種拆法讓讀取能用單條 SELECT + lateral 子查詢一次性把主表 + blobs + writes 全撈回來 (SELECT_SQL:93-118),避免 SQLite 那種 N+1 查詢;同時 blob 列獨立用 BYTEA,不參與 JSONB 反序列化,大物件不會拖累主表讀取。

PostgresSaver 額外支援 psycopg3 的 Pipeline 模式(supports_pipeline:60)——多語句打包下發,服務端不立即回傳結果,適合一次 put 內部的 executemany(blob) + execute(checkpoint) 這種組合。

設計動機

• 三表分離,JSONB 存主體:checkpoints 主表用 JSONB 列裝 Checkpoint TypedDict,checkpoint_blobs 獨立存大物件 (checkpoint_blobs:57-65)。這避免了把整個 channel_values 序列化塞進 JSONB——那種做法在通道多、value 大時效能崩盤。
• 單 SELECT 把 tuple 拼全:SELECT_SQL 用兩個 lateral 子查詢 (array_agg over checkpoint_blobs + array_agg over checkpoint_writes) 一次拿回所有資料 (lateral subqueries:101-117),_load_checkpoint_tuple 在 Python 端解 array 即可,不需要二次查詢。
• UPSERT 用 ON CONFLICT DO UPDATE/NOTHING 區分語意:UPSERT_CHECKPOINT_BLOBS_SQL 用 DO NOTHING (UPSERT_BLOBS:131-135),UPSERT_CHECKPOINTS_SQL 用 DO UPDATE (UPSERT_CHECKPOINTS:137-144),UPSERT_CHECKPOINT_WRITES_SQL 用 DO UPDATE,而 INSERT_CHECKPOINT_WRITES_SQL 用 DO NOTHING (INSERT_WRITES:155-159)。控制 channel 寫要覆蓋最新狀態(用 UPDATE),普通業務寫撞了忽略(用 NOTHING),語意跟 InMemory 的 if inner_key[1] >= 0 ... continue 對齊。
• 遷移列表追加版本號:MIGRATIONS 是 list (MIGRATIONS:43-91),setup() 讀 checkpoint_migrations 表的最大 v,只跑後面的遷移 (setup:85-110)。加列、加索引、改約束都透過追加 migration 字串完成,不改老遷移。
• Pipeline 模式讓單連接真並行:同步 PostgresSaver._cursor 在 pipeline=True 時拿 conn.pipeline() 上下文 (pipeline branch:424-432),一條連接上的多個 cursor 不阻塞彼此;而 pipe 模式 (從 from_conn_string(pipeline=True)) 讓一個連接在多執行緒下復用 (pipe branch:413-423),配合 self.lock 串行 cursor 但不串行語句。

關鍵檔案

• MIGRATIONS:43-91 — 三張表的 CREATE 語句 + WAL/索引/task_path 遷移。
• SELECT_SQL:93-118 — 主查詢,lateral 子查詢把 blobs 和 writes 一起 array_agg 出來。
• SELECT_PENDING_SENDS_SQL:120-129 — 取待發送的 PUSH 任務 writes,按 task_path / task_id / idx 排序。
• UPSERT / INSERT SQL:131-159 — 四條寫入語句,區分 DO UPDATE / DO NOTHING。
• BasePostgresSaver:297 — 共用基類,持有 SQL 模板常數。
• PostgresSaver 類別:40-60 — 同步實作,接 Connection 或 ConnectionPool,帶 threading.Lock。
• from_conn_string:62-83 — 上下文管理器工廠,pipeline=True 走 conn.pipeline()。
• setup:85-110 — 讀 migration 版本,順序跑未應用的遷移並記 v。
• put:263-345 — 把 _DeltaSnapshot / 非原始值挪到 blob_values,主體進 JSONB。
• _cursor:405-442 — 按是否 pipeline / 是否 pipe 三種模式挑上下文。
• AsyncPostgresSaver:40 — 非同步版,aput / aget_tuple / alist 原生 async。

資料流

PostgresSaver.put 是看 saver 怎麼拆 Checkpoint TypedDict 的好例子——它先決定哪些 channel value 內聯進 JSONB 主體、哪些挪到 checkpoint_blobs,然後一次 cursor 內跑 executemany(blob) + execute(checkpoint):

# inline primitive values in checkpoint table
# others are stored in blobs table
blob_values = {}
for k, v in checkpoint["channel_values"].items():
    if isinstance(v, _DeltaSnapshot):
        blob_values[k] = copy["channel_values"].pop(k)
        copy["channel_values"][k] = True
    elif v is None or isinstance(v, (str, int, float, bool)):
        pass
    else:
        blob_values[k] = copy["channel_values"].pop(k)

with self._cursor(pipeline=True) as cur:
    if blob_versions := {
        k: v for k, v in new_versions.items() if k in blob_values
    }:
        cur.executemany(
            self.UPSERT_CHECKPOINT_BLOBS_SQL,
            self._dump_blobs(
                thread_id,
                checkpoint_ns,
                blob_values,
                blob_versions,
            ),
        )
    cur.execute(
        self.UPSERT_CHECKPOINTS_SQL,
        (
            thread_id,
            checkpoint_ns,
            checkpoint["id"],
            checkpoint_id,
            Jsonb(copy),
            Jsonb(get_serializable_checkpoint_metadata(config, metadata)),
        ),
    )

(put body:309-344) 注意 _DeltaSnapshot 走特殊路徑——通道值在主表裡只放一個 True 占位,真實資料進 blob,這是 DeltaChannel 重建時反向走祖先鏈的依據。讀回來時 get_tuple 跑 SELECT_SQL,psycopg 把 JSONB 自動轉 dict、BYTEA 轉 bytes,Python 端再呼叫 _load_checkpoint_tuple 組裝 CheckpointTuple (_load_checkpoint_tuple:552)。

邊界與失敗

• from_conn_string(pipeline=True) 不能跟 ConnectionPool 混用:__init__ 明確 raise (pool vs pipe check:52-55),pipeline 是單連接模式,pool 已經是多連接。
• setup() 必須顯式呼叫一次:docstring 寫明「MUST be called directly by the user the first time checkpointer is used」(setup docstring:85-91)。它建表 + 跑遷移,跑前要給資料庫建好權限。
• MIGRATIONS 改順序會炸:遷移按 list 位置當版本號 (MIGRATIONS:43-91),已經發布過的遷移字串不能改;改了 checkpoint_migrations 表裡 v 跟程式碼裡位置就錯位。
• CREATE INDEX CONCURRENTLY 在事務裡會失敗:MIGRATIONS 裡兩條並行建索引語句 (CONCURRENTLY indexes:82-89) 跟其它遷移放同一個 setup() 上下文裡跑,如果連接不在 autocommit 模式會報錯;from_conn_string 顯式 autocommit=True 是為這個 (autocommit:76-78)。
• supports_pipeline 執行時檢測:Capabilities().has_pipeline() (has_pipeline:60) 在 __init__ 裡查 psycopg 版本和服務端能力,老版 psycopg 或某些託管 Postgres 退化為普通事務模式 (transaction fallback:434-439)。
• DeltaChannel 有專門的快速路徑:get_delta_channel_history 在 PostgresSaver 裡被覆寫 (get_delta_channel_history:444-463) 走兩階段查詢(Stage 1 metadata 分頁 + Stage 2 UNION ALL 拉 writes),不走基類預設實作,自訂 saver 必須同步實作這一介面才能跑 DeltaChannel。

小結

PostgresSaver 是 LangGraph 推薦生產 saver,三表拆分讓「主表 JSONB + 大物件 BYTEA + 增量 writes」各自走最合適的儲存;Pipeline 模式和 async 原生實作給到生產規模的吞吐。介面形狀見 BaseCheckpointSaver,輕量場景見 InMemorySaver + SqliteSaver。

對照官方資料:LangGraph 文件 · README