diff --git a/backend/src/utilities/database.py b/backend/src/utilities/database.py index 3f3abd1..288482b 100644 --- a/backend/src/utilities/database.py +++ b/backend/src/utilities/database.py @@ -25,6 +25,21 @@ def _db_text(value: Any) -> Any: return str(value) +def _serialize_row_for_broadcast(row: Dict[str, Any]) -> Dict[str, Any]: + serialized = dict(row) + + raw_val = serialized.get("raw") + if isinstance(raw_val, (bytes, bytearray)): + serialized["raw_b64"] = base64.b64encode(raw_val).decode("ascii") + serialized.pop("raw", None) + + for key, value in list(serialized.items()): + if hasattr(value, "isoformat"): + serialized[key] = value.isoformat() + + return serialized + + class DatabasePool: """Asyncpg connection pool wrapper used by the packet APIs.""" @@ -87,7 +102,7 @@ class DatabasePool: try: async with self._pool.acquire() as conn: - new_row = await conn.fetchrow( + row = await conn.fetchrow( """ INSERT INTO packets ( packet_uid, @@ -157,7 +172,7 @@ class DatabasePool: dpi_metadata = COALESCE(EXCLUDED.dpi_metadata, packets.dpi_metadata), telemetry_metadata = COALESCE(EXCLUDED.telemetry_metadata, packets.telemetry_metadata), raw = COALESCE(EXCLUDED.raw, packets.raw) - RETURNING id, timestamp + RETURNING * """, pkt_info["packet_uid"], pkt_info.get("ingress_if"), @@ -195,9 +210,9 @@ class DatabasePool: logger.exception("DB upsert failed") return - if new_row: - pkt_info["id"] = new_row["id"] - pkt_info["timestamp"] = new_row["timestamp"].isoformat() + if row: + persisted = _serialize_row_for_broadcast(dict(row)) + pkt_info.update(persisted) if self.broadcaster: try: