db.py

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