Restore shallow checkpointer

- This should definitely be removed soon, but let's give people more time to update
This commit is contained in:
Nuno Campos
2025-06-13 17:37:40 -07:00
parent 5498893780
commit 0cad7019cb
5 changed files with 1030 additions and 10 deletions
+30 -4
View File
@@ -16,7 +16,7 @@ from langgraph.checkpoint.base import (
Checkpoint,
CheckpointMetadata,
)
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver
from langgraph.checkpoint.serde.types import TASKS
from tests.checkpoint_utils import create_checkpoint, empty_checkpoint
from tests.conftest import DEFAULT_POSTGRES_URI
@@ -97,11 +97,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
@@ -161,7 +187,7 @@ def test_data():
}
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
def test_combined_metadata(saver_name: str, test_data) -> None:
with _saver(saver_name) as saver:
config = {
@@ -188,7 +214,7 @@ def test_combined_metadata(saver_name: str, test_data) -> None:
}
@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"]
@@ -237,7 +263,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(