Update app.py

app.py

@@ -1,153 +1,218 @@
-import sqlite3
-import os
+import subprocess
 import sys
+import uuid
+import os
+import tempfile
+import io
+import pandas as pd
 from flask import Flask, request, jsonify
 from flask_cors import CORS
 
 app = Flask(__name__)
 CORS(app)
 
-# Store databases in a temporary directory
-DB_STORAGE_PATH = "/tmp"
+# --- STANDARD SAMPLE DATA (Same for everyone) ---
+SAMPLE_SQL = """
+-- Standard Sample Data
+CREATE TABLE Departments (
+    dept_id INT PRIMARY KEY,
+    dept_name VARCHAR(50),
+    location VARCHAR(50)
+);
 
-def get_db_path(session_id):
-    # Security: Ensure session_id is alphanumeric only to prevent path traversal
-    safe_id = "".join(x for x in session_id if x.isalnum())
-    return os.path.join(DB_STORAGE_PATH, f"{safe_id}.db")
+INSERT INTO Departments VALUES (1, 'Engineering', 'New York');
+INSERT INTO Departments VALUES (2, 'HR', 'London');
+INSERT INTO Departments VALUES (3, 'Marketing', 'San Francisco');
+INSERT INTO Departments VALUES (4, 'Sales', 'Tokyo');
 
-def create_sample_db(path):
-    """Creates a dummy database with Employees and Departments"""
-    if os.path.exists(path):
-        os.remove(path)
-    
-    conn = sqlite3.connect(path)
-    cursor = conn.cursor()
-    
-    # Create Tables
-    cursor.execute('''CREATE TABLE Departments (id INTEGER PRIMARY KEY, name TEXT, budget INTEGER)''')
-    cursor.execute('''CREATE TABLE Employees (id INTEGER PRIMARY KEY, name TEXT, role TEXT, salary INTEGER, dept_id INTEGER, FOREIGN KEY(dept_id) REFERENCES Departments(id))''')
-    
-    # Insert Data
-    cursor.executemany('INSERT INTO Departments VALUES (?,?,?)', [
-        (1, 'Engineering', 150000),
-        (2, 'Marketing', 80000),
-        (3, 'HR', 50000)
-    ])
-    cursor.executemany('INSERT INTO Employees VALUES (?,?,?,?,?)', [
-        (101, 'Alice', 'Engineer', 90000, 1),
-        (102, 'Bob', 'Manager', 95000, 2),
-        (103, 'Charlie', 'Analyst', 60000, 2),
-        (104, 'David', 'Engineer', 88000, 1),
-        (105, 'Eve', 'Recruiter', 55000, 3)
-    ])
-    conn.commit()
-    conn.close()
+CREATE TABLE Employees (
+    emp_id INT PRIMARY KEY,
+    name VARCHAR(50),
+    dept_id INT,
+    role VARCHAR(50),
+    salary DECIMAL(10, 2)
+    -- Note: Foreign keys might be loose depending on DB strictness settings
+);
+
+INSERT INTO Employees VALUES (101, 'Alice Smith', 1, 'Engineer', 90000);
+INSERT INTO Employees VALUES (102, 'Bob Jones', 1, 'Senior Engineer', 110000);
+INSERT INTO Employees VALUES (103, 'Charlie Brown', 2, 'Recruiter', 60000);
+INSERT INTO Employees VALUES (104, 'Diana Prince', 3, 'CMO', 150000);
+INSERT INTO Employees VALUES (105, 'Evan Wright', 1, 'Intern', 40000);
+INSERT INTO Employees VALUES (106, 'Fiona Green', 4, 'Sales Rep', 75000);
+"""
 
 @app.route('/')
 def home():
-    return jsonify({"status": "ok", "message": "SQL Compiler with File Persistence Ready"})
+    return jsonify({"status": "ok", "message": "SQL Compiler API Ready"})
 
-@app.route('/api/session/init', methods=['POST'])
-def init_session():
-    # Allows user to request a fresh sample DB
-    session_id = request.json.get('session_id')
-    if not session_id: return jsonify({'error': 'Missing session_id'}), 400
-    
-    db_path = get_db_path(session_id)
-    create_sample_db(db_path)
-    return jsonify({'status': 'initialized', 'message': 'Sample Database Loaded'})
-
-@app.route('/api/session/upload', methods=['POST'])
-def upload_db():
-    session_id = request.form.get('session_id')
-    file = request.files.get('file')
-    
-    if not session_id or not file:
-        return jsonify({'error': 'Missing session_id or file'}), 400
-        
-    db_path = get_db_path(session_id)
-    file.save(db_path)
-    return jsonify({'status': 'uploaded', 'message': 'Database Uploaded Successfully'})
-
-@app.route('/api/schema', methods=['POST'])
-def get_schema():
-    """Returns list of tables and columns for the sidebar"""
-    session_id = request.json.get('session_id')
-    db_path = get_db_path(session_id)
-    
-    if not os.path.exists(db_path):
-        return jsonify({'error': 'No database found for this session'}), 404
+def run_command(command, input_text=None, timeout=15):
+    """Helper to run shell commands safely"""
+    try:
+        run_process = subprocess.run(
+            command,
+            input=input_text,
+            capture_output=True,
+            text=True,
+            timeout=timeout
+        )
+        output = run_process.stdout
+        if run_process.stderr:
+             output += "\n--- Database Messages/Errors ---\n" + run_process.stderr
+        return output
+    except subprocess.TimeoutExpired:
+        return "Error: Execution timed out (15s limit)."
+    except Exception as e:
+        return f"System Error: {str(e)}"
 
+def file_to_sql(file, filename):
+    """Converts CSV/Excel file objects to SQL CREATE/INSERT strings"""
     try:
-        conn = sqlite3.connect(db_path)
-        cursor = conn.cursor()
+        # 1. Read Data
+        if filename.lower().endswith('.csv'):
+            df = pd.read_csv(file)
+        else:
+            df = pd.read_excel(file)
         
-        # Get all tables
-        cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
-        tables = [row[0] for row in cursor.fetchall()]
+        # 2. Sanitize Table Name (filename without extension, alphanumeric only)
+        table_name = os.path.splitext(filename)[0]
+        table_name = "".join([c if c.isalnum() else "_" for c in table_name])
+        if not table_name: table_name = "uploaded_table"
+
+        # 3. Sanitize Column Names
+        df.columns = ["".join([c if c.isalnum() else "_" for c in str(col)]) for col in df.columns]
         
-        schema = {}
-        for table in tables:
-            cursor.execute(f"PRAGMA table_info({table})")
-            # Returns (cid, name, type, notnull, dflt_value, pk)
-            columns = [{"name": col[1], "type": col[2]} for col in cursor.fetchall()]
-            schema[table] = columns
-            
-        conn.close()
-        return jsonify({'schema': schema})
+        # 4. Generate CREATE TABLE
+        cols = []
+        for col, dtype in zip(df.columns, df.dtypes):
+            sql_type = "TEXT"
+            if "int" in str(dtype): sql_type = "INTEGER"
+            elif "float" in str(dtype): sql_type = "DECIMAL(10,2)"
+            cols.append(f"{col} {sql_type}")
+        
+        create_stmt = f"CREATE TABLE {table_name} ({', '.join(cols)});\n"
+        
+        # 5. Generate INSERT Statements
+        # Note: For very large files, bulk loading is better, but for "Sample" sizes, this is universally compatible
+        inserts = []
+        for _, row in df.iterrows():
+            vals = []
+            for v in row:
+                if pd.isna(v): 
+                    vals.append("NULL")
+                elif isinstance(v, str): 
+                    vals.append(f"'{str(v).replace("'", "''")}'") # Escape single quotes
+                else: 
+                    vals.append(str(v))
+            inserts.append(f"INSERT INTO {table_name} VALUES ({', '.join(vals)});")
+        
+        return create_stmt + "\n".join(inserts) + "\n", table_name
+    
     except Exception as e:
-        return jsonify({'error': str(e)}), 500
+        return f"-- Error converting file: {str(e)}\n", None
 
 @app.route('/api/run-sql', methods=['POST'])
 def run_sql():
-    data = request.json
-    session_id = data.get('session_id')
-    sql_code = data.get('code')
+    sql_code = request.form.get('code', '')
+    db_type = request.form.get('type', 'sqlite').lower()
+    use_sample = request.form.get('use_sample') == 'true'
+    uploaded_file = request.files.get('database')
     
-    if not session_id or not sql_code:
-        return jsonify({'output': 'Missing session_id or SQL code'}), 400
+    if not sql_code: 
+        return jsonify({'output': 'No SQL code provided'}), 400
 
-    db_path = get_db_path(session_id)
+    output = ""
+    session_id = uuid.uuid4().hex[:8]
+    temp_file_path = None
+    
+    # Prepare Initialization SQL (Sample Data + Uploaded CSV/Excel SQL)
+    init_sql = ""
     
-    # If DB doesn't exist yet, create sample
-    if not os.path.exists(db_path):
-        create_sample_db(db_path)
+    if use_sample:
+        init_sql += SAMPLE_SQL + "\n"
+        
+    if uploaded_file:
+        fname = uploaded_file.filename.lower()
+        if fname.endswith('.csv') or fname.endswith('.xlsx') or fname.endswith('.xls'):
+            # Convert Data file to SQL
+            converted_sql, table_name = file_to_sql(uploaded_file, uploaded_file.filename)
+            init_sql += f"-- Imported from {uploaded_file.filename}\n" + converted_sql + "\n"
+            # Reset uploaded_file because we consumed it; we don't need to pass it as a binary DB
+            uploaded_file = None 
 
-    output = []
     try:
-        conn = sqlite3.connect(db_path)
-        cursor = conn.cursor()
-        
-        # Execute script allows multiple statements
-        cursor.executescript(sql_code)
-        
-        # Check if the last query returned data (SELECT)
-        if cursor.description:
-            columns = [desc[0] for desc in cursor.description]
-            rows = cursor.fetchall()
+        # --- 1. SQLITE LOGIC ---
+        if db_type == 'sqlite':
+            db_target = ":memory:"
             
-            # ASCII Table Formatting
-            output.append(f"| {' | '.join(columns)} |")
-            output.append("-" * len(output[-1]))
-            for row in rows:
-                output.append(f"| {' | '.join(map(str, row))} |")
+            # If user uploaded a binary .db file (and didn't convert it above)
+            if uploaded_file and (uploaded_file.filename.endswith('.db') or uploaded_file.filename.endswith('.sqlite')):
+                fd, temp_file_path = tempfile.mkstemp(suffix='.db')
+                os.close(fd)
+                uploaded_file.save(temp_file_path)
+                db_target = temp_file_path
             
-            output.append(f"\n({len(rows)} rows returned)")
+            cmd = ["sqlite3", db_target]
+            header = ".mode column\n.headers on\n" 
+            
+            # Combine: Header + Init SQL (Sample/CSV) + User SQL
+            full_script = header + init_sql + sql_code
+            output = run_command(cmd, input_text=full_script)
+
+        # --- 2. POSTGRESQL LOGIC ---
+        elif db_type in ['postgres', 'postgresql']:
+            schema_name = f"sess_{session_id}"
+            
+            # A. Create Schema
+            run_command(["psql", "-h", "localhost", "-U", "playground", "-d", "playground", "-c", f"CREATE SCHEMA {schema_name};"])
+
+            try:
+                # B. Handle Binary Dump Upload (if any - distinct from CSV/Excel handled in init_sql)
+                if uploaded_file: # If it's a .sql dump
+                    dump_content = uploaded_file.read().decode('utf-8', errors='ignore')
+                    init_sql += f"\n{dump_content}\n"
+
+                # C. Run Everything (Init + User Code) inside Schema
+                # We wrap everything in one transaction block pointing to the schema
+                full_script = f"SET search_path TO {schema_name};\n{init_sql}\n{sql_code}"
+                
+                output = run_command(["psql", "-h", "localhost", "-U", "playground", "-d", "playground"], input_text=full_script)
+
+            finally:
+                run_command(["psql", "-h", "localhost", "-U", "playground", "-d", "playground", "-c", f"DROP SCHEMA {schema_name} CASCADE;"])
+
+        # --- 3. MYSQL LOGIC ---
+        elif db_type == 'mysql':
+            db_name = f"sess_{session_id}"
+
+            # A. Create DB
+            run_command(["mysql", "-h", "localhost", "-u", "playground", "-e", f"CREATE DATABASE {db_name};"])
+
+            try:
+                 # B. Handle Binary Dump Upload
+                if uploaded_file:
+                    dump_content = uploaded_file.read().decode('utf-8', errors='ignore')
+                    init_sql += f"\n{dump_content}\n"
+
+                # C. Run Everything
+                # MySQL accepts multiple statements if we just pipe them
+                full_script = f"USE {db_name};\n{init_sql}\n{sql_code}"
+                
+                output = run_command(["mysql", "-h", "localhost", "-u", "playground", "-t"], input_text=full_script)
+
+            finally:
+                run_command(["mysql", "-h", "localhost", "-u", "playground", "-e", f"DROP DATABASE {db_name};"])
+
         else:
-            output.append(f"Query executed successfully. (Rows affected: {conn.total_changes})")
+            return jsonify({'output': 'Unsupported database type.'}), 400
 
-        conn.commit()
-        conn.close()
-        
-    except sqlite3.Error as e:
-        return jsonify({'error': str(e)})
     except Exception as e:
-        return jsonify({'error': f"System Error: {str(e)}"})
+        output = f"Server Error: {str(e)}"
+    finally:
+        if temp_file_path and os.path.exists(temp_file_path):
+            os.remove(temp_file_path)
 
-    return jsonify({'output': "\n".join(output)})
+    return jsonify({'output': output})
 
 if __name__ == '__main__':
-    # Ensure tmp dir exists
-    if not os.path.exists(DB_STORAGE_PATH):
-        os.makedirs(DB_STORAGE_PATH)
     app.run(host="0.0.0.0", port=7860)