Upload main.py

server/main.py

@@ -1,7 +1,12 @@
 import asyncio
+import csv
+import io
 import os
+import re
+import sqlite3
+from typing import Any
 
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, UploadFile, File
 from fastapi.responses import HTMLResponse
 
 from .environment import SQLQualityEnvironment
@@ -23,6 +28,11 @@ app = FastAPI(
 # Module-level singleton — ONE instance shared across all requests
 _env = SQLQualityEnvironment()
 
+# Module-level custom database connection (separate from RL environment)
+_custom_db_conn: sqlite3.Connection | None = None
+_custom_db_schema: list[dict] = []
+_custom_db_filename: str = ""
+
 # HTTP-level request lock to serialize concurrent requests
 _request_lock = asyncio.Lock()
 
@@ -170,6 +180,266 @@ async def state() -> SQLState:
         raise HTTPException(status_code=422, detail=str(e))
 
 
+# ─── CUSTOM DATABASE ENDPOINTS ─────────────────────────────────────────
+
+def infer_sqlite_type(values: list[str]) -> str:
+    """Infer SQLite column type from a list of string values."""
+    non_empty = [v for v in values if v and v.strip()]
+    if not non_empty:
+        return 'TEXT'
+    try:
+        for v in non_empty:
+            int(v)
+        return 'INTEGER'
+    except ValueError:
+        pass
+    try:
+        for v in non_empty:
+            float(v)
+        return 'REAL'
+    except ValueError:
+        pass
+    return 'TEXT'
+
+
+def get_db_schema(conn: sqlite3.Connection) -> list[dict]:
+    """Extract schema information from a SQLite connection."""
+    schema = []
+    cursor = conn.execute("SELECT name FROM sqlite_master WHERE type='table'")
+    tables = cursor.fetchall()
+    
+    for (table_name,) in tables:
+        # Get row count
+        count_cursor = conn.execute(f'SELECT COUNT(*) FROM "{table_name}"')
+        row_count = count_cursor.fetchone()[0]
+        
+        # Get columns
+        pragma_cursor = conn.execute(f'PRAGMA table_info("{table_name}")')
+        columns = []
+        for row in pragma_cursor.fetchall():
+            col_name = row[1]
+            col_type = row[2]
+            columns.append(f"{col_name} ({col_type})")
+        
+        schema.append({
+            "table_name": table_name,
+            "row_count": row_count,
+            "columns": columns
+        })
+    
+    return schema
+
+
+@app.post("/upload-db")
+async def upload_db(file: UploadFile = File(...)) -> dict:
+    """Upload a .sqlite, .db, or .csv file for custom exploration.
+    
+    Args:
+        file: The uploaded file (.sqlite, .db, or .csv)
+        
+    Returns:
+        dict: Status and schema information for the uploaded database
+        
+    Raises:
+        HTTPException 400: If file extension is invalid
+        HTTPException 413: If file is too large (>15MB)
+        HTTPException 500: If processing fails
+    """
+    global _custom_db_conn, _custom_db_schema, _custom_db_filename
+    
+    # Validate file extension
+    filename = file.filename or ""
+    ext = os.path.splitext(filename)[1].lower()
+    valid_extensions = {'.sqlite', '.db', '.csv'}
+    
+    if ext not in valid_extensions:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Invalid file extension: {ext}. Must be one of: {valid_extensions}"
+        )
+    
+    # Read file content
+    content = await file.read()
+    
+    # Validate file size (15MB limit)
+    max_size = 15 * 1024 * 1024
+    if len(content) > max_size:
+        raise HTTPException(
+            status_code=413,
+            detail=f"File too large: {len(content)} bytes. Maximum is {max_size} bytes (15MB)"
+        )
+    
+    try:
+        # Close existing connection if any
+        if _custom_db_conn is not None:
+            try:
+                _custom_db_conn.close()
+            except Exception:
+                pass
+            _custom_db_conn = None
+        
+        if ext == '.csv':
+            # Parse CSV and create in-memory SQLite database
+            conn = sqlite3.connect(":memory:", check_same_thread=False)
+            
+            # Read CSV content
+            csv_text = content.decode('utf-8')
+            csv_file = io.StringIO(csv_text)
+            reader = csv.DictReader(csv_file)
+            
+            if not reader.fieldnames:
+                raise HTTPException(status_code=400, detail="CSV file has no headers")
+            
+            # Infer types from first 100 rows
+            rows = list(reader)
+            if not rows:
+                raise HTTPException(status_code=400, detail="CSV file is empty")
+            
+            # Create table
+            table_name = os.path.splitext(os.path.basename(filename))[0]
+            # Sanitize table name
+            table_name = re.sub(r'[^a-zA-Z0-9_]', '_', table_name)
+            if not table_name[0].isalpha() and table_name[0] != '_':
+                table_name = '_' + table_name
+            
+            columns = []
+            for col_name in reader.fieldnames:
+                # Sanitize column name
+                safe_col = re.sub(r'[^a-zA-Z0-9_]', '_', col_name)
+                if not safe_col[0].isalpha() and safe_col[0] != '_':
+                    safe_col = '_' + safe_col
+                
+                values = [row.get(col_name, '') for row in rows[:100]]
+                col_type = infer_sqlite_type(values)
+                columns.append(f'"{safe_col}" {col_type}')
+            
+            create_sql = f'CREATE TABLE "{table_name}" ({', '.join(columns)})'
+            conn.execute(create_sql)
+            
+            # Insert rows
+            placeholders = ', '.join(['?' for _ in reader.fieldnames])
+            insert_sql = f'INSERT INTO "{table_name}" VALUES ({placeholders})'
+            
+            for row in rows:
+                values = [row.get(col_name, None) for col_name in reader.fieldnames]
+                conn.execute(insert_sql, values)
+            
+            conn.commit()
+            _custom_db_conn = conn
+            
+        else:
+            # .sqlite or .db file - write to temp and open
+            import tempfile
+            tmp = tempfile.NamedTemporaryFile(suffix=ext, delete=False)
+            tmp.write(content)
+            tmp.close()
+            
+            conn = sqlite3.connect(tmp.name, check_same_thread=False, uri=True)
+            conn.execute("PRAGMA query_only = ON")  # Read-only mode for safety
+            _custom_db_conn = conn
+        
+        # Get schema
+        _custom_db_schema = get_db_schema(_custom_db_conn)
+        _custom_db_filename = filename
+        
+        return {
+            "status": "ok",
+            "filename": filename,
+            "tables": _custom_db_schema
+        }
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to process file: {str(e)}")
+
+
+@app.post("/query-custom")
+async def query_custom(body: dict) -> dict:
+    """Execute a SELECT query against the uploaded custom database.
+    
+    Args:
+        body: JSON body with {"sql": "SELECT ..."}
+        
+    Returns:
+        dict: Query results with rows, columns, and row_count
+        
+    Raises:
+        HTTPException 422: If no database uploaded or SQL is invalid
+    """
+    global _custom_db_conn
+    
+    if _custom_db_conn is None:
+        raise HTTPException(
+            status_code=422,
+            detail="No database uploaded yet. Use /upload-db first."
+        )
+    
+    sql = body.get("sql", "").strip()
+    if not sql:
+        raise HTTPException(status_code=422, detail="No SQL provided")
+    
+    # Validate SQL starts with SELECT (same rule as QUERY action)
+    sql_upper = sql.upper()
+    if not (sql_upper.startswith("SELECT") or sql_upper.startswith("WITH") or 
+            sql_upper.startswith("EXPLAIN") or sql_upper.startswith("PRAGMA")):
+        return {
+            "error": "QUERY action is read-only. This action only accepts SELECT, WITH, EXPLAIN, or PRAGMA.",
+            "message": "SQL Error — check your query. Only SELECT queries are allowed in custom mode."
+        }
+    
+    try:
+        cursor = _custom_db_conn.execute(sql)
+        
+        if cursor.description:
+            columns = [d[0] for d in cursor.description]
+            rows = cursor.fetchall()
+            
+            # Convert rows to dicts
+            row_dicts = []
+            for row in rows:
+                row_dict = {}
+                for i, col in enumerate(columns):
+                    row_dict[col] = row[i]
+                row_dicts.append(row_dict)
+            
+            return {
+                "rows": row_dicts,
+                "columns": columns,
+                "row_count": len(row_dicts)
+            }
+        else:
+            return {
+                "rows": [],
+                "columns": [],
+                "row_count": 0,
+                "message": "Query executed successfully but returned no rows"
+            }
+            
+    except sqlite3.Error as e:
+        return {
+            "error": str(e),
+            "message": "SQL Error — check your query. Check table and column names in the schema panel."
+        }
+
+
+@app.get("/custom-db-status")
+async def custom_db_status() -> dict:
+    """Check if a custom database is loaded and return its schema.
+    
+    Returns:
+        dict: {loaded: true/false, tables: [...]} 
+    """
+    if _custom_db_conn is None:
+        return {"loaded": False}
+    
+    return {
+        "loaded": True,
+        "filename": _custom_db_filename,
+        "tables": _custom_db_schema
+    }
+
+
 def run_server() -> None:
     """
     Entry point for 'uv run server' command.