add utility for database

backend/src/utilities/database.py

@@ -0,0 +1,79 @@
+
+import logging
+from typing import Dict, List, Optional, Any
+
+import asyncpg
+
+# ---- Logging ----------------------------------------------------------
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("af_packet_sniffer")
+
+class DB:
+    def __init__(self, dsn: str, min_size: int = 1, max_size: int = 5):
+        """
+        Create a DB helper bound to the given DSN.
+        Pool is created lazily on first use unless init_pool() is explicitly called.
+        """
+        self._dsn = dsn
+        self._pool: Optional[asyncpg.pool.Pool] = None
+        self._min_size = min_size
+        self._max_size = max_size
+
+    async def init_pool(self) -> None:
+        """Create connection pool if not already created."""
+        if self._pool is None:
+            self._pool = await asyncpg.create_pool(
+                dsn=self._dsn,
+                min_size=self._min_size,
+                max_size=self._max_size,
+            )
+            logger.info("DB pool initialized")
+
+    async def close_pool(self) -> None:
+        """Gracefully close connection pool."""
+        if self._pool:
+            await self._pool.close()
+            self._pool = None
+            logger.info("DB pool closed")
+
+    async def insert_packet(self, pkt_info: Dict[str, Any]) -> None:
+        """
+        Insert packet metadata. Same fields + same SQL statement as before.
+        """
+        if self._pool is None:
+            await self.init_pool()  # lazy init
+
+        try:
+            async with self._pool.acquire() as conn:
+                await conn.execute(
+                    """
+                    INSERT INTO packets(
+                        iface,
+                        src_mac,
+                        dst_mac,
+                        eth_type,
+                        vlan_id,
+                        src_ip,
+                        dst_ip,
+                        ip_proto,
+                        src_port,
+                        dst_port,
+                        length,
+                        raw
+                    ) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12)
+                    """,
+                    pkt_info["iface"],
+                    pkt_info.get("src_mac"),
+                    pkt_info.get("dst_mac"),
+                    pkt_info.get("eth_type"),
+                    pkt_info.get("vlan_id"),
+                    pkt_info.get("src_ip"),
+                    pkt_info.get("dst_ip"),
+                    pkt_info.get("protocol"),
+                    pkt_info.get("src_port"),
+                    pkt_info.get("dst_port"),
+                    pkt_info["length"],
+                    pkt_info["raw"],
+                )
+        except Exception:
+            logger.exception("DB insert failed")