mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-11 12:17:53 +02:00
Restore shallow checkpointer
- This should definitely be removed soon, but let's give people more time to update
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user