/db_export: add ?since

masterbase/app.py

@@ -1,5 +1,6 @@
 """Litestar Application for serving and ingesting data."""
 
+import json
 import logging
 import os
 from datetime import datetime, timezone
@@ -8,7 +9,7 @@
 import requests
 import uvicorn
 from litestar import Litestar, MediaType, Request, WebSocket, get, post
-from litestar.exceptions import HTTPException, PermissionDeniedException
+from litestar.exceptions import HTTPException, PermissionDeniedException, ValidationException
 from litestar.handlers import WebsocketListener
 from litestar.response import Redirect, Response, Stream
 from litestar.status_codes import HTTP_500_INTERNAL_SERVER_ERROR
@@ -172,12 +173,12 @@ async def demodata(request: Request, api_key: str, session_id: str) -> Stream:
 
 
 @get("/db_export", guards=[valid_key_guard, analyst_guard], sync_to_thread=False)
-def db_export(request: Request, api_key: str, table: ExportTable) -> Stream:
+def db_export(request: Request, api_key: str, table: ExportTable, since: datetime | None = None) -> Stream:
     """Return a database export of the requested `table`."""
     engine = request.app.state.engine
     filename = f"{table.value}-{datetime.now()}.csv"
     return Stream(
-        lambda: db_export_chunks(engine, table.value),
+        lambda: db_export_chunks(engine, table.value, since),
         headers={
             "Content-Type": "text/csv",
             "Content-Disposition": f"attachment; filename={filename}",
@@ -378,9 +379,15 @@ def provision_handler(request: Request) -> str:
         """
 
 
-def plain_text_exception_handler(_: Request, exception: Exception) -> Response:
+def readable_exception_handler(_: Request, exception: Exception) -> Response:
     """Handle exceptions subclassed from HTTPException."""
     status_code = getattr(exception, "status_code", HTTP_500_INTERNAL_SERVER_ERROR)
+    if isinstance(exception, ValidationException):
+        return Response(
+            content=json.dumps({"detail": "Validation error in request parameters", "errors": exception.extra}),
+            status_code=status_code,
+            media_type=MediaType.JSON,
+        )
     if isinstance(exception, HTTPException):
         content = exception.detail
     else:
@@ -409,7 +416,7 @@ def plain_text_exception_handler(_: Request, exception: Exception) -> Response:
         report_player,
         db_export,
     ],
-    exception_handlers={Exception: plain_text_exception_handler},
+    exception_handlers={Exception: readable_exception_handler},
     on_shutdown=shutdown_registers,
     opt={"DEVELOPMENT": bool(os.getenv("DEVELOPMENT"))},
 )

masterbase/lib.py

@@ -56,7 +56,7 @@ def make_minio_client(is_secure: bool = False) -> Minio:
     return Minio(f"{host}:{port}", access_key=access_key, secret_key=secret_key, secure=is_secure)
 
 
-def db_export_chunks(engine: Engine, table: str) -> Generator[bytes, None, None]:
+def db_export_chunks(engine: Engine, table: str, since: datetime | None = None) -> Generator[bytes, None, None]:
     """Export the given table as an iterable of csv chunks."""
 
     class Shunt:
@@ -72,7 +72,17 @@ def worker():
         try:
             with engine.connect() as txn:
                 cursor = txn.connection.dbapi_connection.cursor()
-                cursor.copy_expert(f"COPY {table} TO STDOUT DELIMITER ',' CSV HEADER", shunt)
+                if since:
+                    # make a best-effort attempt to match the postgres timestamp format
+                    # this only works up to hourly precision
+                    tzoffset = int(since.utcoffset().total_seconds() / 3600)
+                    sign = "+" if tzoffset >= 0 else "-"
+                    stamp = since.strftime("%Y-%m-%d %H:%M:%S")
+                    stamp += f"{sign}{tzoffset}"
+                    query = f"(SELECT * FROM {table} WHERE created_at >= '{stamp}')"
+                else:
+                    query = table
+                cursor.copy_expert(f"COPY {query} TO STDOUT DELIMITER ',' CSV HEADER", shunt)
                 queue.put(b"")
         except Exception as err:
             queue.put(err)

tests/test_api.py

@@ -2,7 +2,9 @@
 
 import csv
 import io
+import re
 import time
+from datetime import datetime, timedelta, timezone
 from typing import Iterator
 
 import pytest
@@ -55,9 +57,9 @@ def test_client(steam_id: str, api_key: str) -> Iterator[TestClient[Litestar]]:
             conn.execute(sa.text(sql))
             sql = "TRUNCATE TABLE demo_sessions CASCADE;"
             conn.execute(sa.text(sql))
-            sql = "DELETE FROM api_keys WHERE api_key = :api_key;"
+            sql = "TRUNCATE TABLE api_keys CASCADE;"
             conn.execute(sa.text(sql), {"api_key": api_key})
-            sql = "DELETE FROM analyst_steam_ids WHERE steam_id = :steam_id;"
+            sql = "TRUNCATE TABLE analyst_steam_ids CASCADE;"
             conn.execute(sa.text(sql), {"steam_id": steam_id})
             conn.commit()
 
@@ -148,29 +150,39 @@ def test_demo_streaming_no_late(test_client: TestClient[Litestar], api_key: str)
 
 def test_db_exports(test_client: TestClient[Litestar], api_key: str) -> None:
     """Test on-demand exports from the database."""
+
+    def _parse_reports(body):
+        records = csv.DictReader(io.TextIOWrapper(io.BytesIO(body), encoding="utf8"))
+        fields = records.fieldnames
+        assert fields is not None
+        assert set(fields) == {"session_id", "target_steam_id", "reason", "created_at"}
+        return tuple(sorted(records, key=lambda r: r["created_at"]))
+
     session_id = str(_open_mock_session(test_client, api_key).json()["session_id"])
     # Insert mock reports
     expected = []
     for i in range(10):
         reason = "cheater" if i % 2 == 0 else "bot"
         target_steam_id = f"{i:020d}"
         record = {"session_id": session_id, "target_steam_id": target_steam_id, "reason": reason}
+        if i == 4:
+            time.sleep(1.0)  # postgres timestamp comparisons are second-precison
         add_report(test_client.app.state.engine, **record)
-        expected.append((session_id, target_steam_id, reason))
+        expected.append(record)
 
     test_client.get("/close_session", params={"api_key": api_key})
     response = test_client.get("/db_export", params={"api_key": api_key, "table": "reports"})
-    response_records = csv.DictReader(io.TextIOWrapper(io.BytesIO(response.content)))
-    assert response_records.fieldnames is not None and set(response_records.fieldnames).issuperset(
-        {"session_id", "target_steam_id", "reason"}
-    )
-    returned = sorted(
-        (
-            (record["session_id"], record["target_steam_id"], record["reason"])
-            for record in sorted(response_records, key=lambda record: record["created_at"])
-        )
+    returned_full = _parse_reports(response.content)
+    assert tuple(expected) == tuple({k: v for k, v in r.items() if k != "created_at"} for r in returned_full)
+    since = returned_full[4]["created_at"]
+
+    tzone = timezone(timedelta(hours=int(since[-3:])))
+    stamp = datetime.strptime(since[:-3], "%Y-%m-%d %H:%M:%S.%f").replace(tzinfo=tzone)
+    response = test_client.get(
+        "/db_export", params={"api_key": api_key, "table": "reports", "since": stamp.isoformat()}
     )
-    assert tuple(expected) == tuple(returned)
+    returned_since = _parse_reports(response.content)
+    assert returned_full[4:] == returned_since
 
 
 def test_upsert_report_reason(test_client: TestClient[Litestar], api_key: str) -> None: