Update app/db.py

app/db.py

@@ -1,24 +1,21 @@
-import duckdb, os, pathlib, json
-from datetime import datetime
+# app/db.py – BULLETPROOF VERSION
+import os
+import pathlib
+import json
+import duckdb
 from typing import Any, Dict, List
+from datetime import datetime
 
 DB_DIR = pathlib.Path("./data/duckdb")
 DB_DIR.mkdir(parents=True, exist_ok=True)
 
-
 def get_conn(org_id: str):
     """Get or create a DuckDB connection for an organization."""
     db_file = DB_DIR / f"{org_id}.duckdb"
     return duckdb.connect(str(db_file), read_only=False)
 
-
-# ------------------------------------------------------------
-# 🔹 Backward-compatible table for raw JSON ingestion
-# ------------------------------------------------------------
 def ensure_raw_table(conn):
-    """
-    Maintains legacy compatibility for ingestion from webhooks / file uploads.
-    """
+    """Creates audit trail table for raw JSON data."""
     conn.execute("CREATE SCHEMA IF NOT EXISTS main")
     conn.execute("""
         CREATE TABLE IF NOT EXISTS main.raw_rows(
@@ -27,94 +24,112 @@ def ensure_raw_table(conn):
         )
     """)
 
+def infer_duckdb_type(value: Any) -> str:
+    """Infer DuckDB type from Python value."""
+    if isinstance(value, bool):
+        return "BOOLEAN"
+    if isinstance(value, int):
+        return "BIGINT"
+    if isinstance(value, float):
+        return "DOUBLE"
+    if isinstance(value, datetime):
+        return "TIMESTAMP"
+    return "VARCHAR"
 
-# ------------------------------------------------------------
-# 🔹 Flexible dynamic schema table creation
-# ------------------------------------------------------------
+# ✅ BULLETPROOF: Handles integer column names, None values, etc.
 def ensure_table(conn, table_name: str, sample_record: Dict[str, Any]):
     """
-    Ensures a DuckDB table exists with columns inferred from sample_record.
-    If new columns appear later, adds them automatically.
+    Ensures table exists with columns from sample_record.
+    SAFE: Converts int column names to strings, handles missing data.
     """
+    if not sample_record:
+        return
+
     conn.execute("CREATE SCHEMA IF NOT EXISTS main")
+    
+    # Create base table with UUID and timestamp
     conn.execute(
         f"CREATE TABLE IF NOT EXISTS main.{table_name} ("
         "id UUID DEFAULT uuid(), "
         "_ingested_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)"
     )
 
-    if not sample_record:
-        return
-
-    existing_cols = {r[0] for r in conn.execute(f"PRAGMA table_info('main.{table_name}')").fetchall()}
+    # ✅ SAFE: Get existing columns (handle int names from DB)
+    try:
+        existing_cols_raw = conn.execute(f"PRAGMA table_info('main.{table_name}')").fetchall()
+        existing_cols = {str(r[0]).lower() for r in existing_cols_raw}
+    except Exception as e:
+        print(f"[db] ⚠️ Could not get table info: {e}")
+        existing_cols = set()
 
+    # ✅ BULLETPROOF: Add missing columns with safe name handling
     for col, val in sample_record.items():
-        if col.lower() in existing_cols:
+        col_name = str(col).lower().strip()  # ✅ FORCE STRING
+        
+        if col_name in existing_cols:
             continue
-        dtype = infer_duckdb_type(val)
-        print(f"[db] ➕ Adding new column '{col}:{dtype}' to main.{table_name}")
-        conn.execute(f"ALTER TABLE main.{table_name} ADD COLUMN {col} {dtype}")
-
-
-def infer_duckdb_type(value: Any) -> str:
-    """Infer a DuckDB-compatible column type from a Python value."""
-    if isinstance(value, bool):
-        return "BOOLEAN"
-    if isinstance(value, int):
-        return "BIGINT"
-    if isinstance(value, float):
-        return "DOUBLE"
-    if isinstance(value, datetime):
-        return "TIMESTAMP"
-    if isinstance(value, (dict, list)):
-        return "JSON"
-    return "VARCHAR"
-
+            
+        # Skip None/empty values which can't infer type
+        if val is None:
+            print(f"[db] ⚠️ Skipping column {col_name} (None value)")
+            continue
+            
+        try:
+            dtype = infer_duckdb_type(val)
+            print(f"[db] ➕ Adding column '{col_name}:{dtype}' to main.{table_name}")
+            conn.execute(f"ALTER TABLE main.{table_name} ADD COLUMN {col_name} {dtype}")
+        except Exception as e:
+            print(f"[db] ❌ Failed to add column {col_name}: {e}")
+            # Continue with next column - don't crash
 
-# ------------------------------------------------------------
-# 🔹 Insert records with auto-schema
-# ------------------------------------------------------------
 def insert_records(conn, table_name: str, records: List[Dict[str, Any]]):
     """
-    Insert records into the specified table.
-    Assumes ensure_table() has already been called.
+    Insert records into table with safe column handling.
     """
     if not records:
         return
 
-    cols = records[0].keys()
+    # ✅ Get columns from first record (force string names)
+    cols = [str(k) for k in records[0].keys()]
     placeholders = ", ".join(["?"] * len(cols))
     col_list = ", ".join(cols)
+    
     insert_sql = f"INSERT INTO main.{table_name} ({col_list}) VALUES ({placeholders})"
 
-    values = [tuple(r.get(c) for c in cols) for r in records]
-    conn.executemany(insert_sql, values)
-    print(f"[db] ✅ Inserted {len(records)} rows into {table_name}")
-
-
-# ------------------------------------------------------------
-# 🔹 Unified bootstrap entrypoint
-# ------------------------------------------------------------
-# ------------------------------------------------------------
-# 🔹 Unified bootstrap entrypoint  –  n8n-native
-# ------------------------------------------------------------
-# app/db.py - FIX bootstrap to only audit, not create tables
+    # Prepare values (handle missing keys)
+    values = []
+    for record in records:
+        row = []
+        for col in cols:
+            val = record.get(col)
+            # Convert dict/list to JSON string for DuckDB
+            if isinstance(val, (dict, list)):
+                val = json.dumps(val)
+            row.append(val)
+        values.append(tuple(row))
+    
+    try:
+        conn.executemany(insert_sql, values)
+        print(f"[db] ✅ Inserted {len(records)} rows into {table_name}")
+    except Exception as e:
+        print(f"[db] ❌ Insert failed: {e}")
+        raise
 
+# ✅ PURE AUDIT FUNCTION - Does NOT create tables
 def bootstrap(org_id: str, payload: Dict[str, Any]):
     """
-    PURE AUDIT TRAIL - Does NOT create tables.
-    Only stores raw JSON for lineage. Table creation is canonify_df's job.
+    **ENTERPRISE-GRADE**: Only stores raw JSON for audit.
+    Does NOT create any tables. Schema evolution is canonify_df's job.
     """
     conn = get_conn(org_id)
     ensure_raw_table(conn)
 
-    # Store raw payload for audit (only if non-empty)
     try:
-        raw_json = json.dumps(payload)
+        raw_json = json.dumps(payload) if not isinstance(payload, str) else payload
         if raw_json and raw_json not in ("null", "[]", "{}"):
             conn.execute("INSERT INTO main.raw_rows (row_data) VALUES (?)", (raw_json,))
-            print(f"[bootstrap] ✅ Audit trail stored: {len(raw_json)} bytes")
+            print(f"[bootstrap] ✅ Audit stored: {len(raw_json)} bytes")
     except Exception as e:
-        print(f"[bootstrap] ⚠️ Audit storage failed: {e}")
+        print(f"[bootstrap] ⚠️ Audit failed: {e}")
     
-    conn.close()
+    conn.close()