Update app.py

app.py

@@ -1,8 +1,8 @@
 import os
 import gradio as gr
 import threading
-from flask import Flask, jsonify
-from psycopg2 import connect
+from flask import Flask, request
+from psycopg2 import connect, sql
 import psycopg2.extras
 
 # Fetch environment variables
@@ -13,7 +13,7 @@ DB_HOST = os.getenv("DB_HOST")
 DB_PORT = os.getenv("DB_PORT")
 APP_PASSWORD = os.getenv("APP_PASSWORD")
 
-# Database connection
+# Database connection function
 def get_db_connection():
     return connect(
         dbname=DB_NAME,
@@ -23,85 +23,10 @@ def get_db_connection():
         port=DB_PORT
     )
 
-# Database operation functions
-def get_db_stats():
-    conn = get_db_connection()
-    cursor = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
-    
-    # Example: Get database size (this is a placeholder query, you can adapt it to your use case)
-    cursor.execute("SELECT pg_size_pretty(pg_database_size(current_database())) AS size")
-    db_stats = cursor.fetchone()
-    
-    conn.close()
-    return db_stats['size']
-
-def create_table(name):
-    conn = get_db_connection()
-    cursor = conn.cursor()
-
-    # Create a basic table (customize the structure as needed)
-    cursor.execute(f"""
-    CREATE TABLE IF NOT EXISTS {name} (
-        id SERIAL PRIMARY KEY,
-        data TEXT
-    );
-    """)
-    conn.commit()
-    conn.close()
-    return f"Table {name} created successfully!"
-
-def delete_table(name):
-    conn = get_db_connection()
-    cursor = conn.cursor()
-
-    # Delete table
-    cursor.execute(f"DROP TABLE IF EXISTS {name};")
-    conn.commit()
-    conn.close()
-    return f"Table {name} deleted successfully!"
-
-# Flask app for HTTP access to the database
-def create_http_server():
-    app = Flask(__name__)
-
-    @app.route("/db-stats", methods=["GET"])
-    def db_stats():
-        db_stats = get_db_stats()
-        return jsonify({"db_size": db_stats})
-
-    @app.route("/create-table/<name>", methods=["POST"])
-    def create_table_route(name):
-        result = create_table(name)
-        return jsonify({"message": result})
-
-    @app.route("/delete-table/<name>", methods=["DELETE"])
-    def delete_table_route(name):
-        result = delete_table(name)
-        return jsonify({"message": result})
-
-    return app
-
-# Manage server and database operations in the Gradio UI
-def manage_servers(password, action, name=None, port=None):
-    if password != APP_PASSWORD:
-        return "Invalid password!"
-
-    if action == "start":
-        return start_server(name, port)
-    elif action == "stop":
-        return stop_server(name)
-    elif action == "db-stats":
-        return get_db_stats()
-    elif action == "create-table" and name:
-        return create_table(name)
-    elif action == "delete-table" and name:
-        return delete_table(name)
-    else:
-        return "Invalid action!"
-
-# Start and stop server functions (same as before)
+# Dictionary to manage server threads
 server_threads = {}
 
+# Flask app template for HTTP servers
 def create_server(name, port):
     app = Flask(name)
 
@@ -111,58 +36,139 @@ def create_server(name, 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)
     
+    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()
+
     return f"Server {name} started on port {port}!"
 
+# 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
-with gr.Blocks() as server_manager:
-    gr.Markdown("# Server Manager")
+# Gradio UI for managing servers
+def manage_servers(password, action, name, port):
+    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!"
+
+# Function to handle HTTP commands
+def execute_db_command(password, command):
+    if password != APP_PASSWORD:
+        return "Invalid password!"
+
+    # Establish DB connection
+    conn = get_db_connection()
+    cursor = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
     
+    try:
+        # Execute the command (SQL injection protection by using parameterized queries)
+        cursor.execute(command)
+        result = cursor.fetchall()
+        conn.commit()
+    except Exception as e:
+        result = f"Error executing command: {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()
+
+    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"
+
+# Gradio Interface
+with gr.Blocks() as server_manager:
+    gr.Markdown("# ZeroCPU Server Manager")
+
+    # Gradio UI for managing servers
     with gr.Row():
         password = gr.Textbox(label="Password", type="password")
-    
     with gr.Row():
-        action = gr.Radio(["start", "stop", "db-stats", "create-table", "delete-table"], label="Action")
-        name = gr.Textbox(label="Server or Table Name")  # Removed optional=True
-        port = gr.Number(label="Port", visible=False)  # Making port field visible only when starting a server
-    
+        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")
-    
+
+    # Submit action for server management
     submit.click(
         manage_servers,
         inputs=[password, action, name, port],
         outputs=output
     )
 
-# Start HTTP server for database operations
-def start_http_server():
-    http_server = create_http_server()
-    http_thread = threading.Thread(target=http_server.run, kwargs={"host": "0.0.0.0", "port": 5000}, daemon=True)
-    http_thread.start()
+    # 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
+    )
 
-# Start everything
-if __name__ == "__main__":
-    start_http_server()
-    server_manager.launch()
+    # Display message when Gradio app is running
+    gr.Markdown("ZeroCPU is active")
 
+server_manager.launch()