fix: add schema migration to handle missing columns in existing tables

database.py

@@ -1,48 +1,34 @@
-"""
-Database Module
-Handles SQL validation, database operations, and LLM-to-SQL conversion
-"""
-
 import os
-from datetime import datetime
-from sqlalchemy import create_engine, text
-from dotenv import load_dotenv
-from huggingface_hub import InferenceClient
 import traceback
-
-
-"""
-Database Module
-Handles SQL validation, database operations, and LLM-to-SQL conversion
-"""
-
-import os
 from datetime import datetime
-from sqlalchemy import create_engine, text
+
 from dotenv import load_dotenv
 from huggingface_hub import InferenceClient
-import traceback
-
-
-# ==================== INITIALIZATION ====================
+from sqlalchemy import create_engine, text
 
 load_dotenv()
 
+# ================= DATABASE ================= #
+
 DATABASE_URL = os.getenv("DATABASE_URL")
 engine = None
 db_available = False
 
 if DATABASE_URL:
     try:
-        engine = create_engine(DATABASE_URL)
+        engine = create_engine(
+            DATABASE_URL,
+            pool_pre_ping=True
+        )
         db_available = True
         print("✅ Database engine initialized")
     except Exception as e:
-        print(f"⚠️  Database initialization error: {e}")
+        print(f"⚠️ Database initialization error: {e}")
         db_available = False
 else:
-    print("⚠️  DATABASE_URL not set. Database features disabled.")
-    print("   Set DATABASE_URL in Hugging Face Spaces secrets")
+    print("⚠️ DATABASE_URL not set. Database features disabled.")
+
+# ================= HUGGINGFACE ================= #
 
 HF_TOKEN = os.getenv("HF_TOKEN")
 client = None
@@ -55,248 +41,207 @@ if HF_TOKEN:
         )
         print("✅ HuggingFace client initialized")
     except Exception as e:
-        print(f"⚠️  HuggingFace client error: {e}")
+        print(f"⚠️ HuggingFace client error: {e}")
 else:
-    print("⚠️  HF_TOKEN not set. NLP-to-SQL features disabled.")
-    print("   Set HF_TOKEN in Hugging Face Spaces secrets")
-
-# Query history for analytics
-query_history = []
-
-
-# ==================== SQL VALIDATION ====================
-
-ALLOWED_KEYWORDS = [
-    "SELECT",
-    "FROM",
-    "WHERE",
-    "COUNT",
-    "GROUP BY",
-    "ORDER BY",
-    "LIMIT",
-    "AVG",
-    "SUM",
-    "MIN",
-    "MAX",
-    "DISTINCT",
-    "AND",
-    "OR",
-    "IN",
-    "BETWEEN",
-    "LIKE"
-]
+    print("⚠️ HF_TOKEN not set. NLP-to-SQL features disabled.")
 
-BLOCKED_KEYWORDS = [
-    "DROP",
-    "DELETE",
-    "UPDATE",
-    "INSERT",
-    "ALTER",
-    "TRUNCATE",
-    "CREATE",
-    "EXEC",
-    "EXECUTE",
-    "--",
-    "/*"
-]
+# ================= INIT DB ================= #
 
+def init_db():
+    
+    if not db_available or engine is None:
+        print("⚠️ Skipping database initialization - DATABASE_URL not configured")
+        return
 
-def validate_sql(sql):
-    """
-    Validate SQL query for safety
-    Only allow SELECT queries, block dangerous operations
-    """
-    upper_sql = sql.upper()
+    with engine.connect() as conn:
 
-    # Block dangerous keywords
-    for word in BLOCKED_KEYWORDS:
-        if word in upper_sql:
-            return False, f"Dangerous keyword '{word}' detected"
+        # Create table if it doesn't exist
+        conn.execute(text("""
+        CREATE TABLE IF NOT EXISTS vehicle_logs (
 
-    # Must be SELECT query
-    if not upper_sql.strip().startswith("SELECT"):
-        return False, "Only SELECT queries allowed"
+            id BIGSERIAL PRIMARY KEY,
 
-    return True, "Valid"
+            timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
 
+            plate TEXT,
 
-def clean_sql(sql):
-    """Clean SQL output from LLM"""
-    sql = sql.replace("```sql", "")
-    sql = sql.replace("```", "")
-    sql = sql.strip()
+            state TEXT,
 
-    if sql.endswith(";"):
-        return sql
-    
-    return sql + ";"
+            vehicle_type TEXT,
 
+            vehicle_conf FLOAT,
 
-# ==================== DATABASE INITIALIZATION ====================
+            date DATE,
 
-def init_db():
-    """Create vehicle_logs table with improved schema"""
-    if not db_available or engine is None:
-        print("⚠️  Skipping database initialization - DATABASE_URL not configured")
-        return False
-    
-    try:
-        with engine.connect() as conn:
+            hour INTEGER,
+
+            day TEXT
+        )
+        """))
+
+        # Add missing columns if they don't exist
+        # Check and add 'state' column
+        try:
             conn.execute(text("""
-            CREATE TABLE IF NOT EXISTS vehicle_logs (
-                id BIGSERIAL PRIMARY KEY,
-                
-                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                
-                plate TEXT,
-                state TEXT,
-                
-                vehicle_type TEXT,
-                vehicle_conf FLOAT,
-                
-                camera_id TEXT,
-                location TEXT,
-                
-                image_url TEXT,
-                
-                date DATE,
-                hour INTEGER,
-                day TEXT
-            )
+            ALTER TABLE vehicle_logs
+            ADD COLUMN state TEXT;
             """))
-            conn.commit()
-        
-        # Create indexes
-        create_indexes()
-        print("✅ Database initialized successfully")
-        return True
-    
-    except Exception as e:
-        print(f"❌ Database initialization error: {e}")
-        traceback.print_exc()
-        return False
+            print("✅ Added 'state' column")
+        except:
+            pass  # Column already exists
 
+        # Check and add 'vehicle_type' column
+        try:
+            conn.execute(text("""
+            ALTER TABLE vehicle_logs
+            ADD COLUMN vehicle_type TEXT;
+            """))
+            print("✅ Added 'vehicle_type' column")
+        except:
+            pass  # Column already exists
 
-def create_indexes():
-    """Create database indexes for faster queries"""
-    if not db_available or engine is None:
-        return
-    
-    try:
-        with engine.connect() as conn:
-            # Check if indexes exist before creating
-            indexes = [
-                "idx_plate",
-                "idx_state",
-                "idx_vehicle_type",
-                "idx_timestamp",
-                "idx_date"
-            ]
-            
-            for idx_name in indexes:
-                try:
-                    if idx_name == "idx_plate":
-                        conn.execute(text("CREATE INDEX IF NOT EXISTS idx_plate ON vehicle_logs(plate);"))
-                    elif idx_name == "idx_state":
-                        conn.execute(text("CREATE INDEX IF NOT EXISTS idx_state ON vehicle_logs(state);"))
-                    elif idx_name == "idx_vehicle_type":
-                        conn.execute(text("CREATE INDEX IF NOT EXISTS idx_vehicle_type ON vehicle_logs(vehicle_type);"))
-                    elif idx_name == "idx_timestamp":
-                        conn.execute(text("CREATE INDEX IF NOT EXISTS idx_timestamp ON vehicle_logs(timestamp);"))
-                    elif idx_name == "idx_date":
-                        conn.execute(text("CREATE INDEX IF NOT EXISTS idx_date ON vehicle_logs(date);"))
-                except:
-                    pass
-            
-            conn.commit()
-        print("✅ Indexes created successfully")
-    
-    except Exception as e:
-        print(f"⚠️  Index creation warning: {e}")
+        # Check and add 'vehicle_conf' column
+        try:
+            conn.execute(text("""
+            ALTER TABLE vehicle_logs
+            ADD COLUMN vehicle_conf FLOAT;
+            """))
+            print("✅ Added 'vehicle_conf' column")
+        except:
+            pass  # Column already exists
 
+        # Check and add 'date' column
+        try:
+            conn.execute(text("""
+            ALTER TABLE vehicle_logs
+            ADD COLUMN date DATE;
+            """))
+            print("✅ Added 'date' column")
+        except:
+            pass  # Column already exists
 
-# ==================== DATABASE INSERT ====================
+        # Check and add 'hour' column
+        try:
+            conn.execute(text("""
+            ALTER TABLE vehicle_logs
+            ADD COLUMN hour INTEGER;
+            """))
+            print("✅ Added 'hour' column")
+        except:
+            pass  # Column already exists
 
-def save_detection(plate, state, vehicle_type, vehicle_conf, date_str, time_str):
-    """Save detection to database"""
-    if not db_available or engine is None:
-        return False
-    
-    try:
-        # Extract hour and day
-        dt = datetime.strptime(f"{date_str} {time_str}", "%Y-%m-%d %H:%M:%S")
-        hour = dt.hour
-        day = dt.strftime("%A")
-        
-        with engine.connect() as conn:
+        # Check and add 'day' column
+        try:
             conn.execute(text("""
-                INSERT INTO vehicle_logs
-                (plate, state, vehicle_type, vehicle_conf, date, hour, day)
-                VALUES (:plate, :state, :vehicle_type, :vehicle_conf, :date, :hour, :day)
-            """), {
-                "plate": plate,
-                "state": state,
-                "vehicle_type": vehicle_type,
-                "vehicle_conf": float(vehicle_conf),
-                "date": date_str,
-                "hour": hour,
-                "day": day
-            })
-            conn.commit()
-        return True
-    
-    except Exception as e:
-        print(f"❌ Database insert error: {e}")
-        traceback.print_exc()
-        return False
+            ALTER TABLE vehicle_logs
+            ADD COLUMN day TEXT;
+            """))
+            print("✅ Added 'day' column")
+        except:
+            pass  # Column already exists
 
+        conn.commit()
 
-# ==================== HEALTH CHECK ====================
+    print("✅ Database Initialized")
 
-def health_check():
-    """Check database connectivity"""
-    if not db_available or engine is None:
-        return False, "⚠️  Database not configured. Set DATABASE_URL secret."
+
+init_db()
+
+# ================= SAVE ================= #
+
+def save_detection(
+    plate,
+    state,
+    vehicle_type,
+    vehicle_conf,
+    date_str,
+    time_str
+):
     
-    try:
-        with engine.connect() as conn:
-            conn.execute(text("SELECT 1"))
-        return True, "✅ Database Connected"
-    except Exception as e:
-        return False, f"❌ Database Error: {e}"
+    if not db_available or engine is None:
+        print("⚠️ Database not available for saving detection")
+        return
+
+    dt = datetime.strptime(
+        f"{date_str} {time_str}",
+        "%Y-%m-%d %H:%M:%S"
+    )
+
+    hour = dt.hour
+
+    day = dt.strftime("%A")
+
+    with engine.connect() as conn:
 
+        conn.execute(text("""
+        INSERT INTO vehicle_logs
+        (
+            plate,
+            state,
+            vehicle_type,
+            vehicle_conf,
+            date,
+            hour,
+            day
+        )
+        VALUES
+        (
+            :plate,
+            :state,
+            :vehicle_type,
+            :vehicle_conf,
+            :date,
+            :hour,
+            :day
+        )
+        """), {
+
+            "plate": plate,
+            "state": state,
+            "vehicle_type": vehicle_type,
+            "vehicle_conf": float(vehicle_conf),
+            "date": date_str,
+            "hour": hour,
+            "day": day
+        })
+
+        conn.commit()
 
-# ==================== LLM SQL GENERATION ====================
+
+# ================= SQL SAFETY ================= #
+
+BLOCKED = [
+    "DROP",
+    "DELETE",
+    "UPDATE",
+    "INSERT",
+    "ALTER",
+    "TRUNCATE"
+]
+
+def validate_sql(sql):
+
+    upper = sql.upper()
+
+    for word in BLOCKED:
+
+        if word in upper:
+            return False
+
+    return upper.startswith("SELECT")
+
+
+# ================= NLP TO SQL ================= #
 
 def ask_llm(user_query):
-    """Convert natural language to SQL using SQLCoder"""
-    
+
     if client is None:
-        return None, "LLM not configured. Set HF_TOKEN in secrets."
-    
-    schema = """
-    Table: vehicle_logs
-    
-    Columns:
-    - id (BIGSERIAL PRIMARY KEY)
-    - timestamp (TIMESTAMP) - When vehicle was detected
-    - plate (TEXT) - License plate number
-    - state (TEXT) - State code (TN, KA, KL, etc.)
-    - vehicle_type (TEXT) - Type of vehicle (car, truck, bike, etc.)
-    - vehicle_conf (FLOAT) - Detection confidence (0-1)
-    - camera_id (TEXT) - Camera identifier
-    - location (TEXT) - Detection location
-    - image_url (TEXT) - URL to vehicle image
-    - date (DATE) - Date of detection
-    - hour (INTEGER) - Hour of day (0-23)
-    - day (TEXT) - Day of week (Monday, Tuesday, etc.)
-    
-    Note: Indian vehicle plates format - XX DD YYY NNNN (e.g., TN 09 AB 1234)
-    """
+        return "SELECT 1;"  # Dummy query if no client
 
-    prompt = f"""You are an expert PostgreSQL SQL generator for vehicle surveillance database.
+    prompt = f"""
+You are an expert PostgreSQL SQL generator.
 
-Database schema:
+ONLY use this table:
 
 vehicle_logs(
     id,
@@ -305,65 +250,76 @@ vehicle_logs(
     state,
     vehicle_type,
     vehicle_conf,
-    camera_id,
-    location,
-    image_url,
     date,
     hour,
     day
 )
 
-Strict Rules:
-- ONLY generate PostgreSQL SELECT queries
-- NEVER generate DELETE/UPDATE/INSERT/DROP/ALTER queries
-- Always use LIMIT 50 unless specifically counting totals
-- Use COUNT(*) for totals
-- Use GROUP BY for statistics
-- Return SQL ONLY (no markdown, no explanation)
-- Ensure query ends with semicolon
+Rules:
+- ONLY SELECT queries
+- NEVER hallucinate columns
+- state column EXISTS
+- LIMIT 50 max
+- Output ONLY SQL
+
+Examples:
+
+Question:
+Show TN vehicles
 
-Example queries:
-- "How many cars today?" → SELECT COUNT(*) FROM vehicle_logs WHERE vehicle_type='car' AND date=CURRENT_DATE;
-- "Show TN vehicles" → SELECT * FROM vehicle_logs WHERE state='TN' ORDER BY timestamp DESC LIMIT 50;
-- "Top plates" → SELECT plate, COUNT(*) as count FROM vehicle_logs GROUP BY plate ORDER BY count DESC LIMIT 10;
+SQL:
+SELECT * FROM vehicle_logs
+WHERE state='TN'
+ORDER BY timestamp DESC
+LIMIT 50;
 
-Schema:
-{schema}
+Question:
+Top repeated plates
 
-User Question:
+SQL:
+SELECT plate, COUNT(*) AS count
+FROM vehicle_logs
+GROUP BY plate
+ORDER BY count DESC
+LIMIT 10;
+
+Question:
+Hourly traffic
+
+SQL:
+SELECT hour, COUNT(*) AS count
+FROM vehicle_logs
+GROUP BY hour
+ORDER BY hour;
+
+Question:
 {user_query}
 
-Generate SQL query:
+SQL:
 """
 
-    try:
-        response = client.text_generation(
-            prompt,
-            max_new_tokens=150,
-            temperature=0.1
-        )
+    response = client.text_generation(
+        prompt,
+        max_new_tokens=120,
+        temperature=0.1
+    )
 
-        sql_query = clean_sql(response.strip())
-        return sql_query, None
-    
-    except Exception as e:
-        print(f"❌ LLM error: {e}")
-        traceback.print_exc()
-        return None, f"LLM Error: {str(e)}"
+    sql = response.strip()
+
+    sql = sql.replace("```sql", "")
+    sql = sql.replace("```", "")
+    sql = sql.strip()
+
+    if ";" not in sql:
+        sql += ";"
 
+    return sql
 
-# ==================== SQL EXECUTION ====================
+
+# ================= QUERY ================= #
 
 def run_query(user_query):
-    """
-    Main query execution pipeline:
-    1. Convert natural language to SQL
-    2. Validate SQL
-    3. Add result limit
-    4. Execute query
-    5. Return results
-    """
-    
+
     if not db_available or engine is None:
         return {
             "query": user_query,
@@ -371,149 +327,137 @@ def run_query(user_query):
             "sql": None,
             "result": []
         }
-    
-    # Generate SQL from natural language
-    sql_query, llm_error = ask_llm(user_query)
-    
-    if llm_error:
-        return {
-            "query": user_query,
-            "error": llm_error,
-            "sql": None,
-            "result": []
-        }
-    
-    # Validate SQL
-    is_valid, validation_msg = validate_sql(sql_query)
-    if not is_valid:
-        return {
-            "query": user_query,
-            "error": f"SQL Validation Failed: {validation_msg}",
-            "sql": sql_query,
-            "result": []
-        }
-    
-    # Add result limit if not present
-    if "LIMIT" not in sql_query.upper():
-        sql_query = sql_query.replace(";", " LIMIT 50;")
-    
-    # Store in history
-    query_history.append({
-        "query": user_query,
-        "sql": sql_query,
-        "timestamp": datetime.now().isoformat()
-    })
-    
-    # Execute query
+
     try:
+
+        sql = ask_llm(user_query)
+
+        if not validate_sql(sql):
+
+            return {
+                "query": user_query,
+                "error": "Unsafe SQL blocked",
+                "sql": sql,
+                "result": []
+            }
+
         with engine.connect() as conn:
-            result = conn.execute(text(sql_query))
-            rows = [dict(row._mapping) for row in result]
+
+            result = conn.execute(text(sql))
+
+            rows = [
+                dict(r._mapping)
+                for r in result
+            ]
 
         return {
             "query": user_query,
-            "sql": sql_query,
+            "sql": sql,
             "result": rows,
-            "row_count": len(rows)
+            "count": len(rows)
         }
 
     except Exception as e:
-        print(f"❌ Query execution error: {e}")
+
         traceback.print_exc()
+
         return {
             "query": user_query,
-            "error": f"Execution Error: {str(e)}",
-            "sql": sql_query,
+            "error": str(e),
+            "sql": sql if "sql" in locals() else "",
             "result": []
         }
 
 
-# ==================== ANALYTICS QUERIES ====================
+# ================= ANALYTICS ================= #
 
 def get_vehicles_by_state():
-    """Get vehicle count by state"""
+
     if not db_available or engine is None:
         return []
-    
-    try:
-        with engine.connect() as conn:
-            result = conn.execute(text("""
-                SELECT state, COUNT(*) as count
-                FROM vehicle_logs
-                GROUP BY state
-                ORDER BY count DESC
-                LIMIT 20;
-            """))
-            return [dict(row._mapping) for row in result]
-    except Exception as e:
-        print(f"Error: {e}")
-        return []
 
+    with engine.connect() as conn:
+
+        result = conn.execute(text("""
+        SELECT state, COUNT(*) AS count
+        FROM vehicle_logs
+        GROUP BY state
+        ORDER BY count DESC
+        """))
+
+        return [dict(r._mapping) for r in result]
+
+
+def get_hourly_traffic():
 
-def get_hourly_traffic(date=None):
-    """Get traffic by hour"""
     if not db_available or engine is None:
         return []
-    
-    try:
-        date_filter = f"WHERE date = '{date}'" if date else ""
-        
-        with engine.connect() as conn:
-            result = conn.execute(text(f"""
-                SELECT hour, COUNT(*) as count
-                FROM vehicle_logs
-                {date_filter}
-                GROUP BY hour
-                ORDER BY hour;
-            """))
-            return [dict(row._mapping) for row in result]
-    except Exception as e:
-        print(f"Error: {e}")
-        return []
 
+    with engine.connect() as conn:
+
+        result = conn.execute(text("""
+        SELECT hour, COUNT(*) AS count
+        FROM vehicle_logs
+        GROUP BY hour
+        ORDER BY hour
+        """))
+
+        return [dict(r._mapping) for r in result]
+
+
+def get_top_plates():
 
-def get_top_plates(limit=10):
-    """Get most frequently detected plates"""
     if not db_available or engine is None:
         return []
-    
-    try:
-        with engine.connect() as conn:
-            result = conn.execute(text(f"""
-                SELECT plate, COUNT(*) as count
-                FROM vehicle_logs
-                WHERE plate != ''
-                GROUP BY plate
-                ORDER BY count DESC
-                LIMIT {limit};
-            """))
-            return [dict(row._mapping) for row in result]
-    except Exception as e:
-        print(f"Error: {e}")
-        return []
 
+    with engine.connect() as conn:
+
+        result = conn.execute(text("""
+        SELECT plate, COUNT(*) AS count
+        FROM vehicle_logs
+        WHERE plate != ''
+        GROUP BY plate
+        ORDER BY count DESC
+        LIMIT 10
+        """))
+
+        return [dict(r._mapping) for r in result]
+
+
+def get_suspicious_vehicles():
 
-def get_suspicious_vehicles(min_detections=20):
-    """Get vehicles detected multiple times (potential suspicious activity)"""
     if not db_available or engine is None:
         return []
-    
+
+    with engine.connect() as conn:
+
+        result = conn.execute(text("""
+        SELECT plate,
+               COUNT(*) AS detections
+        FROM vehicle_logs
+        GROUP BY plate
+        HAVING COUNT(*) > 5
+        ORDER BY detections DESC
+        LIMIT 20
+        """))
+
+        return [dict(r._mapping) for r in result]
+
+
+# ================= HEALTH ================= #
+
+def health_check():
+
+    if not db_available or engine is None:
+        return False, "⚠️ Database not configured. Set DATABASE_URL in Hugging Face Spaces secrets."
+
     try:
+
         with engine.connect() as conn:
-            result = conn.execute(text(f"""
-                SELECT plate, COUNT(*) as detection_count, state, vehicle_type
-                FROM vehicle_logs
-                WHERE plate != ''
-                GROUP BY plate, state, vehicle_type
-                HAVING COUNT(*) >= {min_detections}
-                ORDER BY detection_count DESC
-                LIMIT 50;
-            """))
-            return [dict(row._mapping) for row in result]
-    except Exception as e:
-        print(f"Error: {e}")
-        return []
+            conn.execute(text("SELECT 1"))
+
+        return True, "✅ Database Connected"
 
+    except Exception as e:
 
-# Initialize database on module load
-init_db()
+        return False, str(e)