ZeroCPU / app.py
wop's picture
Update app.py
c4db877 verified
raw
history blame
4.66 kB
import os
import gradio as gr
import threading
from flask import Flask, jsonify
from psycopg2 import connect
import psycopg2.extras
# Fetch environment variables
DB_NAME = os.getenv("DB_NAME")
DB_USER = os.getenv("DB_USER")
DB_PASSWORD = os.getenv("DB_PASSWORD")
DB_HOST = os.getenv("DB_HOST")
DB_PORT = os.getenv("DB_PORT")
APP_PASSWORD = os.getenv("APP_PASSWORD")
# Database connection
def get_db_connection():
return connect(
dbname=DB_NAME,
user=DB_USER,
password=DB_PASSWORD,
host=DB_HOST,
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)
server_threads = {}
def create_server(name, port):
app = Flask(name)
@app.route("/")
def home():
return f"Server {name} is running on port {port}!"
return app
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
return f"Server {name} started on port {port}!"
def stop_server(name):
if name not in server_threads:
return f"Server {name} is not running!"
server_threads.pop(name)
return f"Server {name} stopped!"
# Gradio UI
with gr.Blocks() as server_manager:
gr.Markdown("# Server Manager")
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
with gr.Row():
submit = gr.Button("Submit")
output = gr.Textbox(label="Output")
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()
# Start everything
if __name__ == "__main__":
start_http_server()
server_manager.launch()