Update app.py

app.py

@@ -1,105 +1,153 @@
-import subprocess
+import sqlite3
+import os
 import sys
-import uuid
 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"
+
+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")
+
+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()
+
 @app.route('/')
 def home():
-    return jsonify({"status": "ok", "message": "API Running (Isolated Sessions Enabled)"})
+    return jsonify({"status": "ok", "message": "SQL Compiler with File Persistence Ready"})
 
-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--- Error Log ---\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)}"
-
-@app.route('/api/run-python', methods=['POST'])
-def run_python():
-    code = request.json.get('code', '')
-    user_input = request.json.get('input', '')
+@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
     
-    if not code: return jsonify({'output': 'No code provided'}), 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
 
-    cmd = [sys.executable, "-c", code]
-    output = run_command(cmd, input_text=user_input)
-    return jsonify({'output': output})
+    try:
+        conn = sqlite3.connect(db_path)
+        cursor = conn.cursor()
+        
+        # Get all tables
+        cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
+        tables = [row[0] for row in cursor.fetchall()]
+        
+        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})
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
 
 @app.route('/api/run-sql', methods=['POST'])
 def run_sql():
     data = request.json
-    sql_code = data.get('code', '')
-    db_type = data.get('type', 'sqlite').lower()
+    session_id = data.get('session_id')
+    sql_code = data.get('code')
     
-    if not sql_code: return jsonify({'output': 'No SQL code provided'}), 400
+    if not session_id or not sql_code:
+        return jsonify({'output': 'Missing session_id or SQL code'}), 400
 
-    output = ""
-    session_id = uuid.uuid4().hex[:8] # Random 8-char ID for isolation
+    db_path = get_db_path(session_id)
+    
+    # If DB doesn't exist yet, create sample
+    if not os.path.exists(db_path):
+        create_sample_db(db_path)
 
+    output = []
     try:
-        # 1. SQLITE (Already Isolated by :memory:)
-        if db_type == 'sqlite':
-            cmd = ["sqlite3", ":memory:"]
-            output = run_command(cmd, input_text=sql_code)
-
-        # 2. POSTGRESQL (Schema Isolation)
-        elif db_type in ['postgres', 'postgresql']:
-            schema_name = f"sess_{session_id}"
+        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()
             
-            # A. Create Temp Schema
-            setup_cmd = ["psql", "-h", "localhost", "-U", "playground", "-d", "playground", "-c", f"CREATE SCHEMA {schema_name};"]
-            run_command(setup_cmd)
-
-            # B. Run User Code (Wrapped in search_path to use the schema)
-            # We modify the SQL to force it into the new schema
-            isolated_sql = f"SET search_path TO {schema_name};\n{sql_code}"
+            # ASCII Table Formatting
+            output.append(f"| {' | '.join(columns)} |")
+            output.append("-" * len(output[-1]))
+            for row in rows:
+                output.append(f"| {' | '.join(map(str, row))} |")
             
-            main_cmd = ["psql", "-h", "localhost", "-U", "playground", "-d", "playground"]
-            output = run_command(main_cmd, input_text=isolated_sql)
-
-            # C. Cleanup (Drop Schema)
-            cleanup_cmd = ["psql", "-h", "localhost", "-U", "playground", "-d", "playground", "-c", f"DROP SCHEMA {schema_name} CASCADE;"]
-            run_command(cleanup_cmd)
-
-        # 3. MYSQL (Database Isolation)
-        elif db_type == 'mysql':
-            db_name = f"sess_{session_id}"
-
-            # A. Create Temp Database
-            setup_cmd = ["mysql", "-h", "localhost", "-u", "playground", "-e", f"CREATE DATABASE {db_name};"]
-            run_command(setup_cmd)
-
-            # B. Run User Code (Against the new DB)
-            main_cmd = ["mysql", "-h", "localhost", "-u", "playground", "-D", db_name, "-t"]
-            output = run_command(main_cmd, input_text=sql_code)
-
-            # C. Cleanup (Drop Database)
-            cleanup_cmd = ["mysql", "-h", "localhost", "-u", "playground", "-e", f"DROP DATABASE {db_name};"]
-            run_command(cleanup_cmd)
-
+            output.append(f"\n({len(rows)} rows returned)")
         else:
-            return jsonify({'output': 'Unsupported database type.'}), 400
+            output.append(f"Query executed successfully. (Rows affected: {conn.total_changes})")
 
+        conn.commit()
+        conn.close()
+        
+    except sqlite3.Error as e:
+        return jsonify({'error': str(e)})
     except Exception as e:
-        output = f"Server Error: {str(e)}"
+        return jsonify({'error': f"System Error: {str(e)}"})
 
-    return jsonify({'output': output})
+    return jsonify({'output': "\n".join(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)