diff --git a/Makefile b/Makefile index e66765a3..f2c64211 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 55dfbd22..5561b36c 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, @@ -57,8 +59,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 @@ -77,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 @@ -110,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]: @@ -150,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(): @@ -250,6 +307,7 @@ def import_all(self, buckets: Dict[str, Any]): raise @_serialized + @_forget_last_event def create_bucket( self, bucket_id: str, @@ -269,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() @@ -292,6 +350,7 @@ def create_bucket( @check_bucket_exists @_serialized + @_forget_last_event def update_bucket( self, bucket_id: str, @@ -304,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, @@ -316,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: @@ -323,7 +383,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 +394,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 +408,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 = [ @@ -358,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. @@ -394,11 +455,12 @@ 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 @_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) @@ -463,31 +525,21 @@ 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 - # 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: @@ -495,9 +547,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: @@ -506,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 @@ -514,37 +564,39 @@ 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: {})".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: - 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 1106b37e..7952bcc8 100644 --- a/aw_server/rest.py +++ b/aw_server/rest.py @@ -15,7 +15,6 @@ Response, current_app, jsonify, - make_response, request, stream_with_context, ) @@ -297,9 +296,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 +361,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 +374,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 @@ -458,9 +459,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 @@ -474,10 +476,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 018616d7..6523b6be 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 = "a794efb9639b03c243c47c85819f701defbd2591" +resolved_reference = "a794efb9639b03c243c47c85819f701defbd2591" + [[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 = "5745f3d46d890fdf28c9c66dfec21e863433f77a9c7bfa1b95966cdab3a25b84" diff --git a/pyproject.toml b/pyproject.toml index 7b7f7ea8..4362040f 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 = "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 00000000..ec6124d2 --- /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() diff --git a/tests/conftest.py b/tests/conftest.py index 7e324047..2c826731 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 00000000..c68b4802 --- /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 85ae4269..8e87c426 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -2,9 +2,13 @@ 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_server.exceptions import NotFound @pytest.fixture() @@ -472,3 +476,56 @@ 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 + + +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") == [] + listing.assert_not_called() + 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")