import os import psycopg2 from psycopg2.extras import RealDictCursor 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(): 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(): """ Return a dict like {table_id: status}. If DB is disabled, return {}. """ 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() 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() def get_alerts(): 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() # 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