From fb458358b0d0218602fb9cd76535c2464f01f628 Mon Sep 17 00:00:00 2001 From: Brayo Date: Thu, 24 Sep 2026 16:48:16 +0300 Subject: [PATCH 1/3] perf: reuse bucket handles and defer debug formatting --- aw_server/api.py | 50 +++++++++++++------------- aw_server/rest.py | 16 +++++---- tests/test_server.py | 83 +++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 116 insertions(+), 33 deletions(-) diff --git a/aw_server/api.py b/aw_server/api.py index 20319ea..61df391 100644 --- a/aw_server/api.py +++ b/aw_server/api.py @@ -43,8 +43,14 @@ def get_device_id() -> str: def check_bucket_exists(f): @functools.wraps(f) def g(self, bucket_id, *args, **kwargs): - if bucket_id not in self.db.buckets(): - raise NotFound("NoSuchBucket", f"There's no bucket named {bucket_id}") + # Datastore caches bucket handles and invalidates them on deletion. + # Reuse that lookup instead of loading every bucket's metadata per call. + try: + self.db[bucket_id] + except KeyError: + raise NotFound( + "NoSuchBucket", f"There's no bucket named {bucket_id}" + ) from None return f(self, bucket_id, *args, **kwargs) return g @@ -198,7 +204,7 @@ def update_bucket( def delete_bucket(self, bucket_id: str) -> None: """Delete a bucket""" self.db.delete_bucket(bucket_id) - logger.debug(f"Deleted bucket '{bucket_id}'") + logger.debug("Deleted bucket '%s'", bucket_id) return None @check_bucket_exists @@ -209,7 +215,7 @@ def get_event( ) -> Optional[Event]: """Get a single event from a bucket""" logger.debug( - f"Received get request for event {event_id} in bucket '{bucket_id}'" + "Received get request for event %s in bucket '%s'", event_id, bucket_id ) event = self.db[bucket_id].get_by_id(event_id) return event.to_json_dict() if event else None @@ -223,7 +229,7 @@ def get_events( end: Optional[datetime] = None, ) -> List[Event]: """Get events from a bucket""" - logger.debug(f"Received get request for events in bucket '{bucket_id}'") + logger.debug("Received get request for events in bucket '%s'", bucket_id) if limit is None: # Let limit = None also mean "no limit" limit = -1 events = [ @@ -255,7 +261,7 @@ def get_eventcount( end: Optional[datetime] = None, ) -> int: """Get eventcount from a bucket""" - logger.debug(f"Received get request for eventcount in bucket '{bucket_id}'") + logger.debug("Received get request for eventcount in bucket '%s'", bucket_id) return self.db[bucket_id].get_eventcount(start, end) @check_bucket_exists @@ -287,13 +293,12 @@ def heartbeat(self, bucket_id: str, heartbeat: Event, pulsetime: float) -> Event Inspired by: https://wakatime.com/developers#heartbeats """ logger.debug( - "Received heartbeat in bucket '{}'\n\ttimestamp: {}, duration: {}, pulsetime: {}\n\tdata: {}".format( - bucket_id, - heartbeat.timestamp, - heartbeat.duration, - pulsetime, - heartbeat.data, - ) + "Received heartbeat in bucket '%s'\n\ttimestamp: %s, duration: %s, pulsetime: %s\n\tdata: %s", + bucket_id, + heartbeat.timestamp, + heartbeat.duration, + pulsetime, + heartbeat.data, ) # The endtime here is set such that in the event that the heartbeat is older than an @@ -319,30 +324,25 @@ def heartbeat(self, bucket_id: str, heartbeat: Event, pulsetime: float) -> Event if merged is not None: # Heartbeat was merged into last_event logger.debug( - "Received valid heartbeat, merging. (bucket: {})".format( - bucket_id - ) + "Received valid heartbeat, merging. (bucket: %s)", bucket_id ) self.last_event[bucket_id] = merged self.db[bucket_id].replace_last(merged) return merged else: logger.info( - "Received heartbeat after pulse window, inserting as new event. (bucket: {})".format( - bucket_id - ) + "Received heartbeat after pulse window, inserting as new event. (bucket: %s)", + bucket_id, ) else: logger.debug( - "Received heartbeat with differing data, inserting as new event. (bucket: {})".format( - bucket_id - ) + "Received heartbeat with differing data, inserting as new event. (bucket: %s)", + bucket_id, ) else: logger.info( - "Received heartbeat, but bucket was previously empty, inserting as new event. (bucket: {})".format( - bucket_id - ) + "Received heartbeat, but bucket was previously empty, inserting as new event. (bucket: %s)", + bucket_id, ) self.db[bucket_id].insert(heartbeat) diff --git a/aw_server/rest.py b/aw_server/rest.py index 0b2c320..79855ac 100644 --- a/aw_server/rest.py +++ b/aw_server/rest.py @@ -214,9 +214,9 @@ def get(self, bucket_id): def post(self, bucket_id): data = request.get_json() logger.debug( - "Received post request for event in bucket '{}' and data: {}".format( - bucket_id, data - ) + "Received post request for event in bucket '%s' and data: %s", + bucket_id, + data, ) if isinstance(data, dict): @@ -251,7 +251,9 @@ class EventResource(Resource): @copy_doc(ServerAPI.get_event) def get(self, bucket_id: str, event_id: int): logger.debug( - f"Received get request for event with id '{event_id}' in bucket '{bucket_id}'" + "Received get request for event with id '%s' in bucket '%s'", + event_id, + bucket_id, ) event = current_app.api.get_event(bucket_id, event_id) if event: @@ -262,9 +264,9 @@ def get(self, bucket_id: str, event_id: int): @copy_doc(ServerAPI.delete_event) def delete(self, bucket_id: str, event_id: int): logger.debug( - "Received delete request for event with id '{}' in bucket '{}'".format( - event_id, bucket_id - ) + "Received delete request for event with id '%s' in bucket '%s'", + event_id, + bucket_id, ) success = current_app.api.delete_event(bucket_id, event_id) return {"success": success}, 200 diff --git a/tests/test_server.py b/tests/test_server.py index 3549f7e..300bdad 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -1,7 +1,13 @@ import random -from datetime import datetime, timedelta +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 + +from aw_server.api import ServerAPI +from aw_server.exceptions import NotFound @pytest.fixture() @@ -182,3 +188,78 @@ def test_query_valid_timeperiod(flask_client): ) assert r.status_code == 200 assert r.json == [1] + + +@pytest.fixture(params=["memory", "peewee", "sqlite"]) +def isolated_api(request, tmp_path, monkeypatch): + """Exercise the storage backends without accessing a user's database/settings.""" + monkeypatch.setattr("aw_server.settings.get_config_dir", lambda _: str(tmp_path)) + monkeypatch.setattr( + "aw_datastore.storages.peewee.get_data_dir", + lambda _: str(tmp_path), + ) + storage = get_storage_methods()[request.param] + kwargs = ( + {} if request.param == "memory" else {"filepath": str(tmp_path / "test.db")} + ) + db = Datastore(storage, testing=True, **kwargs) + try: + yield ServerAPI(db, testing=True) + finally: + if request.param == "peewee": + db.storage_strategy.db.close() + elif request.param == "sqlite": + db.storage_strategy.conn.close() + + +def test_bucket_checks_reuse_datastore_lookup(isolated_api, monkeypatch): + api = isolated_api + api.create_bucket("test", "test", "test", "test") + # Simulate the first access to a bucket that existed before server startup. + api.db.bucket_instances.clear() + listing = Mock(wraps=api.db.buckets) + monkeypatch.setattr(api.db, "buckets", listing) + + assert api.get_events("test") == [] + assert listing.call_count == 1 + listing.reset_mock() + + timestamp = datetime(2024, 1, 1, tzinfo=timezone.utc) + api.heartbeat("test", Event(timestamp=timestamp, data={"app": "test"}), 60) + merged = api.heartbeat( + "test", + Event(timestamp=timestamp + timedelta(seconds=1), data={"app": "test"}), + 60, + ) + assert merged.duration == timedelta(seconds=1) + assert api.get_eventcount("test") == 1 + assert len(api.get_events("test")) == 1 + assert len(api.export_bucket("test")["events"]) == 1 + listing.assert_not_called() + + +def test_bucket_checks_follow_datastore_lifecycle(isolated_api): + api = isolated_api + with pytest.raises(NotFound, match="There's no bucket named test"): + api.get_events("test") + + api.create_bucket("test", "test", "test", "test") + assert api.get_events("test") == [] + api.delete_bucket("test") + with pytest.raises(NotFound, match="There's no bucket named test"): + api.get_events("test") + + # Changes made through Datastore must also be visible to the API. + api.db.create_bucket("test", type="test", client="test", hostname="test") + assert api.get_events("test") == [] + api.db.delete_bucket("test") + with pytest.raises(NotFound): + api.get_events("test") + + +def test_bucket_check_does_not_mask_operation_errors(isolated_api, monkeypatch): + api = isolated_api + api.create_bucket("test", "test", "test", "test") + monkeypatch.setattr(api.db["test"], "get", Mock(side_effect=KeyError("event data"))) + with pytest.raises(KeyError, match="event data"): + api.get_events("test") From 172b7d667761bf079b7a396c0cb1069604b86d31 Mon Sep 17 00:00:00 2001 From: Brayo Date: Thu, 24 Sep 2026 17:06:44 +0300 Subject: [PATCH 2/3] perf: stream exports and cache queries with mutation invalidation --- Makefile | 2 +- aw_server/api.py | 184 +++++++++++++++------- aw_server/query_cache.py | 192 +++++++++++++++++++++++ aw_server/rest.py | 27 ++-- poetry.lock | 24 +-- pyproject.toml | 2 +- tests/conftest.py | 24 +++ tests/test_performance.py | 318 ++++++++++++++++++++++++++++++++++++++ tests/test_server.py | 26 +--- 9 files changed, 699 insertions(+), 100 deletions(-) create mode 100644 aw_server/query_cache.py create mode 100644 tests/test_performance.py diff --git a/Makefile b/Makefile index 2a58b51..5c2721a 100644 --- a/Makefile +++ b/Makefile @@ -24,7 +24,7 @@ test: @# Note that extensive integration tests are also run in the bundle repo, @# for both aw-server and aw-server-rust, but without code coverage. python -c 'import aw_server' - python -m pytest tests/test_server.py tests/test_profile.py tests/test_profile_config.py + python -m pytest tests/test_server.py tests/test_profile.py tests/test_profile_config.py tests/test_performance.py typecheck: python -m mypy aw_server tests --ignore-missing-imports diff --git a/aw_server/api.py b/aw_server/api.py index 61df391..3da07cc 100644 --- a/aw_server/api.py +++ b/aw_server/api.py @@ -1,9 +1,12 @@ +import copy import functools import json import logging from datetime import datetime from pathlib import Path from socket import gethostname +from threading import RLock +from contextlib import closing from typing import ( Any, Dict, @@ -23,6 +26,7 @@ from .__about__ import __version__ from .exceptions import NotFound from .profile import profile_from_env +from .query_cache import QueryCache, ReadTracker from .settings import Settings logger = logging.getLogger(__name__) @@ -56,6 +60,21 @@ def g(self, bucket_id, *args, **kwargs): return g +def bucket_mutation(metadata=False): + """Serialize writes with heartbeat merging and invalidate dependent queries.""" + + def decorate(f): + @functools.wraps(f) + def wrapped(self, bucket_id, *args, **kwargs): + with self._write_lock, self.query_cache.mutation(bucket_id, metadata): + self.last_event.pop(bucket_id, None) + return f(self, bucket_id, *args, **kwargs) + + return wrapped + + return decorate + + class ServerAPI: def __init__(self, db, testing) -> None: self.db = db @@ -63,6 +82,8 @@ def __init__(self, db, testing) -> None: self.testing = testing self.profile = profile_from_env(testing=testing) self.last_event = {} # type: dict + self._write_lock = RLock() + self.query_cache = QueryCache() def get_info(self) -> Dict[str, Any]: """Get server info""" @@ -78,15 +99,7 @@ def get_info(self) -> Dict[str, Any]: def get_buckets(self) -> Dict[str, Dict]: """Get dict {bucket_name: Bucket} of all buckets""" logger.debug("Received get request for buckets") - buckets = self.db.buckets() - for b in buckets: - # TODO: Move this code to aw-core? - last_events = self.db[b].get(limit=1) - if len(last_events) > 0: - last_event = last_events[0] - last_updated = last_event.timestamp + last_event.duration - buckets[b]["last_updated"] = last_updated.isoformat() - return buckets + return self.db.buckets(include_last_updated=True) @check_bucket_exists def get_bucket_metadata(self, bucket_id: str) -> Dict[str, Any]: @@ -106,44 +119,93 @@ def export_bucket(self, bucket_id: str) -> Dict[str, Any]: def export_all(self) -> Dict[str, Any]: """Exports all buckets and their events to a format consistent across versions""" - buckets = self.get_buckets() + buckets = self.db.buckets() exported_buckets = {} for bid in buckets.keys(): exported_buckets[bid] = self.export_bucket(bid) return exported_buckets + def stream_export(self, bucket_id=None): + # Validate before returning the generator so a missing bucket is a 404, + # not an exception after the HTTP response has started. + if bucket_id is not None: + buckets = {bucket_id: self.get_bucket_metadata(bucket_id)} + else: + buckets = self.db.buckets() + + def generate(): + yield '{"buckets":{' + for index, (bid, metadata) in enumerate(buckets.items()): + if index: + yield "," + yield json.dumps(bid) + ":" + yield json.dumps(metadata)[:-1] + ',"events":[' + with closing(self.db[bid].iter_events()) as events: + for event_index, event in enumerate(events): + if event_index: + yield "," + payload = event.to_json_dict() + payload.pop("id", None) + yield json.dumps(payload) + yield "]}" + yield "}}" + + def buffered(): + # Avoid a socket write for every separator/small event while retaining + # only a bounded batch (plus the largest individual event). + with closing(generate()) as fragments: + batch = [] + size = 0 + for fragment in fragments: + batch.append(fragment) + size += len(fragment) + if size >= 64 * 1024: + yield "".join(batch) + batch = [] + size = 0 + if batch: + yield "".join(batch) + + return buffered() + def import_bucket(self, bucket_data: Any): bucket_id = bucket_data["id"] - logger.info(f"Importing bucket {bucket_id}") + with self._write_lock, self.query_cache.mutation(bucket_id, metadata=True): + self.last_event.pop(bucket_id, None) + logger.info(f"Importing bucket {bucket_id}") - # TODO: Check that bucket doesn't already exist - self.db.create_bucket( - bucket_id, - type=bucket_data["type"], - client=bucket_data["client"], - hostname=bucket_data["hostname"], - created=( - bucket_data["created"] - if isinstance(bucket_data["created"], datetime) - else iso8601.parse_date(bucket_data["created"]) - ), - ) + # TODO: Check that bucket doesn't already exist + self.db.create_bucket( + bucket_id, + type=bucket_data["type"], + client=bucket_data["client"], + hostname=bucket_data["hostname"], + created=( + bucket_data["created"] + if isinstance(bucket_data["created"], datetime) + else iso8601.parse_date(bucket_data["created"]) + ), + ) - # scrub IDs from events - # (otherwise causes weird bugs with no events seemingly imported when importing events exported from aw-server-rust, which contains IDs) - for event in bucket_data["events"]: - if "id" in event: - del event["id"] + # scrub IDs from events + # (otherwise causes weird bugs with no events seemingly imported when importing events exported from aw-server-rust, which contains IDs) + for event in bucket_data["events"]: + if "id" in event: + del event["id"] - self.create_events( - bucket_id, - [Event(**e) if isinstance(e, dict) else e for e in bucket_data["events"]], - ) + self.create_events( + bucket_id, + [ + Event(**e) if isinstance(e, dict) else e + for e in bucket_data["events"] + ], + ) def import_all(self, buckets: Dict[str, Any]): for bid, bucket in buckets.items(): self.import_bucket(bucket) + @bucket_mutation(metadata=True) def create_bucket( self, bucket_id: str, @@ -163,7 +225,7 @@ def create_bucket( """ if created is None: created = datetime.now() - if bucket_id in self.db.buckets(): + if self.db.has_bucket(bucket_id): return False if hostname == "!local": info = self.get_info() @@ -181,6 +243,7 @@ def create_bucket( ) return True + @bucket_mutation(metadata=True) @check_bucket_exists def update_bucket( self, @@ -193,13 +256,14 @@ def update_bucket( """Update bucket metadata""" self.db.update_bucket( bucket_id, - type=event_type, + type_id=event_type, client=client, hostname=hostname, data=data, ) return None + @bucket_mutation(metadata=True) @check_bucket_exists def delete_bucket(self, bucket_id: str) -> None: """Delete a bucket""" @@ -237,6 +301,7 @@ def get_events( ] return events + @bucket_mutation(metadata=False) @check_bucket_exists def create_events(self, bucket_id: str, events: List[Event]) -> List[Event]: """Create events for a bucket. Can handle both single events and multiple ones. @@ -264,12 +329,12 @@ def get_eventcount( logger.debug("Received get request for eventcount in bucket '%s'", bucket_id) return self.db[bucket_id].get_eventcount(start, end) + @bucket_mutation(metadata=False) @check_bucket_exists def delete_event(self, bucket_id: str, event_id) -> bool: """Delete a single event from a bucket""" return self.db[bucket_id].delete(event_id) - @check_bucket_exists def heartbeat(self, bucket_id: str, heartbeat: Event, pulsetime: float) -> Event: """ Heartbeats are useful when implementing watchers that simply keep @@ -292,6 +357,15 @@ def heartbeat(self, bucket_id: str, heartbeat: Event, pulsetime: float) -> Event Inspired by: https://wakatime.com/developers#heartbeats """ + with self._write_lock, self.query_cache.mutation(bucket_id) as change: + event = self._heartbeat(bucket_id, heartbeat, pulsetime) + # Invalidate every interval touched by the event, including the old + # duration when a heartbeat is merged. Other days stay cached. + change.span = (event.timestamp, event.timestamp + event.duration) + return event + + @check_bucket_exists + def _heartbeat(self, bucket_id: str, heartbeat: Event, pulsetime: float) -> Event: logger.debug( "Received heartbeat in bucket '%s'\n\ttimestamp: %s, duration: %s, pulsetime: %s\n\tdata: %s", bucket_id, @@ -301,22 +375,13 @@ def heartbeat(self, bucket_id: str, heartbeat: Event, pulsetime: float) -> Event heartbeat.data, ) - # The endtime here is set such that in the event that the heartbeat is older than an - # existing event we should try to merge it with the last event before the heartbeat instead. - # FIXME: This (the endtime=heartbeat.timestamp) gets rid of the "heartbeat was older than last event" - # warning and also causes a already existing "newer" event to be overwritten in the - # replace_last call below. This is problematic. - # Solution: This could be solved if we were able to replace arbitrary events. - # That way we could double check that the event has been applied - # and if it hasn't we simply replace it with the updated counterpart. - last_event = None if bucket_id not in self.last_event: last_events = self.db[bucket_id].get(limit=1) if len(last_events) > 0: last_event = last_events[0] else: - last_event = self.last_event[bucket_id] + last_event = copy.copy(self.last_event[bucket_id]) if last_event: if last_event.data == heartbeat.data: @@ -326,9 +391,13 @@ def heartbeat(self, bucket_id: str, heartbeat: Event, pulsetime: float) -> Event logger.debug( "Received valid heartbeat, merging. (bucket: %s)", bucket_id ) - self.last_event[bucket_id] = merged - self.db[bucket_id].replace_last(merged) - return merged + if merged.id is not None and self.db[bucket_id].replace( + merged.id, merged + ): + self.last_event[bucket_id] = merged + return merged + # The row may have been removed by another datastore user. + self.last_event.pop(bucket_id, None) else: logger.info( "Received heartbeat after pulse window, inserting as new event. (bucket: %s)", @@ -345,12 +414,13 @@ def heartbeat(self, bucket_id: str, heartbeat: Event, pulsetime: float) -> Event bucket_id, ) - self.db[bucket_id].insert(heartbeat) - self.last_event[bucket_id] = heartbeat - return heartbeat + inserted = self.db[bucket_id].insert(heartbeat) + self.last_event[bucket_id] = inserted + return inserted def query2(self, name, query, timeperiods, cache): result = [] + query = "".join(query) for timeperiod in timeperiods: period = timeperiod.split("/")[ :2 @@ -365,8 +435,16 @@ def query2(self, name, query, timeperiods, cache): endtime = iso8601.parse_date(period[1]) except iso8601.ParseError as e: raise QueryException(f"Invalid timeperiod '{timeperiod}': {e}") - query = "".join(query) - result.append(query2.query(name, query, starttime, endtime, self.db)) + if not cache: + result.append(query2.query(name, query, starttime, endtime, self.db)) + continue + key = (name, query, starttime, endtime) + hit, value, revision = self.query_cache.lookup(key) + if not hit: + tracker = ReadTracker(self.db) + value = query2.query(name, query, starttime, endtime, tracker) + self.query_cache.store(key, value, tracker, revision) + result.append(value) return result # TODO: Right now the log format on disk has to be JSON, this is hard to read by humans... diff --git a/aw_server/query_cache.py b/aw_server/query_cache.py new file mode 100644 index 0000000..3ea5970 --- /dev/null +++ b/aw_server/query_cache.py @@ -0,0 +1,192 @@ +"""Bounded query results with dependencies on the actual datastore reads.""" + +import copy +import sys +from collections import OrderedDict +from contextlib import contextmanager +from dataclasses import dataclass, field +from datetime import datetime, timedelta +from threading import RLock +from time import monotonic +from typing import Any, List, Optional, Tuple + +Read = Tuple[str, Optional[datetime], Optional[datetime]] + + +def _size(value, seen=None): + """Account for retained Python objects, including keys and nested payloads.""" + if seen is None: + seen = set() + if id(value) in seen: + return 0 + seen.add(id(value)) + size = sys.getsizeof(value) + if isinstance(value, dict): + size += sum(_size(k, seen) + _size(v, seen) for k, v in value.items()) + elif isinstance(value, (list, tuple, set)): + size += sum(_size(v, seen) for v in value) + return size + + +@dataclass +class Change: + span: Optional[Tuple[datetime, datetime]] = None + + +@dataclass +class Entry: + result: Any + reads: List[Read] + metadata: bool + expires: float + size: int + + +class QueryCache: + def __init__(self, max_entries=128, max_bytes=8 * 1024 * 1024, ttl=300): + self.max_entries = max_entries + self.max_bytes = max_bytes + self.ttl = ttl + self.entries: OrderedDict = OrderedDict() + self.bytes = 0 + self.revision = 0 + self.writers = 0 + self.lock = RLock() + + def _remove(self, key): + self.bytes -= self.entries.pop(key).size + + def lookup(self, key): + with self.lock: + for expired in [ + k for k, e in self.entries.items() if e.expires <= monotonic() + ]: + self._remove(expired) + if self.writers: + return False, None, None + entry = self.entries.get(key) + if entry is not None: + self.entries.move_to_end(key) + return True, copy.deepcopy(entry.result), self.revision + return False, None, self.revision + + def store(self, key, result, tracker, revision): + if revision is None or not tracker.cacheable or self.max_entries <= 0: + return + # Do the potentially expensive copy outside the cache lock. + size = _size((key, result, tracker.reads)) + if size > self.max_bytes: + return + snapshot = copy.deepcopy(result) + with self.lock: + if self.writers or self.revision != revision: + return + if key in self.entries: + self._remove(key) + while self.entries and ( + len(self.entries) >= self.max_entries + or self.bytes + size > self.max_bytes + ): + self._remove(next(iter(self.entries))) + self.entries[key] = Entry( + snapshot, + list(tracker.reads), + tracker.metadata, + monotonic() + self.ttl, + size, + ) + self.bytes += size + + @contextmanager + def mutation(self, bucket_id, metadata=False): + change = Change() + with self.lock: + self.writers += 1 + self.revision += 1 + try: + yield change + except BaseException: + # A failed bulk operation can still have made partial changes. + change.span = None + raise + finally: + with self.lock: + for key, entry in list(self.entries.items()): + affected = metadata and entry.metadata + for bid, start, end in entry.reads: + if bid != bucket_id: + continue + if metadata or change.span is None: + affected = True + else: + lo, hi = sorted(change.span) + affected |= (end is None or lo <= end) and ( + start is None or start <= hi + ) + if affected: + self._remove(key) + self.writers -= 1 + self.revision += 1 + + +@dataclass +class ReadTracker: + datastore: Any + reads: List[Read] = field(default_factory=list) + metadata: bool = False + cacheable: bool = True + + def buckets(self): + self.metadata = True + return self.datastore.buckets() + + def __getitem__(self, bucket_id): + return ReadBucket(self.datastore[bucket_id], self, bucket_id) + + def __getattr__(self, name): + # Custom query functions using other datastore operations must not cache + # a result whose dependencies we cannot describe. + self.cacheable = False + return getattr(self.datastore, name) + + +class ReadBucket: + def __init__(self, bucket, tracker, bucket_id): + self.bucket = bucket + self.tracker = tracker + self.bucket_id = bucket_id + + def _record(self, starttime, endtime): + # Bucket.get rounds start down and end up to millisecond boundaries. + padding = timedelta(milliseconds=1) + try: + start = starttime - padding if starttime else None + except OverflowError: + start = None + try: + end = endtime + padding if endtime else None + except OverflowError: + end = None + self.tracker.reads.append( + ( + self.bucket_id, + start, + end, + ) + ) + + def get(self, limit=-1, starttime=None, endtime=None): + self._record(starttime, endtime) + return self.bucket.get(limit, starttime, endtime) + + def get_eventcount(self, starttime=None, endtime=None): + self._record(starttime, endtime) + return self.bucket.get_eventcount(starttime, endtime) + + def metadata(self): + self.tracker.metadata = True + return self.bucket.metadata() + + def __getattr__(self, name): + self.tracker.cacheable = False + return getattr(self.bucket, name) diff --git a/aw_server/rest.py b/aw_server/rest.py index 79855ac..4132de6 100644 --- a/aw_server/rest.py +++ b/aw_server/rest.py @@ -12,7 +12,8 @@ Blueprint, current_app, jsonify, - make_response, + Response, + stream_with_context, request, ) from flask_restx import Api, Resource, fields @@ -318,15 +319,19 @@ def post(self, bucket_id): class QueryResource(Resource): # TODO Docs @api.expect(query, validate=True) - @api.param("name", "Name of the query (required if using cache)") + @api.param("name", "Name of the query") + @api.param("cache", "Cache query results (default: 1; set to 0 to bypass)") def post(self): name = "" if "name" in request.args: name = request.args["name"] query = request.get_json() + cache_arg = request.args.get("cache", "1").lower() + if cache_arg not in ("0", "1", "false", "true"): + raise BadRequest("InvalidParameter", "cache must be 0, 1, false, or true") try: result = current_app.api.query2( - name, query["query"], query["timeperiods"], False + name, query["query"], query["timeperiods"], cache_arg in ("1", "true") ) return jsonify(result) except QueryException as qe: @@ -342,9 +347,10 @@ class ExportAllResource(Resource): @api.doc(model=buckets_export) @copy_doc(ServerAPI.export_all) def get(self): - buckets_export = current_app.api.export_all() - payload = {"buckets": buckets_export} - response = make_response(json.dumps(payload)) + response = Response( + stream_with_context(current_app.api.stream_export()), + mimetype="application/json", + ) filename = "aw-buckets-export.json" response.headers["Content-Disposition"] = "attachment; filename={}".format( filename @@ -358,10 +364,11 @@ class BucketExportResource(Resource): @api.doc(model=buckets_export) @copy_doc(ServerAPI.export_bucket) def get(self, bucket_id): - bucket_export = current_app.api.export_bucket(bucket_id) - payload = {"buckets": {bucket_export["id"]: bucket_export}} - response = make_response(json.dumps(payload)) - filename = "aw-bucket-export_{}.json".format(bucket_export["id"]) + response = Response( + stream_with_context(current_app.api.stream_export(bucket_id)), + mimetype="application/json", + ) + filename = "aw-bucket-export_{}.json".format(bucket_id) response.headers["Content-Disposition"] = "attachment; filename={}".format( filename ) diff --git a/poetry.lock b/poetry.lock index 5fb760a..e0e1ddc 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.3.2 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.1.3 and should not be changed by hand. [[package]] name = "aniso8601" @@ -59,25 +59,29 @@ name = "aw-core" version = "0.5.18" description = "Core library for ActivityWatch" optional = false -python-versions = "<4.0,>=3.8" +python-versions = "^3.8" groups = ["main"] -files = [ - {file = "aw_core-0.5.18-py3-none-any.whl", hash = "sha256:6e3df7b46e913d631e0dc8b1b10046c111284a9f4240bec8efd279ce4e438cd5"}, - {file = "aw_core-0.5.18.tar.gz", hash = "sha256:dc14867ed15b4529d6e69b1d07df3fb29c91ccf7ab174de3f9100b90a95cc826"}, -] +files = [] +develop = false [package.dependencies] click = "*" deprecation = "*" iso8601 = "*" -jsonschema = ">=4.3,<5.0" +jsonschema = "^4.3" peewee = ">=3.0.0,<5.0.0" platformdirs = "3.10" -rfc3339-validator = ">=0.1.4,<0.2.0" -strict-rfc3339 = ">=0.7,<0.8" +rfc3339-validator = "^0.1.4" +strict-rfc3339 = "^0.7" timeslot = "*" tomlkit = "*" +[package.source] +type = "git" +url = "https://github.com/0xbrayo/aw-core.git" +reference = "951a441021278b3ff6000cfa0673798608c57550" +resolved_reference = "951a441021278b3ff6000cfa0673798608c57550" + [[package]] name = "black" version = "23.10.1" @@ -1192,4 +1196,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [metadata] lock-version = "2.1" python-versions = "^3.8" -content-hash = "370151e1033b3e3d1c635d9d6ce131af2633bbf6642c9d4e3d11baa38258d60f" +content-hash = "222fa6b67a8dc48a53002ff008c5751217de2c1c24fd9eed28c77b50c1b07cf9" diff --git a/pyproject.toml b/pyproject.toml index 02fb20c..3bcfc20 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,7 @@ aw-server = "aw_server:main" [tool.poetry.dependencies] python = "^3.8" -aw-core = "^0.5.18" +aw-core = { git = "https://github.com/0xbrayo/aw-core.git", rev = "951a441021278b3ff6000cfa0673798608c57550" } aw-client = "^0.5.8" flask = "^2.2" flask-restx = "^1.0.3" diff --git a/tests/conftest.py b/tests/conftest.py index 7e32404..2c82673 100755 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,8 @@ import os import pytest +from aw_datastore import Datastore, get_storage_methods +from aw_server.api import ServerAPI from aw_client import ActivityWatchClient from aw_server.server import AWFlask @@ -46,3 +48,25 @@ def aw_client(): for bucket_id in buckets: if bucket_id.startswith("test-"): c.delete_bucket(bucket_id) + + +@pytest.fixture(params=["memory", "peewee", "sqlite"]) +def isolated_api(request, tmp_path, monkeypatch): + """Exercise the storage backends without accessing a user's database/settings.""" + monkeypatch.setattr("aw_server.settings.get_config_dir", lambda _: str(tmp_path)) + monkeypatch.setattr( + "aw_datastore.storages.peewee.get_data_dir", + lambda _: str(tmp_path), + ) + storage = get_storage_methods()[request.param] + kwargs = ( + {} if request.param == "memory" else {"filepath": str(tmp_path / "test.db")} + ) + db = Datastore(storage, testing=True, **kwargs) + try: + yield ServerAPI(db, testing=True) + finally: + if request.param == "peewee": + db.storage_strategy.db.close() + elif request.param == "sqlite": + db.storage_strategy.conn.close() diff --git a/tests/test_performance.py b/tests/test_performance.py new file mode 100644 index 0000000..99687e7 --- /dev/null +++ b/tests/test_performance.py @@ -0,0 +1,318 @@ +import json +import logging +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timedelta, timezone +from unittest.mock import Mock + +import pytest +from aw_core.models import Event +from aw_query.exceptions import QueryException + +from aw_server import api as api_module +from aw_server import rest +from aw_server.query_cache import QueryCache, ReadTracker + +START = datetime(2024, 1, 1, tzinfo=timezone.utc) +PERIOD = "2024-01-01T00:00:00+00:00/2024-01-02T00:00:00+00:00" +NEXT_PERIOD = "2024-01-02T00:00:00+00:00/2024-01-03T00:00:00+00:00" +QUERY = ['RETURN = query_bucket("test");'] + + +def setup_bucket(api): + api.create_bucket("test", "test", "test", "test") + return api.create_events( + "test", [Event(timestamp=START, duration=1, data={"app": "a"})] + )[0] + + +def run(api, query=QUERY, period=PERIOD, cache=True): + return api.query2("test", query, [period], cache)[0] + + +def spy_queries(monkeypatch): + spy = Mock(wraps=api_module.query2.query) + monkeypatch.setattr(api_module.query2, "query", spy) + return spy + + +def test_query_cache_reuses_result_without_sharing_mutable_events( + isolated_api, monkeypatch +): + api = isolated_api + setup_bucket(api) + spy = spy_queries(monkeypatch) + first = run(api) + first[0].data["app"] = "changed" + assert run(api)[0].data["app"] == "a" + assert spy.call_count == 1 + run(api, cache=False) + assert spy.call_count == 2 + run(api, query=["RETURN = 123;"]) + assert spy.call_count == 3 + + +@pytest.mark.parametrize( + "operation", ["insert", "upsert", "delete_event", "delete_bucket", "metadata"] +) +def test_query_cache_invalidates_mutations(isolated_api, monkeypatch, operation): + api = isolated_api + event = setup_bucket(api) + spy = spy_queries(monkeypatch) + run(api) + if operation == "insert": + api.create_events( + "test", [Event(timestamp=START + timedelta(seconds=2), data={"app": "b"})] + ) + elif operation == "upsert": + # Bulk upserts can move events out of a cached period. + api.create_events( + "test", + [ + Event( + id=event.id, timestamp=START + timedelta(days=2), data={"app": "b"} + ), + Event(timestamp=START + timedelta(days=3)), + ], + ) + elif operation == "delete_event": + api.delete_event("test", event.id) + elif operation == "delete_bucket": + api.delete_bucket("test") + api.create_bucket("test", "test", "test", "test") + else: + api.update_bucket("test", data={"label": "new"}) + assert run(api) == run(api, cache=False) + assert spy.call_count == 3 + + +def test_heartbeat_preserves_unaffected_days_and_invalidates_current_day( + isolated_api, monkeypatch +): + api = isolated_api + setup_bucket(api) + current = START + timedelta(days=1, hours=1) + api.heartbeat("test", Event(timestamp=current, data={"app": "live"}), 60) + spy = spy_queries(monkeypatch) + run(api) + run(api, period=NEXT_PERIOD) + api.heartbeat( + "test", + Event(timestamp=current + timedelta(seconds=10), data={"app": "live"}), + 60, + ) + run(api) + assert spy.call_count == 2 + assert run(api, period=NEXT_PERIOD)[0].duration == timedelta(seconds=10) + assert spy.call_count == 3 + + +def test_cache_tracks_overridden_query_time_bounds(isolated_api, monkeypatch): + api = isolated_api + setup_bucket(api) + query = [ + 'STARTTIME = "2024-01-01T00:00:00+00:00"; ENDTIME = "2024-01-02T00:00:00+00:00"; RETURN = query_bucket("test");' + ] + spy = spy_queries(monkeypatch) + run(api, query=query, period=NEXT_PERIOD) + api.heartbeat( + "test", Event(timestamp=START + timedelta(seconds=10), data={"app": "a"}), 60 + ) + assert run(api, query=query, period=NEXT_PERIOD)[0].duration == timedelta( + seconds=10 + ) + assert spy.call_count == 2 + + +def test_cache_invalidates_bucket_discovery_and_import(isolated_api, monkeypatch): + api = isolated_api + setup_bucket(api) + query = ['RETURN = find_bucket("test");'] + spy = spy_queries(monkeypatch) + assert run(api, query) == "test" + api.delete_bucket("test") + with pytest.raises(QueryException): + run(api, query) + api.import_bucket( + { + "id": "test-import", + "type": "test", + "client": "test", + "hostname": "test", + "created": START.isoformat(), + "events": [], + } + ) + assert run(api, query) == "test-import" + assert spy.call_count == 3 + + +def test_query_overlapping_write_is_not_stored(isolated_api, monkeypatch): + api = isolated_api + setup_bucket(api) + original = api_module.query2.query + calls = [] + + def evaluate(*args): + value = original(*args) + calls.append(True) + if len(calls) == 1: + api.delete_event("test", value[0].id) + return value + + monkeypatch.setattr(api_module.query2, "query", evaluate) + assert len(run(api)) == 1 + assert run(api) == [] + assert len(calls) == 2 + + +def test_query_cache_bounds_expiry_and_active_writers(monkeypatch): + cache = QueryCache(max_entries=2, max_bytes=4096, ttl=10) + tracker = ReadTracker(None) + monkeypatch.setattr("aw_server.query_cache.monotonic", lambda: 0) + for key in ("a", "b", "c"): + cache.store(key, [1], tracker, cache.lookup(key)[2]) + assert not cache.lookup("a")[0] + assert cache.lookup("b")[0] + assert cache.bytes <= 4096 + cache.store("large", "x" * 5000, tracker, cache.lookup("large")[2]) + assert not cache.lookup("large")[0] + with cache.mutation("test"): + assert not cache.lookup("b")[0] + cache.store("during_write", [1], tracker, cache.lookup("x")[2]) + assert not cache.lookup("during_write")[0] + monkeypatch.setattr("aw_server.query_cache.monotonic", lambda: 11) + assert not cache.lookup("b")[0] + assert cache.bytes == 0 + + +def test_heartbeat_uses_inserted_id_and_handles_intervening_mutations(isolated_api): + api = isolated_api + api.create_bucket("test", "test", "test", "test") + first = api.heartbeat("test", Event(timestamp=START, data={"app": "a"}), 60) + assert first.id is not None + second = api.heartbeat( + "test", Event(timestamp=START + timedelta(seconds=1), data={"app": "a"}), 60 + ) + assert second.id == first.id + api.delete_event("test", first.id) + replacement = api.heartbeat( + "test", Event(timestamp=START + timedelta(seconds=2), data={"app": "a"}), 60 + ) + assert api.get_eventcount("test") == 1 + assert replacement.duration == timedelta(0) + api.delete_bucket("test") + api.create_bucket("test", "test", "test", "test") + api.heartbeat( + "test", Event(timestamp=START + timedelta(seconds=3), data={"app": "a"}), 60 + ) + assert api.get_eventcount("test") == 1 + + +def test_failed_heartbeat_does_not_advance_cached_duration(isolated_api, monkeypatch): + api = isolated_api + setup_bucket(api) + api.heartbeat( + "test", Event(timestamp=START + timedelta(seconds=1), data={"app": "a"}), 60 + ) + before = api.last_event["test"].duration + monkeypatch.setattr( + api.db["test"], "replace", Mock(side_effect=RuntimeError("write failed")) + ) + with pytest.raises(RuntimeError, match="write failed"): + api.heartbeat( + "test", Event(timestamp=START + timedelta(seconds=5), data={"app": "a"}), 60 + ) + assert api.last_event["test"].duration == before + + +def test_heartbeat_serializes_direct_api_calls(isolated_api): + api = isolated_api + if api.db.storage_strategy.sid != "memory": + pytest.skip( + "Memory backend isolates API synchronization from DB thread support" + ) + api.create_bucket("test", "test", "test", "test") + with ThreadPoolExecutor(max_workers=4) as pool: + list( + pool.map( + lambda _: api.heartbeat( + "test", Event(timestamp=START, data={"app": "a"}), 60 + ), + range(50), + ) + ) + assert api.get_eventcount("test") == 1 + + +def test_stream_export_matches_eager_export_and_is_lazy(isolated_api, monkeypatch): + api = isolated_api + setup_bucket(api) + api.create_bucket("empty", "test", "test", "test") + expected = {"buckets": api.export_all()} + assert json.loads("".join(api.stream_export())) == expected + assert json.loads("".join(api.stream_export("test"))) == { + "buckets": {"test": expected["buckets"]["test"]} + } + monkeypatch.setattr( + api.db.storage_strategy, + "get_events", + Mock(side_effect=AssertionError("eager read")), + ) + # None of the built-in iterators should fall back to materializing get_events. + assert json.loads("".join(api.stream_export())) == expected + + +def test_closing_export_closes_event_iterator(isolated_api, monkeypatch): + api = isolated_api + setup_bucket(api) + consumed, closed = [], [] + + def events(): + try: + for _ in range(10000): + consumed.append(True) + yield Event(timestamp=START) + finally: + closed.append(True) + + monkeypatch.setattr(api.db["test"], "iter_events", events) + stream = api.stream_export("test") + while not consumed: + next(stream) + stream.close() + assert 0 < len(consumed) < 10000 + assert closed == [True] + + +def test_http_stream_export_and_cache_options(flask_client, app): + response = flask_client.get("/api/0/export") + assert response.status_code == 200 + assert response.is_streamed + assert response.mimetype == "application/json" + assert "buckets" in response.json + assert flask_client.get("/api/0/buckets/missing/export").status_code == 404 + body = {"query": ["RETURN = 1;"], "timeperiods": [PERIOD]} + assert flask_client.post("/api/0/query/?cache=0", json=body).json == [1] + assert ( + flask_client.post("/api/0/query/?cache=invalid", json=body).status_code == 400 + ) + + +def test_disabled_debug_logging_does_not_format_payload(isolated_api, monkeypatch, app): + class Payload(dict): + def __str__(self): + raise AssertionError("payload was eagerly formatted") + + api = isolated_api + api.create_bucket("test", "test", "test", "test") + monkeypatch.setattr(api_module.logger, "level", logging.INFO) + api.heartbeat("test", Event(timestamp=START, data=Payload(app="a")), 60) + monkeypatch.setattr(rest.logger, "level", logging.INFO) + with app.test_request_context("/api/0/buckets/test/events", method="POST"): + monkeypatch.setattr( + rest.request._get_current_object(), + "get_json", + lambda: Payload(timestamp=START, data={"app": "a"}), + ) + monkeypatch.setattr(app.api, "create_events", lambda *args: []) + assert rest.EventsResource().post("test") == ([], 200) diff --git a/tests/test_server.py b/tests/test_server.py index 300bdad..a0ba4db 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -4,9 +4,7 @@ import pytest from aw_core.models import Event -from aw_datastore import Datastore, get_storage_methods -from aw_server.api import ServerAPI from aw_server.exceptions import NotFound @@ -190,28 +188,6 @@ def test_query_valid_timeperiod(flask_client): assert r.json == [1] -@pytest.fixture(params=["memory", "peewee", "sqlite"]) -def isolated_api(request, tmp_path, monkeypatch): - """Exercise the storage backends without accessing a user's database/settings.""" - monkeypatch.setattr("aw_server.settings.get_config_dir", lambda _: str(tmp_path)) - monkeypatch.setattr( - "aw_datastore.storages.peewee.get_data_dir", - lambda _: str(tmp_path), - ) - storage = get_storage_methods()[request.param] - kwargs = ( - {} if request.param == "memory" else {"filepath": str(tmp_path / "test.db")} - ) - db = Datastore(storage, testing=True, **kwargs) - try: - yield ServerAPI(db, testing=True) - finally: - if request.param == "peewee": - db.storage_strategy.db.close() - elif request.param == "sqlite": - db.storage_strategy.conn.close() - - def test_bucket_checks_reuse_datastore_lookup(isolated_api, monkeypatch): api = isolated_api api.create_bucket("test", "test", "test", "test") @@ -221,7 +197,7 @@ def test_bucket_checks_reuse_datastore_lookup(isolated_api, monkeypatch): monkeypatch.setattr(api.db, "buckets", listing) assert api.get_events("test") == [] - assert listing.call_count == 1 + listing.assert_not_called() listing.reset_mock() timestamp = datetime(2024, 1, 1, tzinfo=timezone.utc) From c1e39f8c50cf6a6fff5df432a1249ea94aa80ed3 Mon Sep 17 00:00:00 2001 From: Brayo Date: Thu, 24 Sep 2026 17:13:53 +0300 Subject: [PATCH 3/3] test(perf): profile baseline queries and refresh datastore pin --- poetry.lock | 6 +- pyproject.toml | 2 +- scripts/profile-api-query.py | 105 +++++++++++++++++++++++++++++++++++ 3 files changed, 109 insertions(+), 4 deletions(-) create mode 100644 scripts/profile-api-query.py diff --git a/poetry.lock b/poetry.lock index e0e1ddc..1b5cc58 100644 --- a/poetry.lock +++ b/poetry.lock @@ -79,8 +79,8 @@ tomlkit = "*" [package.source] type = "git" url = "https://github.com/0xbrayo/aw-core.git" -reference = "951a441021278b3ff6000cfa0673798608c57550" -resolved_reference = "951a441021278b3ff6000cfa0673798608c57550" +reference = "a794efb9639b03c243c47c85819f701defbd2591" +resolved_reference = "a794efb9639b03c243c47c85819f701defbd2591" [[package]] name = "black" @@ -1196,4 +1196,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [metadata] lock-version = "2.1" python-versions = "^3.8" -content-hash = "222fa6b67a8dc48a53002ff008c5751217de2c1c24fd9eed28c77b50c1b07cf9" +content-hash = "d9a8d45f09ed67bb9ae094fbd286997820bdfc22c35e66657f7098f59c0f0a9c" diff --git a/pyproject.toml b/pyproject.toml index 3bcfc20..58e2828 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,7 @@ aw-server = "aw_server:main" [tool.poetry.dependencies] python = "^3.8" -aw-core = { git = "https://github.com/0xbrayo/aw-core.git", rev = "951a441021278b3ff6000cfa0673798608c57550" } +aw-core = { git = "https://github.com/0xbrayo/aw-core.git", rev = "a794efb9639b03c243c47c85819f701defbd2591" } aw-client = "^0.5.8" flask = "^2.2" flask-restx = "^1.0.3" diff --git a/scripts/profile-api-query.py b/scripts/profile-api-query.py new file mode 100644 index 0000000..ec6124d --- /dev/null +++ b/scripts/profile-api-query.py @@ -0,0 +1,105 @@ +"""Profile queries against a disposable Peewee database. + +Run from the repository root with its environment activated: + python scripts/profile-api-query.py --baseline-ref master + python scripts/profile-api-query.py --cache + +The optional baseline loads aw_server/api.py from that git revision. Both runs +use the same installed aw-core dependency, isolating server/query-path changes. +No ActivityWatch user data is read or modified. +""" + +import argparse +import cProfile +import logging +import pstats +import statistics +import subprocess +import tempfile +import time +import types +from datetime import datetime, timedelta, timezone +from pathlib import Path +from unittest.mock import patch + +from aw_core.models import Event +from aw_datastore import Datastore +from aw_datastore.storages.peewee import PeeweeStorage +from aw_server.api import ServerAPI + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--baseline-ref") + parser.add_argument("--cache", action="store_true") + parser.add_argument("--events", type=int, default=10000) + parser.add_argument("--queries", type=int, default=10) + args = parser.parse_args() + if args.events < 1 or args.queries < 1: + parser.error("events and queries must be positive") + api_class = ServerAPI + if args.baseline_ref: + module = types.ModuleType("aw_server._profile_baseline") + module.__package__ = "aw_server" + source = subprocess.check_output( + ["git", "show", f"{args.baseline_ref}:aw_server/api.py"] + ) + exec(compile(source, "baseline/aw_server/api.py", "exec"), module.__dict__) + api_class = module.ServerAPI + logging.disable(logging.CRITICAL) + with tempfile.TemporaryDirectory(prefix="aw-query-profile-") as tmp: + with patch( + "aw_datastore.storages.peewee.get_data_dir", return_value=tmp + ), patch("aw_server.settings.get_config_dir", return_value=tmp): + db = Datastore( + PeeweeStorage, testing=True, filepath=str(Path(tmp) / "test.db") + ) + api = api_class(db, testing=True) + try: + bucket = db.create_bucket( + "test", type="test", client="test", hostname="test" + ) + start = datetime(2024, 1, 1, tzinfo=timezone.utc) + bucket.insert( + [ + Event( + timestamp=start + timedelta(seconds=i), + duration=1, + data={"app": "test"}, + ) + for i in range(args.events) + ] + ) + period = ( + start.isoformat() + + "/" + + (start + timedelta(seconds=args.events + 1)).isoformat() + ) + query = ['RETURN = sum_durations(query_bucket("test"));'] + + def execute(): + result = api.query2("profile", query, [period], args.cache) + assert result == [timedelta(seconds=args.events)] + + execute() # warm filesystem/page cache and optional result cache + timings = [] + for _ in range(args.queries): + begin = time.perf_counter() + execute() + timings.append(time.perf_counter() - begin) + print( + f"revision={args.baseline_ref or 'working-tree'}, cache={args.cache}, events={args.events}" + ) + print(f"Median query: {1000 * statistics.median(timings):.3f} ms") + profiler = cProfile.Profile() + profiler.enable() + for _ in range(args.queries): + execute() + profiler.disable() + pstats.Stats(profiler).strip_dirs().sort_stats("cumulative").print_stats(18) + finally: + db.storage_strategy.db.close() + + +if __name__ == "__main__": + main()