diff --git a/burr/core/persistence.py b/burr/core/persistence.py index c32bf8e96..87281d23a 100644 --- a/burr/core/persistence.py +++ b/burr/core/persistence.py @@ -337,14 +337,17 @@ def from_values( :param db_path: the path the DB will be stored. :param table_name: the table name to store things under. :param serde_kwargs: kwargs for state serialization/deserialization. - :param connect_kwargs: kwargs to pass to the aiosqlite.connect method. - :return: async sqlite persister instance with an open connection. You are responsible + :param connect_kwargs: kwargs to pass to sqlite3.connect, retained when copying or + unpickling the persister. + :return: sqlite persister instance with an open connection. You are responsible for closing the connection yourself. """ connection = sqlite3.connect( db_path, **connect_kwargs if connect_kwargs is not None else {} ) - return cls(db_path, table_name, serde_kwargs, connection=connection) + return cls( + db_path, table_name, serde_kwargs, connect_kwargs=connect_kwargs, connection=connection + ) def copy(self) -> "Self": return SQLitePersister( diff --git a/tests/core/test_persistence.py b/tests/core/test_persistence.py index 18c65a0d1..2e35617a2 100644 --- a/tests/core/test_persistence.py +++ b/tests/core/test_persistence.py @@ -15,6 +15,9 @@ # specific language governing permissions and limitations # under the License. +import pickle +from concurrent.futures import ThreadPoolExecutor + import pytest from burr.core import State @@ -94,6 +97,31 @@ def test_sqlite_persistence_is_initialized_true_new_connection(tmp_path): p2.cleanup() +@pytest.mark.parametrize("factory", ["from_values", "from_config"]) +@pytest.mark.parametrize("recreate", ["copy", "pickle"]) +def test_sqlite_factory_preserves_connection_options(tmp_path, factory, recreate): + config = { + "db_path": str(tmp_path / "state.db"), + "table_name": "checkpoints", + "connect_kwargs": {"timeout": 0.123, "check_same_thread": False}, + } + persister = ( + SQLLitePersister.from_values(**config) + if factory == "from_values" + else SQLLitePersister.from_config(config) + ) + with persister: + persister.initialize() + persister.save("tenant", "agent", 1, "checkpoint", State({"count": 1}), "completed") + restored = persister.copy() if recreate == "copy" else pickle.loads(pickle.dumps(persister)) + with restored: + assert restored.connection is not persister.connection + assert restored.connection.execute("PRAGMA busy_timeout").fetchone()[0] == 123 + with ThreadPoolExecutor(max_workers=1) as executor: + loaded = executor.submit(restored.load, "tenant", "agent").result(timeout=5) + assert loaded["state"] == State({"count": 1}) + + def test_sqlite_persister_load_without_initialize_raises_runtime_error(): """Test that calling load() without initialize() raises a clear RuntimeError.""" persister = SQLLitePersister(db_path=":memory:", table_name="test_table")