31 lines
784 B
Python
31 lines
784 B
Python
from __future__ import annotations
|
|
|
|
import threading
|
|
|
|
import psycopg
|
|
from langgraph.checkpoint.postgres import PostgresSaver
|
|
from psycopg.rows import dict_row
|
|
|
|
from app.modules.shared.db import database_url
|
|
|
|
_CHECKPOINTER: PostgresSaver | None = None
|
|
_LOCK = threading.Lock()
|
|
|
|
|
|
def _conn_string() -> str:
|
|
url = database_url()
|
|
if url.startswith("postgresql+psycopg"):
|
|
return url.replace("postgresql+psycopg", "postgresql", 1)
|
|
return url
|
|
|
|
|
|
def get_checkpointer() -> PostgresSaver:
|
|
global _CHECKPOINTER
|
|
with _LOCK:
|
|
if _CHECKPOINTER is None:
|
|
conn = psycopg.connect(_conn_string(), autocommit=True, row_factory=dict_row)
|
|
cp = PostgresSaver(conn)
|
|
cp.setup()
|
|
_CHECKPOINTER = cp
|
|
return _CHECKPOINTER
|