perf(dashboard): keyset pagination for streaming session export
OFFSET paging made the streaming export O(n^2) on huge transcripts; after_id keyset paging keeps each page seek O(1). Adds after_id to SessionDB.get_messages (ascending-only, guarded against latest/offset combos).
This commit is contained in:
@@ -746,13 +746,15 @@ async def export_session_endpoint(session_id: str, profile: Optional[str] = None
|
||||
)
|
||||
yield metadata[:-1] + ',"messages":['
|
||||
|
||||
offset = 0
|
||||
# Keyset pagination (id > last_seen): O(n) total over the
|
||||
# transcript, vs OFFSET's O(n²) on huge sessions.
|
||||
last_id = None
|
||||
first = True
|
||||
while True:
|
||||
messages = db.get_messages(
|
||||
sid,
|
||||
limit=500,
|
||||
offset=offset,
|
||||
after_id=last_id if last_id is not None else 0,
|
||||
)
|
||||
for message in messages:
|
||||
if not first:
|
||||
@@ -765,7 +767,9 @@ async def export_session_endpoint(session_id: str, profile: Optional[str] = None
|
||||
first = False
|
||||
if len(messages) < 500:
|
||||
break
|
||||
offset += len(messages)
|
||||
last_id = messages[-1].get("id")
|
||||
if last_id is None:
|
||||
break # defensive: cannot keyset without row ids
|
||||
|
||||
yield "]}"
|
||||
finally:
|
||||
|
||||
+11
-1
@@ -7421,6 +7421,7 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin)
|
||||
limit: Optional[int] = None,
|
||||
offset: int = 0,
|
||||
latest: bool = False,
|
||||
after_id: Optional[int] = None,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Load messages for a session in insertion order.
|
||||
|
||||
@@ -7440,13 +7441,22 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin)
|
||||
order. ``offset`` alone (without ``limit``) also pages — SQLite
|
||||
requires a LIMIT clause for OFFSET, so it's emitted as ``LIMIT -1``
|
||||
(unbounded).
|
||||
|
||||
``after_id`` enables keyset pagination (``id > after_id``): O(1)
|
||||
page seeks on huge transcripts where OFFSET degrades to O(n) per
|
||||
page. Ascending order only (incompatible with ``latest``/``offset``).
|
||||
"""
|
||||
if after_id is not None and (latest or offset):
|
||||
raise ValueError("after_id is incompatible with latest/offset paging")
|
||||
active_clause = "" if include_inactive else " AND active = 1"
|
||||
keyset_clause = " AND id > ?" if after_id is not None else ""
|
||||
sql = (
|
||||
"SELECT * FROM messages WHERE session_id = ?"
|
||||
f"{active_clause} ORDER BY id {'DESC' if latest else 'ASC'}"
|
||||
f"{active_clause}{keyset_clause} ORDER BY id {'DESC' if latest else 'ASC'}"
|
||||
)
|
||||
params: list = [session_id]
|
||||
if after_id is not None:
|
||||
params.append(after_id)
|
||||
if limit is not None or offset:
|
||||
# SQLite's OFFSET requires LIMIT; -1 means "no limit".
|
||||
sql += " LIMIT ? OFFSET ?"
|
||||
|
||||
@@ -1895,7 +1895,7 @@ class TestWebServerEndpoints:
|
||||
original_get_messages = SessionDB.get_messages
|
||||
|
||||
def tracked_get_messages(self, session_id, *args, **kwargs):
|
||||
calls.append((kwargs.get("limit"), kwargs.get("offset")))
|
||||
calls.append((kwargs.get("limit"), kwargs.get("after_id")))
|
||||
return original_get_messages(self, session_id, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(SessionDB, "get_messages", tracked_get_messages)
|
||||
|
||||
@@ -3775,6 +3775,19 @@ class TestGetMessagesPagination:
|
||||
assert [m["content"] for m in page2] == ["msg-2", "msg-3", "msg-4", "msg-5"]
|
||||
assert [m["content"] for m in page3] == ["msg-0", "msg-1"]
|
||||
|
||||
def test_after_id_keyset_pages_forward_in_insertion_order(self, db):
|
||||
self._seed(db)
|
||||
page1 = db.get_messages("s1", limit=4, after_id=0)
|
||||
assert [m["content"] for m in page1] == ["msg-0", "msg-1", "msg-2", "msg-3"]
|
||||
page2 = db.get_messages("s1", limit=4, after_id=page1[-1]["id"])
|
||||
assert [m["content"] for m in page2] == ["msg-4", "msg-5", "msg-6", "msg-7"]
|
||||
page3 = db.get_messages("s1", limit=4, after_id=page2[-1]["id"])
|
||||
assert [m["content"] for m in page3] == ["msg-8", "msg-9"]
|
||||
with pytest.raises(ValueError):
|
||||
db.get_messages("s1", limit=4, after_id=0, latest=True)
|
||||
with pytest.raises(ValueError):
|
||||
db.get_messages("s1", limit=4, after_id=0, offset=2)
|
||||
|
||||
def test_resume_safety_counts_active_rows_across_lineage(self, db):
|
||||
db.create_session(session_id="root", source="cli")
|
||||
db.append_messages_batch(
|
||||
|
||||
Reference in New Issue
Block a user