test nfstream fix

backend/src/utilities/database.py

@@ -4,6 +4,7 @@ import asyncio
 import base64
 import json
 import logging
+from datetime import datetime, timezone
 from typing import Any, Dict, List, Optional
 
 import asyncpg
@@ -266,6 +267,101 @@ class DatabasePool:
             except Exception:
                 logger.exception("Failed to publish pkt_info to broadcaster")
 
+    async def backfill_flow_metadata(
+        self,
+        *,
+        iface: str,
+        src_ip: str,
+        dst_ip: str,
+        src_port: int,
+        dst_port: int,
+        protocol: int,
+        first_seen_ms: int,
+        last_seen_ms: int,
+        enrichment: Dict[str, Any],
+        window_ms: int,
+    ) -> int:
+        """Update recent packet rows for a flow after enrichment arrives asynchronously."""
+        if self._pool is None:
+            await self.init_pool()
+
+        lower_bound = datetime.fromtimestamp(max(first_seen_ms - window_ms, 0) / 1000.0, tz=timezone.utc)
+        upper_bound = datetime.fromtimestamp(max(last_seen_ms + window_ms, 0) / 1000.0, tz=timezone.utc)
+        dpi_metadata = enrichment.get("dpi_metadata")
+
+        try:
+            async with self._pool.acquire() as conn:
+                rows = await conn.fetch(
+                    """
+                    UPDATE packets
+                    SET
+                        updated_at = NOW(),
+                        app_protocol = COALESCE(packets.app_protocol, $9),
+                        app_master_protocol = COALESCE(packets.app_master_protocol, $10),
+                        app_category = COALESCE(packets.app_category, $11),
+                        app_confidence = COALESCE(packets.app_confidence, $12),
+                        app_hostname = COALESCE(packets.app_hostname, $13),
+                        app_is_encrypted = COALESCE(packets.app_is_encrypted, $14),
+                        dpi_metadata = CASE
+                            WHEN $15::jsonb IS NULL THEN packets.dpi_metadata
+                            WHEN packets.dpi_metadata IS NULL THEN $15::jsonb
+                            ELSE packets.dpi_metadata || $15::jsonb
+                        END
+                    WHERE
+                        ip_proto_raw = $1
+                        AND (capture_iface = $2 OR ingress_if = $2 OR egress_if = $2)
+                        AND timestamp BETWEEN $7 AND $8
+                        AND (
+                            (src_ip = $3::inet AND dst_ip = $4::inet AND src_port = $5 AND dst_port = $6)
+                            OR
+                            (src_ip = $4::inet AND dst_ip = $3::inet AND src_port = $6 AND dst_port = $5)
+                        )
+                        AND (
+                            packets.app_protocol IS NULL
+                            OR packets.app_master_protocol IS NULL
+                            OR packets.app_category IS NULL
+                            OR packets.app_confidence IS NULL
+                            OR packets.app_hostname IS NULL
+                            OR packets.app_is_encrypted IS NULL
+                            OR ($15::jsonb IS NOT NULL)
+                        )
+                    RETURNING *
+                    """,
+                    protocol,
+                    iface,
+                    src_ip,
+                    dst_ip,
+                    src_port,
+                    dst_port,
+                    lower_bound,
+                    upper_bound,
+                    enrichment.get("app_protocol"),
+                    enrichment.get("app_master_protocol"),
+                    enrichment.get("app_category"),
+                    enrichment.get("app_confidence"),
+                    enrichment.get("app_hostname"),
+                    enrichment.get("app_is_encrypted"),
+                    json.dumps(dpi_metadata) if dpi_metadata is not None else None,
+                )
+        except Exception:
+            logger.exception("DB flow metadata backfill failed")
+            return 0
+
+        if not rows:
+            return 0
+
+        updated_count = 0
+        for row in rows:
+            serialized = _serialize_row_for_broadcast(dict(row))
+            updated_count += 1
+            if self.broadcaster:
+                try:
+                    self.broadcaster.sync_publish(serialized)
+                except Exception:
+                    logger.exception("Failed to publish flow-enriched packet row")
+
+        return updated_count
+
     async def fetch_latest(self, limit: int) -> List[PacketDBModel]:
         """Fetch newest packet rows as validated `PacketDBModel` instances."""
         if self._pool is None: