fix: lint S3/run.py, SQS multi-record parse, BIP-0042 tests and CI job

Author: vaquarkhan

.github/workflows/python-package.yml

@@ -69,12 +69,35 @@ jobs:
 
       - name: Install dependencies
         run: |
-          python -m pip install -e ".[tests,tracking-client,graphviz,tracking-server-s3]"
+          python -m pip install -e ".[tests,tracking-client,graphviz]"
 
       - name: Run tests
         run: |
           python -m pytest tests --ignore=tests/integrations/persisters
 
+  test-tracking-server-s3:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ['3.9', '3.10', '3.11', '3.12']
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          python -m pip install -e ".[tests,tracking-client,tracking-server-s3]"
+
+      - name: Run S3 tracking server tests
+        run: |
+          python -m pytest tests/tracking/test_bip0042_s3_buffering.py -v
+
   validate-examples:
     runs-on: ubuntu-latest
     steps:

burr/tracking/server/run.py

@@ -186,22 +186,15 @@ def get_app_spec():
 
 if app_spec.indexing:
     # Only use polling when not in event-driven mode
-    if not (
-        isinstance(backend, EventDrivenBackendMixin) and backend.is_event_driven()
-    ):
-        update_interval = (
-            backend.update_interval_milliseconds() / 1000 if app_spec.indexing else None
-        )
+    _event_driven = isinstance(backend, EventDrivenBackendMixin) and backend.is_event_driven()
+    if not _event_driven:
         sync_index = repeat_every(
             seconds=backend.update_interval_milliseconds() / 1000,
             wait_first=True,
             logger=logger,
         )(sync_index)
 
 if app_spec.snapshotting:
-    snapshot_interval = (
-        backend.snapshot_interval_milliseconds() / 1000 if app_spec.snapshotting else None
-    )
     save_snapshot = repeat_every(
         seconds=backend.snapshot_interval_milliseconds() / 1000,
         wait_first=True,

burr/tracking/server/s3/backend.py

@@ -33,8 +33,8 @@
 import fastapi
 import pydantic
 from aiobotocore import session
-from pydantic import field_validator
 from fastapi import FastAPI
+from pydantic import field_validator
 from pydantic_settings import BaseSettings
 from tortoise import functions, transactions
 from tortoise.contrib.fastapi import RegisterTortoise
@@ -205,7 +205,40 @@ def timestamp_to_reverse_alphabetical(timestamp: datetime) -> str:
     return inverted_str + "-" + timestamp.isoformat()
 
 
-class SQLiteS3Backend(BackendBase, IndexingBackendMixin, SnapshottingBackendMixin, EventDrivenBackendMixin):
+def _parse_sqs_message_events(
+    body: dict,
+) -> Optional[List[Tuple[str, datetime.datetime]]]:
+    """Parse EventBridge-wrapped or native S3 notification bodies from SQS.
+
+    Returns None if the format is not recognized. Multiple S3 records in one
+    message yield one tuple per record.
+    """
+    if "detail" in body:
+        return [
+            (
+                body["detail"]["object"]["key"],
+                datetime.datetime.fromisoformat(body["time"].replace("Z", "+00:00")),
+            )
+        ]
+    if "Records" in body:
+        out: List[Tuple[str, datetime.datetime]] = []
+        for record in body["Records"]:
+            out.append(
+                (
+                    record["s3"]["object"]["key"],
+                    datetime.datetime.fromisoformat(record["eventTime"].replace("Z", "+00:00")),
+                )
+            )
+        return out
+    return None
+
+
+class SQLiteS3Backend(
+    BackendBase,
+    IndexingBackendMixin,
+    SnapshottingBackendMixin,
+    EventDrivenBackendMixin,
+):
     def __init__(
         self,
         s3_bucket: str,
@@ -790,27 +823,14 @@ async def start_event_consumer(self) -> None:
                         for message in messages:
                             try:
                                 body = json.loads(message["Body"])
-                                s3_key = None
-                                event_time = None
-
-                                # Handle EventBridge wrapped S3 events
-                                if "detail" in body:
-                                    s3_key = body["detail"]["object"]["key"]
-                                    event_time = datetime.datetime.fromisoformat(
-                                        body["time"].replace("Z", "+00:00")
-                                    )
-                                elif "Records" in body:
-                                    record = body["Records"][0]
-                                    s3_key = record["s3"]["object"]["key"]
-                                    event_time = datetime.datetime.fromisoformat(
-                                        record["eventTime"].replace("Z", "+00:00")
-                                    )
-                                else:
-                                    logger.warning(f"Unknown message format: {body}")
+                                events = _parse_sqs_message_events(body)
+                                if events is None:
+                                    logger.warning("Unknown message format: %s", body)
                                     continue
 
-                                if s3_key and s3_key.endswith(".jsonl"):
-                                    await self._handle_s3_event(s3_key, event_time)
+                                for s3_key, event_time in events:
+                                    if s3_key and s3_key.endswith(".jsonl"):
+                                        await self._handle_s3_event(s3_key, event_time)
 
                                 await sqs_client.delete_message(
                                     QueueUrl=self._sqs_queue_url,
@@ -865,7 +885,6 @@ async def indexing_jobs(
 
 if __name__ == "__main__":
     os.environ["BURR_LOAD_SNAPSHOT_ON_START"] = "True"
-    import asyncio
 
     be = SQLiteS3Backend.from_settings(S3Settings())
     # coro = be.snapshot()  # save to s3