Update db.py

db.py

@@ -1,74 +1,47 @@
 import os
 import psycopg2
+from psycopg2.extras import RealDictCursor
 
-DB_URL = os.getenv("DATABASE_URL")  # Set this in Hugging Face Space Secrets
+DB_URL = os.getenv("DB_URL")  # e.g. postgres://user:pass@host:5432/dbname
+ENABLE_DB = os.getenv("ENABLE_DB", "0") == "1"
 
+def db_enabled() -> bool:
+    return ENABLE_DB and bool(DB_URL)
 
 def get_connection():
-    """
-    Returns a Postgres connection if DATABASE_URL is set,
-    otherwise raises RuntimeError (caught in fallback functions).
-    """
-    if not DB_URL:
-        raise RuntimeError("DATABASE_URL is not set. Using fallback (mock mode).")
-    return psycopg2.connect(DB_URL)
-
+    if not db_enabled():
+        raise RuntimeError("DB disabled or DB_URL not set")
+    return psycopg2.connect(DB_URL, connect_timeout=5)
 
 def get_table_status():
     """
-    Example: Fetch current table occupancy status.
-    Returns mock data if DB not available.
+    Return a dict like {table_id: status}. If DB is disabled, return {}.
     """
-    try:
-        conn = get_connection()
-        cur = conn.cursor()
-        cur.execute("SELECT table_id, occupied FROM tables ORDER BY table_id;")
+    if not db_enabled():
+        return {}
+    with get_connection() as conn, conn.cursor() as cur:
+        cur.execute("SELECT table_id, status FROM tables")
         rows = cur.fetchall()
-        cur.close()
-        conn.close()
-        return [{"table_id": r[0], "occupied": r[1]} for r in rows]
-
-    except Exception as e:
-        print("⚠️ DB not available, using mock table status:", e)
-        # fallback mock data
-        return [
-            {"table_id": 1, "occupied": True},
-            {"table_id": 2, "occupied": False},
-            {"table_id": 3, "occupied": True},
-        ]
-
-
-def log_customer_visit(customer_name: str):
-    """
-    Log when a customer is detected/recognized.
-    """
-    try:
-        conn = get_connection()
-        cur = conn.cursor()
-        cur.execute("INSERT INTO visits (customer_name) VALUES (%s);", (customer_name,))
+        return {row[0]: row[1] for row in rows}
+
+def log_customer_visit(face_id: str, timestamp, table_id: int):
+    if not db_enabled():
+        return
+    with get_connection() as conn, conn.cursor() as cur:
+        cur.execute(
+            "INSERT INTO visits (face_id, timestamp, table_id) VALUES (%s, %s, %s)",
+            (face_id, timestamp, table_id),
+        )
         conn.commit()
-        cur.close()
-        conn.close()
-    except Exception as e:
-        print(f"⚠️ Could not log visit to DB ({e}). Customer:", customer_name)
-
 
 def get_alerts():
-    """
-    Fetch alerts from DB if available, otherwise return mock data.
-    """
-    try:
-        conn = get_connection()
-        cur = conn.cursor()
-        cur.execute("SELECT id, message FROM alerts ORDER BY id DESC LIMIT 10;")
+    if not db_enabled():
+        return []
+    with get_connection() as conn, conn.cursor(cursor_factory=RealDictCursor) as cur:
+        cur.execute("SELECT type, trigger_time FROM alerts WHERE status='active'")
         rows = cur.fetchall()
-        cur.close()
-        conn.close()
-        return [{"id": r[0], "message": r[1]} for r in rows]
-
-    except Exception as e:
-        print("⚠️ DB not available, using mock alerts:", e)
-        return [
-            {"id": 1, "message": "Mock Alert: No DB connection"},
-            {"id": 2, "message": "System running in fallback mode"},
-        ]
+        # Convert datetimes to isoformat for JSON friendliness
+        for r in rows:
+            if "trigger_time" in r and r["trigger_time"] is not None:
+                r["trigger_time"] = r["trigger_time"].isoformat()
+        return rows