Update app.py

app.py

@@ -1,9 +1,6 @@
 import os
 import subprocess
 import tempfile
-import sqlite3
-import psycopg2
-import pymysql
 import sys
 from flask import Flask, request, jsonify
 from flask_cors import CORS
@@ -13,163 +10,96 @@ CORS(app)
 
 @app.route('/')
 def home():
-    return jsonify({
-        "status": "ok",
-        "message": "Playground Backend Running. Use /api/run-python or /api/run-sql"
-    })
+    return jsonify({"status": "ok", "message": "API Running. POST to /api/run-python or /api/run-sql"})
 
-# --- PYTHON RUNNER ---
-@app.route('/api/run-python', methods=['POST'])
-def run_python():
-    code = request.json.get('code', '')
-    user_input = request.json.get('input', '')
-    if not code:
-        return jsonify({'error': 'No code provided.'}), 400
-    
-    temp_file_path = None
+def run_command(command, input_text=None, timeout=15):
+    """Helper to run shell commands safely"""
     try:
-        with tempfile.NamedTemporaryFile(suffix=".py", delete=False, mode='w', encoding='utf-8') as temp_file:
-            temp_file.write(code)
-            temp_file_path = temp_file.name
-
         run_process = subprocess.run(
-            [sys.executable, temp_file_path], # Use current python executable
-            input=user_input,
+            command,
+            input=input_text,
             capture_output=True,
             text=True,
-            timeout=15
+            timeout=timeout
         )
-        
-        # Return stdout if successful, or stderr if failed
-        if run_process.returncode != 0:
-             return jsonify({'output': run_process.stderr or run_process.stdout})
-        else:
-             return jsonify({'output': run_process.stdout})
-
+        # Return stdout if success, stderr if fail (or both)
+        output = run_process.stdout
+        if run_process.stderr:
+             # Append stderr to output so user sees errors
+             output += "\n--- Error Log ---\n" + run_process.stderr
+        return output
     except subprocess.TimeoutExpired:
-        return jsonify({'output': 'Error: Code execution timed out (15s limit).'})
+        return "Error: Execution timed out (15s limit)."
     except Exception as e:
-        return jsonify({'output': f"Error: {str(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', '') # Optional stdin input
+    
+    if not code: return jsonify({'output': 'No code provided'}), 400
+
+    # Write code to temp file
+    temp_path = None
+    try:
+        with tempfile.NamedTemporaryFile(suffix=".py", delete=False, mode='w', encoding='utf-8') as f:
+            f.write(code)
+            temp_path = f.name
+        
+        # Run python using the current environment's python executable
+        output = run_command([sys.executable, temp_path], input_text=user_input)
+        return jsonify({'output': output})
+        
     finally:
-        if temp_file_path and os.path.exists(temp_file_path):
-            os.remove(temp_file_path)
+        if temp_path and os.path.exists(temp_path):
+            os.remove(temp_path)
 
-# --- SQL RUNNER ---
 @app.route('/api/run-sql', methods=['POST'])
 def run_sql():
     data = request.json
     sql_code = data.get('code', '')
-    # Default to sqlite if not specified
     db_type = data.get('type', 'sqlite').lower()
 
-    if not sql_code:
-        return jsonify({'error': 'No SQL code provided.'}), 400
-
-    output_log = ""
+    if not sql_code: return jsonify({'output': 'No SQL code provided'}), 400
 
+    # Write SQL to temp file (safer for some CLI tools)
+    temp_sql_path = None
     try:
-        # 1. SQLITE (In-Memory)
+        with tempfile.NamedTemporaryFile(suffix=".sql", delete=False, mode='w', encoding='utf-8') as f:
+            f.write(sql_code)
+            temp_sql_path = f.name
+
+        output = ""
+        
         if db_type == 'sqlite':
-            # Create a fresh in-memory DB every time
-            con = sqlite3.connect(':memory:')
-            cur = con.cursor()
-            try:
-                # executescript allows multiple statements (CREATE + INSERT + SELECT)
-                cur.executescript(sql_code)
-                output_log = "Executed successfully.\n"
-                
-                # Simple heuristic: If user didn't print anything, try to verify tables
-                # But sqlite3 executescript doesn't return rows. 
-                # Use a specific fetch loop if you want to support SELECTs explicitly
-                # or tell users to use PRINT in python. 
-                # For SQL playground, we usually just run it. 
-                # To show data, we can iterate cursor if the last statement was SELECT.
-                # (Complex to detect in executescript, so we just return success message)
-                output_log += "(Note: SQLite script ran. If you need results, ensure your frontend handles 'select' logic or splits queries)"
-            except Exception as e:
-                output_log = f"SQLite Error: {e}"
-            finally:
-                con.close()
+            # Run sqlite3 against an in-memory DB
+            # We pass the SQL content via stdin
+            cmd = ["sqlite3", ":memory:"]
+            output = run_command(cmd, input_text=sql_code)
 
-        # 2. POSTGRESQL
         elif db_type in ['postgres', 'postgresql']:
-            # Connect using the 'playground' user we created in start.sh
-            # No password needed because of 'trust' in pg_hba.conf
-            con = psycopg2.connect(
-                host="localhost",
-                user="playground",
-                database="playground"
-            )
-            con.autocommit = True
-            cur = con.cursor()
-            try:
-                # Psycopg2 execute cannot run multiple statements easily separated by ;
-                # unless we use simple execute. 
-                cur.execute(sql_code)
-                
-                if cur.description:
-                    rows = cur.fetchall()
-                    colnames = [desc[0] for desc in cur.description]
-                    output_log = f"Columns: {colnames}\n"
-                    for row in rows:
-                        output_log += f"{row}\n"
-                else:
-                    output_log = f"Command executed. Rows affected: {cur.rowcount}"
-            except Exception as e:
-                output_log = f"PostgreSQL Error: {e}"
-            finally:
-                con.close()
+            # Run psql using the 'playground' user defined in start.sh
+            # -f runs the file provided
+            cmd = ["psql", "-h", "localhost", "-U", "playground", "-d", "playground", "-f", temp_sql_path]
+            output = run_command(cmd)
 
-        # 3. MYSQL / MARIADB
         elif db_type == 'mysql':
-            # Connect using the 'playground' user from start.sh
-            con = pymysql.connect(
-                host="localhost",
-                user="playground",
-                database="playground",
-                cursorclass=pymysql.cursors.DictCursor,
-                client_flag=pymysql.constants.CLIENT.MULTI_STATEMENTS
-            )
-            try:
-                with con.cursor() as cur:
-                    cur.execute(sql_code)
-                    # Fetch results from all statements
-                    while True:
-                        if cur.description:
-                            rows = cur.fetchall()
-                            output_log += f"--- Result ---\n{rows}\n"
-                        else:
-                            output_log += f"Affected rows: {cur.rowcount}\n"
-                        if not cur.nextset():
-                            break
-            except Exception as e:
-                output_log = f"MySQL Error: {e}"
-            finally:
-                con.close()
+            # Run mysql using the 'playground' user defined in start.sh
+            # -t forces ASCII table output even in non-interactive mode
+            cmd = ["mysql", "-h", "localhost", "-u", "playground", "playground", "-t"]
+            # Pass SQL via stdin
+            output = run_command(cmd, input_text=sql_code)
 
         else:
-            return jsonify({'output': 'Error: Unsupported database type'}), 400
+            return jsonify({'output': 'Unsupported database type. Use sqlite, postgres, or mysql.'}), 400
 
-        return jsonify({'output': output_log})
+        return jsonify({'output': output})
 
-    except Exception as e:
-        return jsonify({'output': f"System Error: {str(e)}"})
+    finally:
+        if temp_sql_path and os.path.exists(temp_sql_path):
+            os.remove(temp_sql_path)
 
 if __name__ == '__main__':
     # Hugging Face expects port 7860
-    app.run(host="0.0.0.0", port=7860)
-
-### How to use this API from your HTML/Frontend
-
-**1. For Python (Existing)**
-URL: `https://your-space-url.hf.space/api/run-python`
-Body: `{"code": "print(1+1)"}`
-
-**2. For SQL (New)**
-URL: `https://your-space-url.hf.space/api/run-sql`
-Body:
-```json
-{
-  "type": "mysql",
-  "code": "CREATE TABLE test (id INT); INSERT INTO test VALUES (1); SELECT * FROM test;"
-}
+    app.run(host="0.0.0.0", port=7860)