From 1a57031366cdfe3c0062a4216d4c169b722b1763 Mon Sep 17 00:00:00 2001 From: MIRS571 <259312884+MIRS571@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:00:55 +0800 Subject: [PATCH] fix: normalize default partitions in in-memory persisters --- burr/core/persistence.py | 38 +++++++++--- docs/concepts/state-persistence.rst | 15 +++++ tests/core/test_persistence.py | 89 +++++++++++++++++++++++++++++ 3 files changed, 134 insertions(+), 8 deletions(-) diff --git a/burr/core/persistence.py b/burr/core/persistence.py index d0775b55c..d1e698d0e 100644 --- a/burr/core/persistence.py +++ b/burr/core/persistence.py @@ -615,18 +615,27 @@ def __setstate__(self, state): class InMemoryPersister(BaseStatePersister): - """In-memory persister for testing purposes. This is not recommended for production use.""" + """In-memory persister for testing purposes. This is not recommended for production use. + + None and the empty string identify the same default partition, as in SQLitePersister. + Non-empty partition keys are preserved. + """ def __init__(self): self._storage = defaultdict(lambda: defaultdict(list)) def load( - self, partition_key: str, app_id: Optional[str], sequence_id: Optional[int] = None, **kwargs + self, + partition_key: Optional[str], + app_id: Optional[str], + sequence_id: Optional[int] = None, + **kwargs, ) -> Optional[PersistedStateData]: # If no app_id provided, return None if app_id is None: return None + partition_key = "" if partition_key is None else partition_key if not (states := self._storage.get(partition_key, {}).get(app_id)): return None @@ -639,7 +648,8 @@ def load( # Return the latest state for this sequence_id, if exists return matching_states[-1] if matching_states else None - def list_app_ids(self, partition_key: str, **kwargs) -> list[str]: + def list_app_ids(self, partition_key: Optional[str], **kwargs) -> list[str]: + partition_key = "" if partition_key is None else partition_key return list(self._storage[partition_key].keys()) def save( @@ -652,9 +662,10 @@ def save( status: Literal["completed", "failed"], **kwargs, ): + partition_key = "" if partition_key is None else partition_key # Create a PersistedStateData entry persisted_state: PersistedStateData = { - "partition_key": partition_key or "", + "partition_key": partition_key, "app_id": app_id, "sequence_id": sequence_id, "position": position, @@ -668,18 +679,27 @@ def save( class AsyncInMemoryPersister(AsyncBaseStatePersister): - """Sync in-memory persister for testing purposes. This is not recommended for production use.""" + """Async in-memory persister for testing purposes. This is not recommended for production use. + + None and the empty string identify the same default partition, as in SQLitePersister. + Non-empty partition keys are preserved. + """ def __init__(self): self._storage = defaultdict(lambda: defaultdict(list)) async def load( - self, partition_key: str, app_id: Optional[str], sequence_id: Optional[int] = None, **kwargs + self, + partition_key: Optional[str], + app_id: Optional[str], + sequence_id: Optional[int] = None, + **kwargs, ) -> Optional[PersistedStateData]: # If no app_id provided, return None if app_id is None: return None + partition_key = "" if partition_key is None else partition_key if not (states := self._storage.get(partition_key, {}).get(app_id)): return None @@ -692,7 +712,8 @@ async def load( # Return the latest state for this sequence_id, if exists return matching_states[-1] if matching_states else None - async def list_app_ids(self, partition_key: str, **kwargs) -> list[str]: + async def list_app_ids(self, partition_key: Optional[str], **kwargs) -> list[str]: + partition_key = "" if partition_key is None else partition_key return list(self._storage[partition_key].keys()) async def save( @@ -705,9 +726,10 @@ async def save( status: Literal["completed", "failed"], **kwargs, ): + partition_key = "" if partition_key is None else partition_key # Create a PersistedStateData entry persisted_state: PersistedStateData = { - "partition_key": partition_key or "", + "partition_key": partition_key, "app_id": app_id, "sequence_id": sequence_id, "position": position, diff --git a/docs/concepts/state-persistence.rst b/docs/concepts/state-persistence.rst index e6b38ffb8..b8d421e65 100644 --- a/docs/concepts/state-persistence.rst +++ b/docs/concepts/state-persistence.rst @@ -60,6 +60,21 @@ Burr `applications` are, by default, keyed on two entities: In the case of a chatbot, the ``app_id`` could be a uuid, and the ``partition_key`` could be the user's ID or email, etc. Note that ``partition_key`` can be `None` if this is not relevant. A UUID is always generated for the ``app_id`` if not provided. +For ``InMemoryPersister``, ``AsyncInMemoryPersister``, and ``SQLitePersister``, +``None`` and the empty string ``""`` identify the same default partition when +saving, loading, or listing applications. Checkpoints loaded from this default +partition use ``""`` as their partition key. Non-empty keys, including +whitespace-only strings, remain distinct. + +.. note:: + + Earlier versions of the in-memory persisters treated ``None`` and ``""`` as + separate storage keys. They now share one application history: using both + aliases for the same ``app_id`` does not create two isolated applications. + In-memory loads still select the last saved checkpoint, or the last save for + a requested sequence ID. Use distinct non-empty keys if you need separate + partitions. + You set these values using the :py:meth:`with_identifiers() ` method. Note: to access application-level metadata such as ``app_id``, ``partition_key``, ``sequence_id``, and ``action_name`` in your running application, you can have the :py:class:`ApplicationContext ` diff --git a/tests/core/test_persistence.py b/tests/core/test_persistence.py index 1944bbb13..2c7cf32cb 100644 --- a/tests/core/test_persistence.py +++ b/tests/core/test_persistence.py @@ -128,6 +128,51 @@ def test_memory_persistence_missing_load_preserves_app_ids(partition_key, has_sa assert persistence.load(partition_key, "missing")["state"] == State({"step": 3}) +@pytest.mark.parametrize("save_key", [None, ""]) +@pytest.mark.parametrize("load_key", [None, ""]) +def test_persistence_default_partition_aliases(persistence, save_key, load_key): + """Default partition aliases can list and reload the same checkpoint.""" + if hasattr(persistence, "initialize"): + persistence.initialize() + state = State({"step": 1}) + persistence.save(save_key, "app", 1, "position", state, "completed") + + loaded = persistence.load(load_key, "app", 1) + assert loaded is not None + assert loaded["partition_key"] == "" + assert loaded["state"] == state + reloaded = persistence.load(loaded["partition_key"], "app") + assert reloaded is not None + assert reloaded["state"] == state + assert persistence.list_app_ids(load_key) == ["app"] + assert persistence.load(load_key, "missing") is None + assert persistence.load(load_key, "app", 99) is None + assert persistence.list_app_ids(load_key) == ["app"] + + +@pytest.mark.parametrize("first_key,second_key", [(None, ""), ("", None)]) +def test_memory_default_partition_aliases_share_history(first_key, second_key): + """Mixed aliases retain history and last-save semantics, not separate apps.""" + persistence = InMemoryPersister() + persistence.save(first_key, "app", 1, "first", State({"step": 1}), "completed") + persistence.save(second_key, "app", 2, "second", State({"step": 2}), "completed") + persistence.save(first_key, "app", 1, "retry", State({"step": 3}), "completed") + persistence.save(second_key, "other", 1, "other", State({"step": 4}), "completed") + + for key in ("tenant", " ", "None"): + persistence.save(key, "app", 1, "named", State({"partition": key}), "completed") + loaded = persistence.load(key, "app") + assert loaded["partition_key"] == key + assert loaded["state"] == State({"partition": key}) + + for key in (None, ""): + assert set(persistence.list_app_ids(key)) == {"app", "other"} + assert persistence.load(key, "app")["state"] == State({"step": 3}) + assert persistence.load(key, "app", 1)["state"] == State({"step": 3}) + assert persistence.load(key, "app", 2)["state"] == State({"step": 2}) + assert persistence.load(key, "other")["state"] == State({"step": 4}) + + def test_persistence_is_initialized_false(initializing_persistence): assert not initializing_persistence.is_initialized() @@ -401,6 +446,50 @@ async def test_async_persister_methods_none_partition_key( # these operations are stateful (i.e., read/write to a db) +@pytest.mark.parametrize("save_key", [None, ""]) +@pytest.mark.parametrize("load_key", [None, ""]) +async def test_async_persistence_default_partition_aliases(async_persistence, save_key, load_key): + """Default partition aliases can list and reload the same checkpoint.""" + state = State({"step": 1}) + await async_persistence.save(save_key, "app", 1, "position", state, "completed") + + loaded = await async_persistence.load(load_key, "app", 1) + assert loaded is not None + assert loaded["partition_key"] == "" + assert loaded["state"] == state + reloaded = await async_persistence.load(loaded["partition_key"], "app") + assert reloaded is not None + assert reloaded["state"] == state + assert await async_persistence.list_app_ids(load_key) == ["app"] + assert await async_persistence.load(load_key, "missing") is None + assert await async_persistence.load(load_key, "app", 99) is None + assert await async_persistence.list_app_ids(load_key) == ["app"] + + +@pytest.mark.parametrize("first_key,second_key", [(None, ""), ("", None)]) +async def test_async_memory_default_partition_aliases_share_history( + async_persistence, first_key, second_key +): + """Mixed aliases retain history and last-save semantics, not separate apps.""" + await async_persistence.save(first_key, "app", 1, "first", State({"step": 1}), "completed") + await async_persistence.save(second_key, "app", 2, "second", State({"step": 2}), "completed") + await async_persistence.save(first_key, "app", 1, "retry", State({"step": 3}), "completed") + await async_persistence.save(second_key, "other", 1, "other", State({"step": 4}), "completed") + + for key in ("tenant", " ", "None"): + await async_persistence.save(key, "app", 1, "named", State({"partition": key}), "completed") + loaded = await async_persistence.load(key, "app") + assert loaded["partition_key"] == key + assert loaded["state"] == State({"partition": key}) + + for key in (None, ""): + assert set(await async_persistence.list_app_ids(key)) == {"app", "other"} + assert (await async_persistence.load(key, "app"))["state"] == State({"step": 3}) + assert (await async_persistence.load(key, "app", 1))["state"] == State({"step": 3}) + assert (await async_persistence.load(key, "app", 2))["state"] == State({"step": 2}) + assert (await async_persistence.load(key, "other"))["state"] == State({"step": 4}) + + async def test_AsyncSQLitePersister_from_values(): await asyncio.sleep(0.00001) connection = await aiosqlite.connect(":memory:")