From 96d7a4633f2dfe954ba6413518979773402d889c Mon Sep 17 00:00:00 2001 From: Brayo Date: Thu, 24 Sep 2026 16:48:16 +0300 Subject: [PATCH 1/4] 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 55dfbd2..2d75f38 100644 --- a/aw_server/api.py +++ b/aw_server/api.py @@ -57,8 +57,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 @@ -323,7 +329,7 @@ def delete_bucket(self, bucket_id: str) -> None: finally: self._invalidate_all() self.last_event.pop(bucket_id, None) - logger.debug(f"Deleted bucket '{bucket_id}'") + logger.debug("Deleted bucket '%s'", bucket_id) return None @check_bucket_exists @@ -334,7 +340,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 @@ -348,7 +354,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 = [ @@ -394,7 +400,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 @@ -463,13 +469,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 @@ -495,9 +500,7 @@ 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 ) affected = [event_range(last_event), event_range(merged)] if merged.id is None and self.query_cache: @@ -522,21 +525,18 @@ def heartbeat(self, bucket_id: str, heartbeat: Event, pulsetime: float) -> Event 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, ) try: diff --git a/aw_server/rest.py b/aw_server/rest.py index 1106b37..17af4b5 100644 --- a/aw_server/rest.py +++ b/aw_server/rest.py @@ -297,9 +297,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): @@ -362,7 +362,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: @@ -373,9 +375,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 85ae426..564be5c 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -2,9 +2,15 @@ import io import json 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() @@ -472,3 +478,78 @@ def event(offset, label): events = flask_client.get(f"/api/0/buckets/{bucket}/events").json assert len(events) == 1 assert events[0]["duration"] == 0.5 + + +@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 10c71bc646564c7031b06ebe061ef91c4a979922 Mon Sep 17 00:00:00 2001 From: Brayo Date: Thu, 24 Sep 2026 17:06:44 +0300 Subject: [PATCH 2/4] perf: stream exports and merge heartbeats by event ID --- Makefile | 2 +- aw_server/api.py | 108 ++++++++++++++++++++-------- aw_server/rest.py | 19 ++--- poetry.lock | 22 +++--- pyproject.toml | 2 +- tests/conftest.py | 24 +++++++ tests/test_performance.py | 148 ++++++++++++++++++++++++++++++++++++++ tests/test_server.py | 26 +------ 8 files changed, 279 insertions(+), 72 deletions(-) create mode 100644 tests/test_performance.py diff --git a/Makefile b/Makefile index e66765a..f2c6421 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 tests/test_query_cache.py tests/test_import.py tests/test_host_header.py tests/test_database_concurrency.py + python -m pytest tests/test_server.py tests/test_profile.py tests/test_profile_config.py tests/test_query_cache.py tests/test_import.py tests/test_host_header.py tests/test_database_concurrency.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 2d75f38..5561b36 100644 --- a/aw_server/api.py +++ b/aw_server/api.py @@ -1,3 +1,4 @@ +import copy import functools import json import logging @@ -5,6 +6,7 @@ from datetime import datetime from pathlib import Path from socket import gethostname +from contextlib import closing from typing import ( Any, Dict, @@ -83,6 +85,19 @@ def g(self, *args, **kwargs): return g +def _forget_last_event(f): + """Drop the cached heartbeat state for the bucket before a non-heartbeat + write, so a later heartbeat re-reads the stored last event instead of + merging into (and replacing by ID) a stale one. Use under _serialized.""" + + @functools.wraps(f) + def g(self, bucket_id, *args, **kwargs): + self.last_event.pop(bucket_id, None) + return f(self, bucket_id, *args, **kwargs) + + return g + + class ServerAPI: def __init__(self, db, testing, query_cache: bool = True) -> None: self.db = db @@ -116,15 +131,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]: @@ -156,15 +163,59 @@ def export_bucket_csv( 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() + @_serialized def import_bucket(self, bucket_data: Any): bucket_id = bucket_data["id"] + self.last_event.pop(bucket_id, None) logger.info(f"Importing bucket {bucket_id}") if bucket_id in self.db.buckets(): @@ -256,6 +307,7 @@ def import_all(self, buckets: Dict[str, Any]): raise @_serialized + @_forget_last_event def create_bucket( self, bucket_id: str, @@ -275,7 +327,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() @@ -298,6 +350,7 @@ def create_bucket( @check_bucket_exists @_serialized + @_forget_last_event def update_bucket( self, bucket_id: str, @@ -310,7 +363,7 @@ def update_bucket( try: self.db.update_bucket( bucket_id, - type=event_type, + type_id=event_type, client=client, hostname=hostname, data=data, @@ -322,6 +375,7 @@ def update_bucket( @check_bucket_exists @_serialized + @_forget_last_event def delete_bucket(self, bucket_id: str) -> None: """Delete a bucket""" try: @@ -364,6 +418,7 @@ def get_events( @check_bucket_exists @_serialized + @_forget_last_event 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. @@ -405,6 +460,7 @@ def get_eventcount( @check_bucket_exists @_serialized + @_forget_last_event def delete_event(self, bucket_id: str, event_id) -> bool: """Delete a single event from a bucket""" old = self.db[bucket_id].get_by_id(event_id) @@ -477,22 +533,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: @@ -509,7 +556,7 @@ def heartbeat(self, bucket_id: str, heartbeat: Event, pulsetime: float) -> Event # out-of-order heartbeat), so invalidate its range too. stored = self.db[bucket_id].get(limit=1) affected += [event_range(e) for e in stored] - self.last_event[bucket_id] = merged + replaced = False try: # Replace the event we merged into by id. replace_last() # targets whichever event sorts last by timestamp, which @@ -517,12 +564,17 @@ def heartbeat(self, bucket_id: str, heartbeat: Event, pulsetime: float) -> Event # it gets overwritten with a copy of this one, leaving # same-timestamp duplicates (aw-watcher-afk#61). if merged.id is not None: - self.db[bucket_id].replace(merged.id, merged) + replaced = self.db[bucket_id].replace(merged.id, merged) else: self.db[bucket_id].replace_last(merged) + replaced = True finally: self._invalidate(affected) - return merged + if replaced: + 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)", @@ -540,11 +592,11 @@ def heartbeat(self, bucket_id: str, heartbeat: Event, pulsetime: float) -> Event ) try: - self.db[bucket_id].insert(heartbeat) + inserted = self.db[bucket_id].insert(heartbeat) finally: self._invalidate([event_range(heartbeat)]) - self.last_event[bucket_id] = heartbeat - return heartbeat + self.last_event[bucket_id] = inserted + return inserted def query2(self, name, query, timeperiods, cache: bool = True): """Run a query for each timeperiod. diff --git a/aw_server/rest.py b/aw_server/rest.py index 17af4b5..ada2971 100644 --- a/aw_server/rest.py +++ b/aw_server/rest.py @@ -15,7 +15,8 @@ Response, current_app, jsonify, - make_response, + Response, + stream_with_context, request, stream_with_context, ) @@ -460,9 +461,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 @@ -476,10 +478,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 018616d..eb7939f 100644 --- a/poetry.lock +++ b/poetry.lock @@ -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 = "24.8.0" @@ -1207,4 +1211,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.1" python-versions = "^3.8" -content-hash = "48be42890bb2039993350e7078dfb3c783eff55955609c2c11e52d0368d60915" +content-hash = "7bf53a2e1b6643bf036c6796c4d4965417ef75f9cb48691ea1421952016a28e0" diff --git a/pyproject.toml b/pyproject.toml index 7b7f7ea..4a677c3 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..c68b480 --- /dev/null +++ b/tests/test_performance.py @@ -0,0 +1,148 @@ +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_server import api as api_module +from aw_server import rest + +START = datetime(2024, 1, 1, tzinfo=timezone.utc) + + +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 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(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 + + +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 564be5c..8e87c42 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -7,9 +7,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 @@ -480,28 +478,6 @@ def event(offset, label): assert events[0]["duration"] == 0.5 -@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") @@ -511,7 +487,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 4cf3033e8426ffc995c3aa31b14f9800ccc1eb01 Mon Sep 17 00:00:00 2001 From: Brayo Date: Thu, 24 Sep 2026 17:13:53 +0300 Subject: [PATCH 3/4] 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 eb7939f..6523b6b 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" @@ -1211,4 +1211,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.1" python-versions = "^3.8" -content-hash = "7bf53a2e1b6643bf036c6796c4d4965417ef75f9cb48691ea1421952016a28e0" +content-hash = "5745f3d46d890fdf28c9c66dfec21e863433f77a9c7bfa1b95966cdab3a25b84" diff --git a/pyproject.toml b/pyproject.toml index 4a677c3..4362040 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() From a1b18aaf2121a3df9f397b3818c69d3a94285734 Mon Sep 17 00:00:00 2001 From: Brayo Date: Mon, 19 Oct 2026 14:58:40 +0300 Subject: [PATCH 4/4] fix(rest): drop duplicate flask imports left by rebase --- aw_server/rest.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/aw_server/rest.py b/aw_server/rest.py index ada2971..7952bcc 100644 --- a/aw_server/rest.py +++ b/aw_server/rest.py @@ -15,8 +15,6 @@ Response, current_app, jsonify, - Response, - stream_with_context, request, stream_with_context, )