mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-13 21:27:52 +02:00
checkpoint postgres: add a shallow checkpointer (#2826)
This PR adds a "shallow" version of `PostgresSaver` checkpointer that ONLY stores the most recent checkpoint and does NOT retain any history. It is meant to be a light-weight drop-in replacement for the PostgresSaver that supports most of the LangGraph persistence functionality with the exception of time travel.
This commit is contained in:
@@ -16,7 +16,7 @@ from langgraph.checkpoint.base import (
|
||||
create_checkpoint,
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver
|
||||
from tests.conftest import DEFAULT_POSTGRES_URI
|
||||
|
||||
|
||||
@@ -91,11 +91,37 @@ def _base_saver():
|
||||
conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _shallow_saver():
|
||||
"""Fixture for regular connection mode testing with a shallow checkpointer."""
|
||||
database = f"test_{uuid4().hex[:16]}"
|
||||
# create unique db
|
||||
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
||||
conn.execute(f"CREATE DATABASE {database}")
|
||||
try:
|
||||
with Connection.connect(
|
||||
DEFAULT_POSTGRES_URI + database,
|
||||
autocommit=True,
|
||||
prepare_threshold=0,
|
||||
row_factory=dict_row,
|
||||
) as conn:
|
||||
checkpointer = ShallowPostgresSaver(conn)
|
||||
checkpointer.setup()
|
||||
yield checkpointer
|
||||
finally:
|
||||
# drop unique db
|
||||
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
||||
conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _saver(name: str):
|
||||
if name == "base":
|
||||
with _base_saver() as saver:
|
||||
yield saver
|
||||
elif name == "shallow":
|
||||
with _shallow_saver() as saver:
|
||||
yield saver
|
||||
elif name == "pool":
|
||||
with _pool_saver() as saver:
|
||||
yield saver
|
||||
@@ -155,7 +181,7 @@ def test_data():
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
|
||||
def test_search(saver_name: str, test_data) -> None:
|
||||
with _saver(saver_name) as saver:
|
||||
configs = test_data["configs"]
|
||||
@@ -198,7 +224,7 @@ def test_search(saver_name: str, test_data) -> None:
|
||||
} == {"", "inner"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
|
||||
def test_null_chars(saver_name: str, test_data) -> None:
|
||||
with _saver(saver_name) as saver:
|
||||
config = saver.put(
|
||||
|
||||
Reference in New Issue
Block a user