| 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 |
| |
| 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() |
| |
| return PooledConnectionProxy(conn, _supabase_pool) |
| except Exception as e: |
| print(f"Failed to fetch connection from pool: {e}", flush=True) |
| |
| |
| 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 |
|
|
| |
| def run_db_migrations(): |
| conn = get_supabase_connection() |
| if conn: |
| try: |
| with conn.cursor() as cur: |
| |
| 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!") |
| |
| |
| 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_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: |
| |
| 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() |
|
|
| |
|
|
| 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() |
|
|
|
|
| |
| 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) |
|
|
|
|
