from typing import Any import pytest from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, create_checkpoint, empty_checkpoint, ) from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver class TestAsyncSqliteSaver: @pytest.fixture(autouse=True) def setup(self) -> None: # objects for test setup self.config_1: RunnableConfig = { "configurable": { "thread_id": "thread-1", "checkpoint_id": "1", "checkpoint_ns": "", } } self.config_2: RunnableConfig = { "configurable": { "thread_id": "thread-2", "checkpoint_id": "2", "checkpoint_ns": "", } } self.config_3: RunnableConfig = { "configurable": { "thread_id": "thread-2", "checkpoint_id": "2-inner", "checkpoint_ns": "inner", } } self.chkpnt_1: Checkpoint = empty_checkpoint() self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1) self.chkpnt_3: Checkpoint = empty_checkpoint() self.metadata_1: CheckpointMetadata = { "source": "input", "step": 2, "writes": {}, "score": 1, } self.metadata_2: CheckpointMetadata = { "source": "loop", "step": 1, "writes": {"foo": "bar"}, "score": None, } self.metadata_3: CheckpointMetadata = {} async def test_combined_metadata(self) -> None: async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: config: RunnableConfig = { "configurable": { "thread_id": "thread-2", "checkpoint_ns": "", "__super_private_key": "super_private_value", }, "metadata": {"run_id": "my_run_id"}, } await saver.aput(config, self.chkpnt_2, self.metadata_2, {}) checkpoint = await saver.aget_tuple(config) assert checkpoint is not None and checkpoint.metadata == { **self.metadata_2, "run_id": "my_run_id", } async def test_asearch(self) -> None: async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: await saver.aput(self.config_1, self.chkpnt_1, self.metadata_1, {}) await saver.aput(self.config_2, self.chkpnt_2, self.metadata_2, {}) await saver.aput(self.config_3, self.chkpnt_3, self.metadata_3, {}) # call method / assertions query_1 = {"source": "input"} # search by 1 key query_2 = { "step": 1, "writes": {"foo": "bar"}, } # search by multiple keys query_3: dict[str, Any] = {} # search by no keys, return all checkpoints query_4 = {"source": "update", "step": 1} # no match search_results_1 = [c async for c in saver.alist(None, filter=query_1)] assert len(search_results_1) == 1 assert search_results_1[0].metadata == self.metadata_1 search_results_2 = [c async for c in saver.alist(None, filter=query_2)] assert len(search_results_2) == 1 assert search_results_2[0].metadata == self.metadata_2 search_results_3 = [c async for c in saver.alist(None, filter=query_3)] assert len(search_results_3) == 3 search_results_4 = [c async for c in saver.alist(None, filter=query_4)] assert len(search_results_4) == 0 # search by config (defaults to checkpoints across all namespaces) search_results_5 = [ c async for c in saver.alist({"configurable": {"thread_id": "thread-2"}}) ] assert len(search_results_5) == 2 assert { search_results_5[0].config["configurable"]["checkpoint_ns"], search_results_5[1].config["configurable"]["checkpoint_ns"], } == {"", "inner"} # Test limit param search_results_6 = [ c async for c in saver.alist( {"configurable": {"thread_id": "thread-2"}}, limit=1 ) ] assert len(search_results_6) == 1 assert search_results_6[0].config["configurable"]["thread_id"] == "thread-2" # Test before param search_results_7 = [ c async for c in saver.alist(None, before=search_results_5[1].config) ] assert len(search_results_7) == 1 assert search_results_7[0].config["configurable"]["thread_id"] == "thread-1" async def test_limit_parameter_sql_injection_prevention(self) -> None: """Test that the limit parameter properly uses parameterized queries to prevent SQL injection.""" async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: # Setup: Create multiple checkpoints for i in range(5): config: RunnableConfig = { "configurable": { "thread_id": f"thread-{i}", "checkpoint_ns": "", } } checkpoint = empty_checkpoint() metadata: CheckpointMetadata = {"index": i} await saver.aput(config, checkpoint, metadata, {}) # Test that limit works correctly with valid integer results = [c async for c in saver.alist(None, limit=2)] assert len(results) == 2 # Test that limit=0 returns no results results = [c async for c in saver.alist(None, limit=0)] assert len(results) == 0 # Test that limit=None returns all results results = [c async for c in saver.alist(None, limit=None)] assert len(results) == 5 # Test explicit SQL injection attempt via limit parameter # Even if type checking is bypassed and a malicious string is passed, # the parameterized query will treat it as a value, not SQL code # This would cause an error (can't convert string to int for LIMIT), # which is the correct secure behavior malicious_limits = [ "1; DROP TABLE checkpoints; --", "1 OR 1=1", "999999 UNION SELECT * FROM checkpoints", ] for malicious_limit in malicious_limits: # The parameterized query should safely reject non-integer limits # or convert them in a way that prevents SQL injection try: # Bypass type checking by casting results = [ c async for c in saver.alist(None, limit=malicious_limit) # type: ignore ] # If it doesn't raise an error, it should at least not execute the injection # SQLite's parameter binding will try to convert the string to an integer # which will either fail or treat it as 0 except Exception: # Expected: SQLite should reject invalid limit values pass # Verify the checkpoints table still exists and has all data # (would have been dropped if injection succeeded) results = [c async for c in saver.alist(None, limit=None)] assert len(results) == 5