Update app.py

app.py

@@ -1,9 +1,9 @@
 import os
-import gradio as gr
-import threading
-from flask import Flask, request
+from flask import Flask, request, jsonify
 from psycopg2 import connect, sql
 import psycopg2.extras
+import gradio as gr
+import threading
 
 # Fetch environment variables
 DB_NAME = os.getenv("DB_NAME")
@@ -23,152 +23,64 @@ def get_db_connection():
         port=DB_PORT
     )
 
-# Dictionary to manage server threads
-server_threads = {}
-
-# Flask app template for HTTP servers
-def create_server(name, port):
-    app = Flask(name)
-
-    @app.route("/")
-    def home():
-        return f"Server {name} is running on port {port}!"
-
-    return app
-
-# Function to start a server
-def start_server(name, port):
-    if name in server_threads:
-        return f"Server {name} is already running!"
-    
-    app = create_server(name, port)
-
-    def run():
-        app.run(host="0.0.0.0", port=port, debug=False, use_reloader=False)
-
-    thread = threading.Thread(target=run, daemon=True)
-    thread.start()
-    server_threads[name] = thread
-
-    # Update database status
-    conn = get_db_connection()
-    with conn.cursor() as cursor:
-        cursor.execute(
-            "INSERT INTO servers (name, port, status) VALUES (%s, %s, %s) ON CONFLICT (name) DO UPDATE SET status = EXCLUDED.status",
-            (name, port, "running")
-        )
-        conn.commit()
-    conn.close()
+# Create Flask app
+app = Flask(__name__)
 
-    return f"Server {name} started on port {port}!"
+# API endpoint for running SQL commands
+@app.route("/run_sql", methods=["POST"])
+def run_sql():
+    # Get the password and command from the request
+    password = request.json.get("password")
+    command = request.json.get("command")
 
-# Function to stop a server (dummy implementation for now)
-def stop_server(name):
-    if name not in server_threads:
-        return f"Server {name} is not running!"
-
-    # Terminate the thread (not directly supported, requires external handling)
-    server_threads.pop(name)
-    return f"Server {name} stopped!"
-
-# Gradio UI for managing servers
-def manage_servers(password, action, name, port):
+    # Check if the password is correct
     if password != APP_PASSWORD:
-        return "Invalid password!"
-    if action == "start":
-        return start_server(name, port)
-    elif action == "stop":
-        return stop_server(name)
-    else:
-        return "Invalid action!"
+        return jsonify({"error": "Invalid password!"}), 401
 
-# Function to handle HTTP commands
-def execute_db_command(password, command):
-    if password != APP_PASSWORD:
-        return "Invalid password!"
+    # Validate SQL command input
+    if not command:
+        return jsonify({"error": "No SQL command provided!"}), 400
 
-    # Establish DB connection
-    conn = get_db_connection()
-    cursor = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
-    
+    # Execute the SQL command
+    conn = None
+    result = None
     try:
-        # Execute the command (SQL injection protection by using parameterized queries)
+        # Connect to the database
+        conn = get_db_connection()
+        cursor = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
+
+        # Execute the SQL command (this uses parameterized queries for security)
         cursor.execute(command)
-        result = cursor.fetchall()
-        conn.commit()
+        if command.strip().lower().startswith("select"):
+            result = cursor.fetchall()  # If the command is a SELECT, fetch the results
+        else:
+            conn.commit()  # For non-SELECT commands, commit the changes
+            result = {"message": "Command executed successfully!"}
     except Exception as e:
-        result = f"Error executing command: {e}"
+        result = {"error": f"Error executing command: {str(e)}"}
     finally:
-        conn.close()
-    
-    return result
-
-# Gradio UI for managing servers and interacting with DB
-def get_db_stats():
-    conn = get_db_connection()
-    cursor = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
-    
-    # CockroachDB query to get database size (total used space in bytes)
-    cursor.execute("""
-    SELECT 
-        db_name,
-        SUM(range_size) AS total_size_bytes
-    FROM crdb_internal.database_statistics
-    WHERE db_name = current_database()
-    GROUP BY db_name;
-    """)
-    db_stats = cursor.fetchone()
+        if conn:
+            conn.close()
 
-    conn.close()
-    if db_stats:
-        return f"Database {db_stats['db_name']} size: {db_stats['total_size_bytes']} bytes"
-    else:
-        return "Unable to fetch database size"
+    return jsonify(result)
 
-# Gradio Interface
-with gr.Blocks() as server_manager:
-    gr.Markdown("# ZeroCPU Server Manager")
+# Define Gradio interface
+def gradio_interface():
+    def greet(name):
+        return f"Hello {name}!"
 
-    # Gradio UI for managing servers
-    with gr.Row():
-        password = gr.Textbox(label="Password", type="password")
-    with gr.Row():
-        action = gr.Radio(["start", "stop"], label="Action")
-        name = gr.Textbox(label="Server Name")
-        port = gr.Number(label="Port")
-    with gr.Row():
-        submit = gr.Button("Submit")
-    output = gr.Textbox(label="Output")
+    interface = gr.Interface(fn=greet, inputs="text", outputs="text")
+    interface.launch(share=True, inline=True)
 
-    # Submit action for server management
-    submit.click(
-        manage_servers,
-        inputs=[password, action, name, port],
-        outputs=output
-    )
-
-    # Gradio UI for viewing DB stats and executing commands
-    with gr.Row():
-        command_input = gr.Textbox(label="Database Command", placeholder="Enter SQL command")
-        command_submit = gr.Button("Execute Command")
-    command_output = gr.Textbox(label="Command Output")
-
-    # Submit action for executing DB commands
-    command_submit.click(
-        execute_db_command,
-        inputs=[password, command_input],
-        outputs=command_output
-    )
-
-    # Show the database size
-    with gr.Row():
-        db_stats_output = gr.Textbox(label="Database Stats")
-    gr.Button("Show Database Stats").click(
-        get_db_stats,
-        outputs=db_stats_output
-    )
+# Starting Flask app in a separate thread
+def run_flask():
+    app.run(host="0.0.0.0", port=5050)
 
-    # Display message when Gradio app is running
-    gr.Markdown("ZeroCPU is active")
+# Run Flask and Gradio in parallel using threading
+if __name__ == "__main__":
+    flask_thread = threading.Thread(target=run_flask)
+    flask_thread.daemon = True  # Make sure it will close when the main thread exits
+    flask_thread.start()
 
-server_manager.launch()
+    # Run Gradio in the main thread
+    gradio_interface()