Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions burr/core/persistence.py
Original file line number Diff line number Diff line change
Expand Up @@ -449,19 +449,22 @@ def is_initialized(self) -> bool:
return self._initialized

def list_app_ids(self, partition_key: Optional[str], **kwargs) -> list[str]:
"""Lists the app_ids for a given partition_key, each once, most recently saved first."""
"""Lists each app once, newest timestamp first, breaking ties by insertion order."""
partition_key = (
partition_key if partition_key is not None else SQLitePersister.PARTITION_KEY_DEFAULT
)

cursor = self.connection.cursor()
try:
# Only rows at each app's latest timestamp may break an ordering tie.
cursor.execute(
f"SELECT app_id FROM {self.table_name} "
f"WHERE partition_key = ? "
f"GROUP BY app_id "
f"ORDER BY MAX(created_at) DESC",
(partition_key,),
f"SELECT saved.app_id FROM {self.table_name} AS saved "
f"JOIN (SELECT app_id, MAX(created_at) AS created_at FROM {self.table_name} "
f"WHERE partition_key = ? GROUP BY app_id) AS latest "
f"ON saved.app_id = latest.app_id AND saved.created_at IS latest.created_at "
f"WHERE saved.partition_key = ? GROUP BY saved.app_id "
f"ORDER BY MAX(saved.created_at) DESC, MAX(saved.rowid) DESC",
(partition_key, partition_key),
)
except sqlite3.OperationalError as e:
if "no such table" in str(e):
Expand Down
15 changes: 9 additions & 6 deletions burr/integrations/persisters/b_aiosqlite.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,20 +188,23 @@ async def is_initialized(self) -> bool:
return self._initialized

async def list_app_ids(self, partition_key: Optional[str] = None, **kwargs) -> list[str]:
"""Lists the app_ids for a given partition_key, each once, most recently saved first."""
"""Lists each app once, newest timestamp first, breaking ties by insertion order."""
partition_key = (
partition_key
if partition_key is not None
else AsyncSQLitePersister.PARTITION_KEY_DEFAULT
)

cursor = await self.connection.cursor()
# Only rows at each app's latest timestamp may break an ordering tie.
await cursor.execute(
f"SELECT app_id FROM {self.table_name} "
f"WHERE partition_key = ? "
f"GROUP BY app_id "
f"ORDER BY MAX(created_at) DESC",
(partition_key,),
f"SELECT saved.app_id FROM {self.table_name} AS saved "
f"JOIN (SELECT app_id, MAX(created_at) AS created_at FROM {self.table_name} "
f"WHERE partition_key = ? GROUP BY app_id) AS latest "
f"ON saved.app_id = latest.app_id AND saved.created_at IS latest.created_at "
f"WHERE saved.partition_key = ? GROUP BY saved.app_id "
f"ORDER BY MAX(saved.created_at) DESC, MAX(saved.rowid) DESC",
(partition_key, partition_key),
)
app_ids = [row[0] for row in await cursor.fetchall()]
return app_ids
Expand Down
59 changes: 59 additions & 0 deletions tests/core/test_persistence.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,65 @@ def test_sqlite_persister_list_app_ids_is_unique_and_newest_first():
assert persister.list_app_ids("pk") == ["a-app", "z-app"]


@pytest.mark.parametrize("partition_key", [None, "", "tenant", " "])
@pytest.mark.parametrize("latest_app", ["a-app", "z-app"])
def test_sqlite_list_app_ids_breaks_timestamp_ties(tmp_path, partition_key, latest_app):
db_path = str(tmp_path / "app-order.db")
first_app = "z-app" if latest_app == "a-app" else "a-app"
read_keys = (None, "") if partition_key in (None, "") else (partition_key,)
with SQLLitePersister(db_path=db_path, table_name="test_table") as persister:
persister.initialize()
assert persister.list_app_ids(partition_key) == []
for app_id, sequence_id in [(latest_app, 0), (first_app, 0), (latest_app, 1)]:
persister.save(
partition_key, app_id, sequence_id, "step", State({"app": app_id}), "completed"
)
# A later row for the same app in another partition must not affect the order.
persister.save("other-partition", first_app, 2, "step", State({}), "completed")
# Avoid wall-clock races: all saves have exactly the same timestamp.
persister.connection.execute("UPDATE test_table SET created_at = '2026-01-01 00:00:00'")
persister.connection.commit()
for read_key in read_keys:
assert persister.list_app_ids(read_key) == [latest_app, first_app]

with SQLLitePersister(db_path=db_path, table_name="test_table") as persister:
for read_key in read_keys:
app_ids = persister.list_app_ids(read_key)
assert app_ids == [latest_app, first_app]
assert (persister.load(read_key, None))["app_id"] == app_ids[0]
assert (persister.load(read_key, app_ids[0]))["state"] == State({"app": latest_app})
assert persister.list_app_ids("missing-partition") == []


@pytest.mark.parametrize(
"timestamps,expected",
[
(
["2026-01-02 00:00:00", "2026-01-01 00:00:00", "2026-01-01 00:00:00"],
["a-app", "z-app"],
),
(
["2026-01-02 00:00:00", "2026-01-02 00:00:00", "2026-01-01 00:00:00"],
["z-app", "a-app"],
),
],
ids=["timestamp-before-insertion-order", "tie-break-uses-latest-timestamp-row"],
)
def test_sqlite_list_app_ids_prioritizes_latest_timestamp_rows(timestamps, expected):
with SQLLitePersister(db_path=":memory:", table_name="test_table") as persister:
persister.initialize()
saves = [("a-app", 0), ("z-app", 0), ("a-app", 1)]
for (app_id, sequence_id), created_at in zip(saves, timestamps):
persister.save("pk", app_id, sequence_id, "step", State({}), "completed")
persister.connection.execute(
"UPDATE test_table SET created_at = ? WHERE app_id = ? AND sequence_id = ?",
(created_at, app_id, sequence_id),
)
persister.connection.commit()
assert persister.list_app_ids("pk") == expected
assert (persister.load("pk", None))["app_id"] == expected[0]


@pytest.mark.parametrize("partition_key", [None, "partition_key"])
@pytest.mark.parametrize("has_saved_app", [False, True])
def test_memory_persistence_missing_load_preserves_app_ids(partition_key, has_saved_app):
Expand Down
69 changes: 69 additions & 0 deletions tests/integrations/persisters/test_b_aiosqlite.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,75 @@ async def test_async_persistence_list_app_ids_is_unique_and_newest_first(async_p
assert await async_persistence.list_app_ids("pk") == ["a-app", "z-app"]


@pytest.mark.parametrize("partition_key", [None, "", "tenant", " "])
@pytest.mark.parametrize("latest_app", ["a-app", "z-app"])
async def test_async_sqlite_list_app_ids_breaks_timestamp_ties(tmp_path, partition_key, latest_app):
db_path = str(tmp_path / "app-order.db")
first_app = "z-app" if latest_app == "a-app" else "a-app"
read_keys = (None, "") if partition_key in (None, "") else (partition_key,)
async with AsyncSQLitePersister.from_values(
db_path=db_path, table_name="test_table"
) as persister:
await persister.initialize()
assert await persister.list_app_ids(partition_key) == []
for app_id, sequence_id in [(latest_app, 0), (first_app, 0), (latest_app, 1)]:
await persister.save(
partition_key, app_id, sequence_id, "step", State({"app": app_id}), "completed"
)
# A later row for the same app in another partition must not affect the order.
await persister.save("other-partition", first_app, 2, "step", State({}), "completed")
# Avoid wall-clock races: all saves have exactly the same timestamp.
await persister.connection.execute(
"UPDATE test_table SET created_at = '2026-01-01 00:00:00'"
)
await persister.connection.commit()
for read_key in read_keys:
assert await persister.list_app_ids(read_key) == [latest_app, first_app]

async with AsyncSQLitePersister.from_values(
db_path=db_path, table_name="test_table"
) as persister:
for read_key in read_keys:
app_ids = await persister.list_app_ids(read_key)
assert app_ids == [latest_app, first_app]
assert (await persister.load(read_key, None))["app_id"] == app_ids[0]
assert (await persister.load(read_key, app_ids[0]))["state"] == State(
{"app": latest_app}
)
assert await persister.list_app_ids("missing-partition") == []


@pytest.mark.parametrize(
"timestamps,expected",
[
(
["2026-01-02 00:00:00", "2026-01-01 00:00:00", "2026-01-01 00:00:00"],
["a-app", "z-app"],
),
(
["2026-01-02 00:00:00", "2026-01-02 00:00:00", "2026-01-01 00:00:00"],
["z-app", "a-app"],
),
],
ids=["timestamp-before-insertion-order", "tie-break-uses-latest-timestamp-row"],
)
async def test_async_sqlite_list_app_ids_prioritizes_latest_timestamp_rows(timestamps, expected):
async with AsyncSQLitePersister.from_values(
db_path=":memory:", table_name="test_table"
) as persister:
await persister.initialize()
saves = [("a-app", 0), ("z-app", 0), ("a-app", 1)]
for (app_id, sequence_id), created_at in zip(saves, timestamps):
await persister.save("pk", app_id, sequence_id, "step", State({}), "completed")
await persister.connection.execute(
"UPDATE test_table SET created_at = ? WHERE app_id = ? AND sequence_id = ?",
(created_at, app_id, sequence_id),
)
await persister.connection.commit()
assert await persister.list_app_ids("pk") == expected
assert (await persister.load("pk", None))["app_id"] == expected[0]


@pytest.mark.parametrize("partition_key", [None, "partition"])
@pytest.mark.parametrize("latest_app,latest_sequence", [("a-app", 12), ("z-app", 0)])
async def test_async_sqlite_load_latest_breaks_timestamp_ties(
Expand Down
Loading