db.py

import psycopg2
import os

# Database connection using environment variable
DB_URL = os.getenv("DB_URL")  # Ensure this is set, e.g., postgres://user:password@host:port/dbname

# Initialize connection as None
conn = None
cur = None

try:
    if DB_URL:
        conn = psycopg2.connect(DB_URL)
        cur = conn.cursor()
    else:
        print("Warning: DB_URL environment variable not set. Database functions will be disabled.")
except psycopg2.Error as e:
    print(f"Warning: Database connection failed: {e}. Database functions will be disabled.")

def get_table_status():
    if not cur:
        return {"error": "Database connection not available"}
    try:
        cur.execute("SELECT table_id, status FROM tables")
        return {row[0]: row[1] for row in cur.fetchall()}
    except psycopg2.Error as e:
        print(f"Error fetching table status: {e}")
        return {"error": f"Database query failed: {e}"}

def log_customer_visit(face_id, timestamp, table_id):
    if not cur:
        print("Cannot log visit: Database connection not available")
        return
    try:
        cur.execute("INSERT INTO visits (face_id, timestamp, table_id) VALUES (%s, %s, %s)",
                    (face_id, timestamp, table_id))
        conn.commit()
    except psycopg2.Error as e:
        print(f"Error logging visit: {e}")
        conn.rollback()

def get_alerts():
    if not cur:
        return [{"type": "error", "time": "N/A", "message": "Database connection not available"}]
    try:
        cur.execute("SELECT type, trigger_time FROM alerts WHERE status='active'")
        return [{"type": r[0], "time": r[1].isoformat()} for r in cur.fetchall()]
    except psycopg2.Error as e:
        print(f"Error fetching alerts: {e}")
        return [{"type": "error", "time": "N/A", "message": f"Database query failed: {e}"}]