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)