snapshot of https://github.com/ActivityWatch/aw-core.git (master, 2026-09-16; blobs > 131072B omitted)

tests/test_peewee.py

@@ -0,0 +1,66 @@
+"""
+Peewee-storage-specific tests, mainly for the schema/index migration in
+auto_migrate() and for guarding against query-plan regressions.
+"""
+
+import logging
+
+from aw_datastore.storages import PeeweeStorage
+from aw_datastore.storages import peewee as peewee_storage
+
+from . import context  # noqa: F401
+
+logging.basicConfig(level=logging.DEBUG)
+
+
+def _indexes(db) -> set:
+    return {row[1] for row in db.execute_sql("PRAGMA index_list(eventmodel)")}
+
+
+def test_index_migration():
+    """auto_migrate() replaces the old single-column indexes with the composite one."""
+    storage = PeeweeStorage(testing=True)
+    db = storage.db
+    path = db.database
+
+    assert "eventmodel_bucket_id_timestamp" in _indexes(db)
+
+    # Recreate the pre-migration state: single-column indexes, no composite
+    db.execute_sql("DROP INDEX eventmodel_bucket_id_timestamp")
+    db.execute_sql("CREATE INDEX eventmodel_bucket_id ON eventmodel (bucket_id)")
+    db.execute_sql("CREATE INDEX eventmodel_timestamp ON eventmodel (timestamp)")
+
+    peewee_storage.auto_migrate(path)
+
+    indexes = _indexes(db)
+    assert "eventmodel_bucket_id_timestamp" in indexes
+    assert "eventmodel_bucket_id" not in indexes
+    assert "eventmodel_timestamp" not in indexes
+
+
+def test_query_plan_uses_composite_index():
+    """
+    The last-event query (issued by replace_last on every merged heartbeat) and
+    range queries must be served by the composite index in a single pass.
+
+    Without it, SQLite reads *all* of a bucket's rows and sorts them in a temp
+    B-tree even for LIMIT 1, which on multi-million-event buckets meant seconds
+    of CPU/IO per heartbeat (and terabytes read per day).
+    """
+    storage = PeeweeStorage(testing=True)
+    db = storage.db
+
+    queries = [
+        # replace_last/_get_last
+        "SELECT * FROM eventmodel WHERE bucket_id = 1"
+        " ORDER BY timestamp DESC LIMIT 1",
+        # get_events with a time range
+        "SELECT * FROM eventmodel WHERE bucket_id = 1"
+        " AND timestamp >= '2026-01-01' AND timestamp <= '2026-01-02'"
+        " ORDER BY timestamp DESC LIMIT 100",
+    ]
+    for query in queries:
+        rows = db.execute_sql(f"EXPLAIN QUERY PLAN {query}").fetchall()
+        plan = "\n".join(str(row) for row in rows)
+        assert "eventmodel_bucket_id_timestamp" in plan, plan
+        assert "TEMP B-TREE" not in plan, plan