Blame for database.py

import os
import json
import psycopg2
import time
from typing import Any
from psycopg2.extras import RealDictCursor
from dotenv import load_dotenv

load_dotenv(override=True)

from psycopg2.pool import ThreadedConnectionPool

_supabase_pool = None

class PooledConnectionProxy:
    def __init__(self, conn, pool):
        self._conn = conn
        self._pool = pool
        self._released = False

    def __getattr__(self, name):
        return getattr(self._conn, name)

    def __enter__(self):
        return self._conn.__enter__()

    def __exit__(self, exc_type, exc_val, exc_tb):
        return self._conn.__exit__(exc_type, exc_val, exc_tb)

    def close(self):
        if not self._released:
            self._released = True
            try:
                self._pool.putconn(self._conn)
            except Exception as e:
                try:
                    self._conn.close()
                except:
                    pass


def get_supabase_connection():
    global _supabase_pool
    url = os.getenv("DATABASE_URL")
    if not url:
        return None
        
    if _supabase_pool is None:
        try:
            sslmode = "require" if ("supabase" in url and "sslmode" not in url) else None
            # Initialize pool with 2 minimum and 15 maximum concurrent connections
            if sslmode:
                _supabase_pool = ThreadedConnectionPool(2, 15, url, sslmode=sslmode)
            else:
                _supabase_pool = ThreadedConnectionPool(2, 15, url)
            print("Supabase database connection pool initialized successfully.", flush=True)
        except Exception as e:
            print(f"Failed to initialize Supabase connection pool: {e}", flush=True)
            
    if _supabase_pool is not None:
        try:
            conn = _supabase_pool.getconn()
            # Wrap the connection in the proxy to avoid 'close' read-only attribute error on C extension types
            return PooledConnectionProxy(conn, _supabase_pool)
        except Exception as e:
            print(f"Failed to fetch connection from pool: {e}", flush=True)
            
    # Fallback to direct connection if pool fails
    try:
        if "supabase" in url and "sslmode" not in url:
            conn = psycopg2.connect(url, sslmode="require")
        else:
            conn = psycopg2.connect(url)
        return conn
    except Exception as e:
        print(f"Failed to connect to Supabase (direct fallback): {e}")
        return None

# ── Auto-Migration Script ──
def run_db_migrations():
    conn = get_supabase_connection()
    if conn:
        try:
            with conn.cursor() as cur:
                # 1. Migrate user_id column in machine_states
                cur.execute("""
                    SELECT column_name 
                    FROM information_schema.columns 
                    WHERE table_name='machine_states' AND column_name='user_id';
                """)
                col_exists = cur.fetchone()
                
                if not col_exists:
                    print("[Migration] Adding 'user_id' column to 'machine_states'...")
                    cur.execute("ALTER TABLE machine_states ADD COLUMN user_id VARCHAR NOT NULL DEFAULT 'global';")
                    cur.execute("ALTER TABLE machine_states DROP CONSTRAINT IF EXISTS machine_states_pkey CASCADE;")
                    cur.execute("ALTER TABLE machine_states ADD PRIMARY KEY (id, user_id);")
                    conn.commit()
                    print("[Migration] Column 'user_id' and composite primary key configured successfully!")
                else:
                    print("[Migration] 'user_id' column already exists. Running cleanup and constraint check...")
                    cur.execute("UPDATE machine_states SET user_id = 'global' WHERE user_id IS NULL;")
                    cur.execute("ALTER TABLE machine_states ALTER COLUMN user_id SET NOT NULL;")
                    cur.execute("ALTER TABLE machine_states DROP CONSTRAINT IF EXISTS machine_states_pkey CASCADE;")
                    cur.execute("ALTER TABLE machine_states ADD PRIMARY KEY (id, user_id);")
                    conn.commit()
                    print("[Migration] Database cleanup and constraint migration complete!")
                
                # 2. Migrate sent_emails table in Supabase
                cur.execute("""
                    SELECT EXISTS (
                        SELECT 1 
                        FROM information_schema.tables 
                        WHERE table_schema = 'public' 
                          AND table_name = 'sent_emails'
                    );
                """)
                table_exists = cur.fetchone()[0]
                if not table_exists:
                    print("[Migration] Creating 'sent_emails' table in Supabase...")
                    cur.execute("""
                        CREATE TABLE sent_emails (
                            id VARCHAR PRIMARY KEY,
                            recipient VARCHAR,
                            sender VARCHAR,
                            subject VARCHAR,
                            timestamp VARCHAR,
                            body_html TEXT,
                            simulated BOOLEAN,
                            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                        );
                    """)
                    conn.commit()
                    print("[Migration] 'sent_emails' table created successfully!")
        except Exception as e:
            print(f"[Migration] Database update failed: {e}")
        finally:
            conn.close()

# Run migration on startup
run_db_migrations()

def load_db(user_id=None):
    """Load machines data directly from Supabase PostgreSQL."""
    machines = []
    conn = get_supabase_connection()
    if conn:
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                # Query strictly using the active user_id (fallbacks to 'global')
                cur.execute("SELECT * FROM machine_states WHERE user_id = %s;", (user_id or 'global',))
                rows = cur.fetchall()
                for r in rows:
                    upcoming = r["upcoming_tasks"]
                    history = r["history_logs"]
                    alerts = r["safety_alerts"]
                    
                    machines.append({
                        "id": r["id"],
                        "name": r.get("name") or r["id"],
                        "status": r["status"] or "Operational",
                        "health_score": r["health_score"] if r["health_score"] is not None else 90,
                        "last_maintenance": r["last_maintenance"],
                        "upcoming_tasks": upcoming if isinstance(upcoming, list) else json.loads(upcoming or "[]"),
                        "history_logs": history if isinstance(history, list) else json.loads(history or "[]"),
                        "safety_alerts": alerts if isinstance(alerts, list) else json.loads(alerts or "[]")
                    })
        except Exception as e:
            print(f"Failed to load machine states from Supabase: {e}")
        finally:
            conn.close()

    return {"machines": machines}

def get_machine_state(machine_id: str, user_id: str = None) -> dict:
    """Fetch specific machine state directly from Supabase PostgreSQL using machine_id and user_id."""
    conn = get_supabase_connection()
    if conn:
        try:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                cur.execute(
                    "SELECT * FROM machine_states WHERE id = %s AND (user_id = %s OR user_id = 'global');",
                    (machine_id, user_id or 'global')
                )
                r = cur.fetchone()
                if r:
                    upcoming = r["upcoming_tasks"]
                    history = r["history_logs"]
                    alerts = r["safety_alerts"]
                    return {
                        "id": r["id"],
                        "name": r.get("name") or r["id"],
                        "status": r["status"] or "Operational",
                        "health_score": r["health_score"] if r["health_score"] is not None else 90,
                        "last_maintenance": r["last_maintenance"],
                        "upcoming_tasks": upcoming if isinstance(upcoming, list) else json.loads(upcoming or "[]"),
                        "history_logs": history if isinstance(history, list) else json.loads(history or "[]"),
                        "safety_alerts": alerts if isinstance(alerts, list) else json.loads(alerts or "[]")
                    }
        except Exception as e:
            print(f"Failed to fetch machine state for '{machine_id}' from Supabase: {e}")
        finally:
            conn.close()
    return {}

def save_db(data, user_id=None):
    """Save dynamic modifications (logs, tasks, health score, etc.) back to Supabase PostgreSQL."""
    conn = get_supabase_connection()
    if not conn:
        print("Failed to save DB: No Supabase connection.")
        return
    try:
        with conn.cursor() as cur:
            for m in data.get("machines", []):
                cur.execute("""
                    INSERT INTO machine_states (id, user_id, status, health_score, last_maintenance, upcoming_tasks, history_logs, safety_alerts)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                    ON CONFLICT (id, user_id) 
                    DO UPDATE SET 
                        status = EXCLUDED.status,
                        health_score = EXCLUDED.health_score,
                        last_maintenance = EXCLUDED.last_maintenance,
                        upcoming_tasks = EXCLUDED.upcoming_tasks,
                        history_logs = EXCLUDED.history_logs,
                        safety_alerts = EXCLUDED.safety_alerts
                """, (
                    m["id"],
                    user_id or 'global',
                    m.get("status", "Operational"),
                    m.get("health_score", 90),
                    m.get("last_maintenance"),
                    json.dumps(m.get("upcoming_tasks", [])),
                    json.dumps(m.get("history_logs", [])),
                    json.dumps(m.get("safety_alerts", []))
                ))
            conn.commit()
    except Exception as e:
        print(f"Failed to save machines state to Supabase: {e}")
    finally:
        conn.close()

# ── RAG Chunks Persistence (PostgreSQL) ──

def save_chunks_to_db(machine_id: Any = None, chunks: list = None) -> None:
    """
    Upsert all enriched chunks for a document into the rag_chunks table.
    Supports both signatures:
      - save_chunks_to_db(machine_id, chunks)
      - save_chunks_to_db(chunks)
    """
    if chunks is None:
        if isinstance(machine_id, list):
            chunks = machine_id
            machine_id = chunks[0].get("machine_id", "") if chunks else ""
        else:
            chunks = []

    if not chunks:
        print("[DB Chunks] No chunks provided to save_chunks_to_db — skipping.")
        return

    if not machine_id and chunks:
        machine_id = chunks[0].get("machine_id", "")

    conn = get_supabase_connection()
    if not conn:
        print("[DB Chunks] No Supabase connection — skipping cloud save.")
        return
    try:
        with conn.cursor() as cur:
            for c in chunks:
                m_id = c.get("machine_id") or machine_id
                cur.execute("""
                    INSERT INTO rag_chunks (
                        chunk_id, machine_id, filename, page, chunk_index, text,
                        document_type, section, heading, section_number,
                        previous_chunk, next_chunk
                    ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                    ON CONFLICT (chunk_id) DO UPDATE SET
                        text            = EXCLUDED.text,
                        section         = EXCLUDED.section,
                        heading         = EXCLUDED.heading,
                        section_number  = EXCLUDED.section_number,
                        previous_chunk  = EXCLUDED.previous_chunk,
                        next_chunk      = EXCLUDED.next_chunk
                """, (
                    c["chunk_id"],
                    m_id,
                    c["filename"],
                    c["page"],
                    c.get("chunk_index", 0),
                    c["text"],
                    c.get("document_type", ""),
                    c.get("section", ""),
                    c.get("heading", ""),
                    c.get("section_number", ""),
                    c.get("previous_chunk"),
                    c.get("next_chunk"),
                ))
            conn.commit()
            print(f"[DB Chunks] Saved {len(chunks)} chunks to Supabase for machine '{machine_id}'.", flush=True)
    except Exception as e:
        print(f"[DB Chunks] Failed to save chunks: {e}", flush=True)
    finally:
        conn.close()


def load_chunks_from_db(machine_id: str) -> list:
    """
    Load all chunks for a machine from the rag_chunks table.
    Returns chunks ordered by filename, page, chunk_index (sequential order).
    """
    conn = get_supabase_connection()
    if not conn:
        print("[DB Chunks] No Supabase connection — cannot load chunks from cloud.")
        return []
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute("""
                SELECT * FROM rag_chunks
                WHERE machine_id = %s
                ORDER BY filename, page, chunk_index
            """, (machine_id,))
            rows = cur.fetchall()
            chunks = []
            for r in rows:
                chunks.append({
                    "chunk_id":       r["chunk_id"],
                    "machine_id":     r["machine_id"],
                    "filename":       r["filename"],
                    "page":           r["page"],
                    "chunk_index":    r["chunk_index"],
                    "text":           r["text"],
                    "document_type":  r["document_type"] or "",
                    "section":        r["section"] or "",
                    "heading":        r["heading"] or "",
                    "section_number": r["section_number"] or "",
                    "previous_chunk": r["previous_chunk"],
                    "next_chunk":     r["next_chunk"],
                    "score":          0.0
                })
            print(f"[DB Chunks] Loaded {len(chunks)} chunks from Supabase for machine '{machine_id}'.", flush=True)
            return chunks
    except Exception as e:
        print(f"[DB Chunks] Failed to load chunks: {e}", flush=True)
        return []
    finally:
        conn.close()


def delete_chunks_from_db(machine_id: str, filename: str) -> None:
    """
    Delete all chunks for a specific document from the rag_chunks table.
    Called when a user removes a document.
    """
    conn = get_supabase_connection()
    if not conn:
        return
    try:
        with conn.cursor() as cur:
            cur.execute(
                "DELETE FROM rag_chunks WHERE machine_id = %s AND filename = %s",
                (machine_id, filename)
            )
            deleted = cur.rowcount
            conn.commit()
            print(f"[DB Chunks] Deleted {deleted} chunks from Supabase for '{filename}' (machine '{machine_id}').", flush=True)
    except Exception as e:
        print(f"[DB Chunks] Failed to delete chunks for {filename}: {e}", flush=True)
    finally:
        conn.close()


def delete_all_chunks_from_db(machine_id: str) -> None:
    """
    Delete ALL chunks for a machine from the rag_chunks table.
    Called when the entire machine cache is cleared.
    """
    conn = get_supabase_connection()
    if not conn:
        return
    try:
        with conn.cursor() as cur:
            cur.execute("DELETE FROM rag_chunks WHERE machine_id = %s", (machine_id,))
            deleted = cur.rowcount
            conn.commit()
            print(f"[DB Chunks] Deleted all {deleted} chunks from Supabase for machine '{machine_id}'.", flush=True)
    except Exception as e:
        print(f"[DB Chunks] Failed to delete all chunks: {e}", flush=True)
    finally:
        conn.close()


# ── Redis Connection Setup ──
import redis

redis_client = None
REDIS_URL = os.getenv("REDIS_URL")
if REDIS_URL:
    try:
        redis_client = redis.from_url(REDIS_URL, decode_responses=True)
        print(f"Supabase DB Manager: Connected to Redis successfully at {REDIS_URL}", flush=True)
    except Exception as e:
        print(f"Supabase DB Manager: Failed to initialize Redis connection: {e}", flush=True)