diff --git a/Makefile b/Makefile index 2a58b51..af15665 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_database_concurrency.py typecheck: python -m mypy aw_server tests --ignore-missing-imports diff --git a/aw_server/server.py b/aw_server/server.py index 073406f..7ca0308 100644 --- a/aw_server/server.py +++ b/aw_server/server.py @@ -6,6 +6,7 @@ import aw_datastore import flask.json.provider from aw_datastore import Datastore +from aw_datastore.storages.peewee import PeeweeStorage from flask import ( Blueprint, Flask, @@ -58,6 +59,19 @@ def __init__( if storage_method is None: storage_method = aw_datastore.get_storage_methods()["memory"] db = Datastore(storage_method, testing=testing) + if isinstance(db.storage_strategy, PeeweeStorage): + database = db.storage_strategy.db + # aw-core enables WAL for concurrent readers and writers. Close + # its initialization connection before serving request threads. + database.close() + + @self.teardown_request + def close_database(error): + # Peewee opens connections lazily and keeps them thread-local. + # Release cursors/locks even if handling a request failed. + if not database.is_closed(): + database.close() + self.api = ServerAPI(db=db, testing=testing) self.register_blueprint(root) diff --git a/tests/test_database_concurrency.py b/tests/test_database_concurrency.py new file mode 100644 index 0000000..dbc7d4c --- /dev/null +++ b/tests/test_database_concurrency.py @@ -0,0 +1,94 @@ +"""Exercise real SQLite locks across separate request-thread connections.""" + +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +import pytest +from aw_datastore.storages.peewee import PeeweeStorage + +from aw_server.server import AWFlask + + +@pytest.fixture() +def sqlite_app(tmp_path): + app = AWFlask( + "127.0.0.1", + testing=True, + storage_method=partial(PeeweeStorage, filepath=str(tmp_path / "events.db")), + cors_origins=[], + ) + database = app.api.db.storage_strategy.db + with app.test_client() as client: + response = client.post( + "/api/0/buckets/test", + json={"client": "test", "type": "test", "hostname": "test"}, + ) + assert response.status_code == 200 + try: + yield app, database + finally: + database.close() + + +def request_in_thread(app, database, method, path, **kwargs): + # Keep failures fast on rollback-journal SQLite, rather than waiting for + # its default busy timeout while the main thread holds a transaction. + database.connect(reuse_if_open=True) + database.execute_sql("PRAGMA busy_timeout = 100") + try: + response = getattr(app.test_client(), method)(path, **kwargs) + return response.status_code, response.get_json(), database.is_closed() + finally: + database.close() + + +def test_read_during_write(sqlite_app): + app, database = sqlite_app + with ThreadPoolExecutor(max_workers=1) as executor: + with database.atomic("EXCLUSIVE"): + database.execute_sql("UPDATE bucketmodel SET name = 'pending'") + future = executor.submit( + request_in_thread, app, database, "get", "/api/0/buckets/test/events" + ) + status, events, closed = future.result(timeout=5) + assert status == 200 + assert events == [] + assert closed + + +def test_write_during_read(sqlite_app): + app, database = sqlite_app + with ThreadPoolExecutor(max_workers=1) as executor: + with database.atomic(): + # Retain a read snapshot until the concurrent writer has committed. + database.execute_sql("SELECT * FROM bucketmodel").fetchall() + future = executor.submit( + request_in_thread, + app, + database, + "post", + "/api/0/buckets/test/events", + json=[{"timestamp": "2024-01-01T00:00:00Z", "duration": 1, "data": {}}], + ) + status, events, closed = future.result(timeout=5) + assert status == 200 + assert len(events) == 1 + assert closed + response = app.test_client().get("/api/0/buckets/test/events") + assert len(response.get_json()) == 1 + + +def test_connection_closed_after_error(sqlite_app, monkeypatch): + app, database = sqlite_app + + def database_error(*args, **kwargs): + database.execute_sql("SELECT * FROM bucketmodel").fetchall() + raise RuntimeError("request failed") + + monkeypatch.setattr(app.api, "get_events", database_error) + with ThreadPoolExecutor(max_workers=1) as executor: + status, _, closed = executor.submit( + request_in_thread, app, database, "get", "/api/0/buckets/test/events" + ).result(timeout=5) + assert status == 500 + assert closed