From 951a441021278b3ff6000cfa0673798608c57550 Mon Sep 17 00:00:00 2001 From: Brayo Date: Thu, 24 Sep 2026 17:03:10 +0300 Subject: [PATCH 1/2] perf(datastore): add streaming reads and direct event updates --- aw_core/util.py | 2 +- aw_datastore/datastore.py | 12 ++- aw_datastore/storages/abstract.py | 23 +++++- aw_datastore/storages/memory.py | 12 ++- aw_datastore/storages/peewee.py | 66 +++++++++++++-- aw_datastore/storages/sqlite.py | 63 ++++++++++++-- tests/test_storage_performance.py | 133 ++++++++++++++++++++++++++++++ 7 files changed, 294 insertions(+), 17 deletions(-) create mode 100644 tests/test_storage_performance.py diff --git a/aw_core/util.py b/aw_core/util.py index 9efe5d8f..09a41b37 100644 --- a/aw_core/util.py +++ b/aw_core/util.py @@ -19,6 +19,6 @@ def assert_version(required_version: Tuple[int, ...] = (3, 5)): # pragma: no co ( "Python version {} not supported, you need to upgrade your Python" + " version to at least {}." - ).format(required_version) + ).format(actual_version, required_version) ) logger.debug(f"Python version: {_version_info_tuple()}") diff --git a/aw_datastore/datastore.py b/aw_datastore/datastore.py index 3cb09090..9cb3645a 100644 --- a/aw_datastore/datastore.py +++ b/aw_datastore/datastore.py @@ -34,7 +34,7 @@ def __getitem__(self, bucket_id: str) -> "Bucket": # If this bucket doesn't have a initialized object, create it if bucket_id not in self.bucket_instances: # If the bucket exists in the database, create an object representation of it - if bucket_id in self.buckets(): + if self.has_bucket(bucket_id): bucket = Bucket(self, bucket_id) self.bucket_instances[bucket_id] = bucket else: @@ -72,7 +72,12 @@ def delete_bucket(self, bucket_id: str): del self.bucket_instances[bucket_id] return self.storage_strategy.delete_bucket(bucket_id) - def buckets(self): + def has_bucket(self, bucket_id: str) -> bool: + return self.storage_strategy.has_bucket(bucket_id) + + def buckets(self, include_last_updated: bool = False): + if include_last_updated: + return self.storage_strategy.buckets_with_last_updated() return self.storage_strategy.buckets() @@ -85,6 +90,9 @@ def __init__(self, datastore: Datastore, bucket_id: str) -> None: def metadata(self) -> dict: return self.ds.storage_strategy.get_metadata(self.bucket_id) + def iter_events(self): + return self.ds.storage_strategy.iter_events(self.bucket_id) + def get( self, limit: int = -1, diff --git a/aw_datastore/storages/abstract.py b/aw_datastore/storages/abstract.py index b6bb49ca..cd2e7ea9 100644 --- a/aw_datastore/storages/abstract.py +++ b/aw_datastore/storages/abstract.py @@ -1,6 +1,6 @@ from abc import ABCMeta, abstractmethod from datetime import datetime -from typing import Dict, List, Optional +from typing import Dict, Iterator, List, Optional from aw_core.models import Event @@ -21,6 +21,27 @@ def __init__(self, testing: bool) -> None: def buckets(self) -> Dict[str, dict]: raise NotImplementedError + def has_bucket(self, bucket_id: str) -> bool: + return bucket_id in self.buckets() + + def buckets_with_last_updated(self) -> Dict[str, dict]: + buckets = self.buckets() + for bucket_id, metadata in buckets.items(): + events = self.get_events(bucket_id, 1) + if events: + metadata["last_updated"] = ( + events[0].timestamp + events[0].duration + ).isoformat() + return buckets + + def iter_events(self, bucket_id: str) -> Iterator[Event]: + """Iterate a bucket in the same order as an unbounded get_events call. + + Disk backends override this to avoid materializing all events. Callers + must close the iterator if they stop consuming it early. + """ + yield from self.get_events(bucket_id, -1) + @abstractmethod def create_bucket( self, diff --git a/aw_datastore/storages/memory.py b/aw_datastore/storages/memory.py index 32e39380..fab242b9 100644 --- a/aw_datastore/storages/memory.py +++ b/aw_datastore/storages/memory.py @@ -80,6 +80,14 @@ def buckets(self): buckets[bucket_id] = self.get_metadata(bucket_id) return buckets + def has_bucket(self, bucket_id: str) -> bool: + return bucket_id in self.db + + def iter_events(self, bucket_id): + # Only sort references; copy each event when it is consumed. + for event in sorted(self.db[bucket_id], key=lambda e: e.timestamp)[::-1]: + yield copy.deepcopy(event) + def get_event( self, bucket_id: str, @@ -132,7 +140,7 @@ def get_eventcount( def get_metadata(self, bucket_id: str): if bucket_id in self._metadata: - return self._metadata[bucket_id] + return copy.deepcopy(self._metadata[bucket_id]) else: raise ValueError("Bucket did not exist, could not get metadata") @@ -180,6 +188,8 @@ def replace(self, bucket_id, event_id, event): event = copy.copy(event) event.id = event_id self.db[bucket_id][idx] = event + return True + return False def replace_last(self, bucket_id, event): # NOTE: This does not actually get the most recent event, only the last inserted diff --git a/aw_datastore/storages/peewee.py b/aw_datastore/storages/peewee.py index 1223ef42..5cd9e954 100644 --- a/aw_datastore/storages/peewee.py +++ b/aw_datastore/storages/peewee.py @@ -208,6 +208,52 @@ def update_bucket_keys(self) -> None: def buckets(self) -> Dict[str, Dict[str, Any]]: return {bucket.id: bucket.json() for bucket in BucketModel.select()} + def has_bucket(self, bucket_id: str) -> bool: + return bucket_id in self.bucket_keys + + def buckets_with_last_updated(self) -> Dict[str, Dict[str, Any]]: + # The correlated seek uses (bucket_id, timestamp), including empty buckets. + latest = ( + EventModel.select(EventModel.id) + .where(EventModel.bucket == BucketModel.key) + .order_by(EventModel.timestamp.desc()) + .limit(1) + ) + query = BucketModel.select( + BucketModel, + EventModel.timestamp.alias("last_timestamp"), + EventModel.duration.alias("last_duration"), + ).join(EventModel, peewee.JOIN.LEFT_OUTER, on=(EventModel.id == latest)) + buckets = {} + for row in query.objects().iterator(): + metadata = row.json() + if row.last_timestamp is not None: + event = Event( + timestamp=row.last_timestamp, duration=float(row.last_duration) + ) + metadata["last_updated"] = ( + event.timestamp + event.duration + ).isoformat() + buckets[row.id] = metadata + return buckets + + def iter_events(self, bucket_id): + cursor = self.db.execute_sql( + "SELECT id, timestamp, duration, datastr FROM eventmodel " + "WHERE bucket_id = ? ORDER BY timestamp DESC", + (self.bucket_keys[bucket_id],), + ) + try: + for event_id, timestamp, duration, datastr in cursor: + yield Event( + id=event_id, + timestamp=timestamp, + duration=float(duration), + data=json.loads(datastr), + ) + finally: + cursor.close() + def create_bucket( self, bucket_id: str, @@ -347,13 +393,19 @@ def delete(self, bucket_id, event_id): ) def replace(self, bucket_id, event_id, event): - e = self._get_event(bucket_id, event_id) - e.timestamp = event.timestamp - e.duration = event.duration.total_seconds() - e.datastr = json.dumps(event.data) - e.save() - event.id = e.id - return event + updated = ( + EventModel.update( + timestamp=event.timestamp, + duration=event.duration.total_seconds(), + datastr=json.dumps(event.data), + ) + .where(EventModel.id == event_id) + .where(EventModel.bucket == self.bucket_keys[bucket_id]) + .execute() + ) + if updated: + event.id = event_id + return bool(updated) def get_event( self, diff --git a/aw_datastore/storages/sqlite.py b/aw_datastore/storages/sqlite.py index f139697a..a8f97184 100644 --- a/aw_datastore/storages/sqlite.py +++ b/aw_datastore/storages/sqlite.py @@ -149,6 +149,55 @@ def buckets(self): } return buckets + def has_bucket(self, bucket_id: str) -> bool: + return ( + self.conn.execute( + "SELECT 1 FROM buckets WHERE id = ?", (bucket_id,) + ).fetchone() + is not None + ) + + def buckets_with_last_updated(self): + # Match get_events' ordering by endtime for this backend. + rows = self.conn.execute( + "SELECT b.id, b.name, b.type, b.client, b.hostname, b.created, " + "b.datastr, e.endtime FROM buckets b LEFT JOIN events e ON e.id = " + "(SELECT id FROM events WHERE bucketrow = b.rowid " + "AND endtime >= 0 AND starttime <= ? ORDER BY endtime DESC LIMIT 1)", + (MAX_TIMESTAMP,), + ) + try: + buckets = {} + for row in rows: + metadata = dict( + zip( + ("id", "name", "type", "client", "hostname", "created"), row[:6] + ) + ) + metadata["data"] = json.loads(row[6] or "{}") + if row[7] is not None: + metadata["last_updated"] = datetime.fromtimestamp( + row[7] / 1000000, timezone.utc + ).isoformat() + buckets[row[0]] = metadata + return buckets + finally: + rows.close() + + def iter_events(self, bucket_id): + self.commit() + cursor = self.conn.execute( + "SELECT id, starttime, endtime, datastr FROM events " + "WHERE bucketrow = (SELECT rowid FROM buckets WHERE id = ?) " + "AND endtime >= 0 AND starttime <= ? ORDER BY endtime DESC", + (bucket_id, MAX_TIMESTAMP), + ) + try: + for row in cursor: + yield _rows_to_events([row])[0] + finally: + cursor.close() + def create_bucket( self, bucket_id: str, @@ -301,14 +350,18 @@ def replace(self, bucket_id, event_id, event) -> bool: endtime = starttime + (event.duration.total_seconds() * 1000000) datastr = json.dumps(event.data) query = """UPDATE events - SET bucketrow = (SELECT rowid FROM buckets WHERE id = ?), - starttime = ?, + SET starttime = ?, endtime = ?, datastr = ? - WHERE id = ?""" - self.conn.execute(query, [bucket_id, starttime, endtime, datastr, event_id]) + WHERE id = ? AND bucketrow = + (SELECT rowid FROM buckets WHERE id = ?)""" + cursor = self.conn.execute( + query, [starttime, endtime, datastr, event_id, bucket_id] + ) self.conditional_commit(1) - return True + if cursor.rowcount: + event.id = event_id + return bool(cursor.rowcount) def get_event( self, diff --git a/tests/test_storage_performance.py b/tests/test_storage_performance.py new file mode 100644 index 00000000..bb7738bb --- /dev/null +++ b/tests/test_storage_performance.py @@ -0,0 +1,133 @@ +"""Storage performance contracts, using only disposable databases.""" + +from datetime import datetime, timedelta, timezone +from unittest.mock import Mock + +import pytest +from aw_core.models import Event +from aw_datastore import Datastore, get_storage_methods + + +@pytest.fixture(params=["memory", "peewee", "sqlite"]) +def store(request, tmp_path, monkeypatch): + monkeypatch.setattr( + "aw_datastore.storages.peewee.get_data_dir", lambda _: str(tmp_path) + ) + kwargs = ( + {} if request.param == "memory" else {"filepath": str(tmp_path / "test.db")} + ) + ds = Datastore(get_storage_methods()[request.param], testing=True, **kwargs) + yield ds + if request.param == "peewee": + ds.storage_strategy.db.close() + elif request.param == "sqlite": + ds.storage_strategy.conn.close() + + +def create(ds, bid): + return ds.create_bucket(bid, type="test", client="test", hostname="test") + + +def event(seconds=0, duration=1, label="test"): + return Event( + timestamp=datetime(2024, 1, 1, tzinfo=timezone.utc) + + timedelta(seconds=seconds), + duration=duration, + data={"label": label}, + ) + + +def test_cold_lookup_does_not_enumerate_buckets(store, monkeypatch): + create(store, "a") + store.bucket_instances.clear() + monkeypatch.setattr(store, "buckets", Mock(side_effect=AssertionError("full scan"))) + assert store.has_bucket("a") + assert not store.has_bucket("missing") + assert store["a"].metadata()["id"] == "a" + store.delete_bucket("a") + assert not store.has_bucket("a") + + +def test_listing_matches_latest_event_and_includes_empty_buckets(store): + a = create(store, "a") + create(store, "empty") + a.insert([event(0, duration=30), event(10), event(10, label="tie")]) + latest = a.get(1)[0] + result = store.buckets(include_last_updated=True) + assert ( + result["a"]["last_updated"] == (latest.timestamp + latest.duration).isoformat() + ) + assert "last_updated" not in result["empty"] + assert "last_updated" not in store.buckets()["a"] + + +def test_listing_uses_one_select(store): + if store.storage_strategy.sid == "memory": + return + for i in range(10): + create(store, str(i)).insert(event(i)) + storage = store.storage_strategy + conn = storage.db.connection() if storage.sid == "peewee" else storage.conn + statements = [] + conn.set_trace_callback(statements.append) + try: + assert len(store.buckets(include_last_updated=True)) == 10 + finally: + conn.set_trace_callback(None) + assert sum(s.lstrip().upper().startswith("SELECT") for s in statements) == 1 + + +def test_replace_is_bucket_scoped_and_reports_missing_events(store): + a, b = create(store, "a"), create(store, "b") + original = a.insert(event()) + assert not b.replace(original.id, event(1)) + assert a.get(1)[0].timestamp == original.timestamp + assert not a.replace(9999, event()) + replacement = event(2) + assert a.replace(original.id, replacement) + assert replacement.id == original.id or store.storage_strategy.sid == "memory" + assert a.get(1)[0].timestamp == replacement.timestamp + + +def test_replace_does_not_select_event_first(store): + if store.storage_strategy.sid == "memory": + return + bucket = create(store, "a") + original = bucket.insert(event()) + storage = store.storage_strategy + conn = storage.db.connection() if storage.sid == "peewee" else storage.conn + statements = [] + conn.set_trace_callback(statements.append) + try: + assert bucket.replace(original.id, event(1)) + finally: + conn.set_trace_callback(None) + assert not any(s.lstrip().upper().startswith("SELECT") for s in statements) + assert sum(s.lstrip().upper().startswith("UPDATE") for s in statements) == 1 + + +def test_iterator_preserves_event_data_order_and_duplicate_timestamps(store): + bucket = create(store, "a") + bucket.insert([event(2), event(0, duration=30), event(2, label="tie")]) + expected = [e.to_json_dict() for e in bucket.get()] + iterator = bucket.iter_events() + try: + assert [e.to_json_dict() for e in iterator] == expected + finally: + iterator.close() + assert list(create(store, "empty").iter_events()) == [] + + +def test_disk_iterator_does_not_materialize_get_events(store, monkeypatch): + if store.storage_strategy.sid == "memory": + return + bucket = create(store, "a") + bucket.insert([event(i) for i in range(100)]) + monkeypatch.setattr( + store.storage_strategy, "get_events", Mock(side_effect=AssertionError("eager")) + ) + iterator = bucket.iter_events() + assert next(iterator).timestamp == event(99).timestamp + iterator.close() + # A subsequent write succeeds after an early close. + bucket.insert(event(100)) From a794efb9639b03c243c47c85819f701defbd2591 Mon Sep 17 00:00:00 2001 From: Brayo Date: Thu, 24 Sep 2026 17:13:18 +0300 Subject: [PATCH 2/2] fix(datastore): refresh bucket keys on cold existence checks --- aw_datastore/storages/peewee.py | 13 ++++++++++++- tests/test_storage_performance.py | 23 +++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/aw_datastore/storages/peewee.py b/aw_datastore/storages/peewee.py index 5cd9e954..861d4dd5 100644 --- a/aw_datastore/storages/peewee.py +++ b/aw_datastore/storages/peewee.py @@ -209,7 +209,18 @@ def buckets(self) -> Dict[str, Dict[str, Any]]: return {bucket.id: bucket.json() for bucket in BucketModel.select()} def has_bucket(self, bucket_id: str) -> bool: - return bucket_id in self.bucket_keys + key = ( + BucketModel.select(BucketModel.key) + .where(BucketModel.id == bucket_id) + .scalar() + ) + if key is None: + self.bucket_keys.pop(bucket_id, None) + return False + # A cold lookup may discover a bucket created by another connection. + # Update the key used by subsequent event reads as well as existence. + self.bucket_keys[bucket_id] = key + return True def buckets_with_last_updated(self) -> Dict[str, Dict[str, Any]]: # The correlated seek uses (bucket_id, timestamp), including empty buckets. diff --git a/tests/test_storage_performance.py b/tests/test_storage_performance.py index bb7738bb..eebfdf66 100644 --- a/tests/test_storage_performance.py +++ b/tests/test_storage_performance.py @@ -1,6 +1,7 @@ """Storage performance contracts, using only disposable databases.""" from datetime import datetime, timedelta, timezone +import sqlite3 from unittest.mock import Mock import pytest @@ -131,3 +132,25 @@ def test_disk_iterator_does_not_materialize_get_events(store, monkeypatch): iterator.close() # A subsequent write succeeds after an early close. bucket.insert(event(100)) + + +def test_peewee_cold_lookup_refreshes_externally_created_bucket_key(store): + storage = store.storage_strategy + if storage.sid != "peewee": + return + with sqlite3.connect(storage.db.database) as other: + other.execute( + "INSERT INTO bucketmodel (id, type, client, hostname, created, datastr) " + "VALUES (?, ?, ?, ?, ?, ?)", + ("external", "test", "test", "test", "2024-01-01T00:00:00+00:00", "{}"), + ) + assert "external" not in storage.bucket_keys + bucket = store["external"] + assert bucket.get() == [] + assert bucket.metadata()["id"] == "external" + with sqlite3.connect(storage.db.database) as other: + other.execute("DELETE FROM bucketmodel WHERE id = ?", ("external",)) + store.bucket_instances.clear() + with pytest.raises(KeyError): + store["external"] + assert "external" not in storage.bucket_keys