app.py

# app.py
import gradio as gr
from sql_agent import generate_sql
from db_executor import execute_sql
from schema_reader import get_schema_hint
import pandas as pd
import psycopg2

# 全局变量存储当前连接配置和schema
current_db_config = None
current_schema_hint = None

def test_db_connection(host, port, dbname, user, password):
    """测试数据库连接"""
    try:
        conn = psycopg2.connect(
            host=host,
            port=port,
            dbname=dbname,
            user=user,
            password=password
        )
        conn.close()
        return "✅ 数据库连接成功！"
    except Exception as e:
        return f"❌ 连接失败：{str(e)}"

def update_db_config(host, port, dbname, user, password):
    """更新数据库配置并获取schema"""
    global current_db_config, current_schema_hint
    
    try:
        # 测试连接
        conn = psycopg2.connect(
            host=host,
            port=port,
            dbname=dbname,
            user=user,
            password=password
        )
        conn.close()
        
        # 更新配置
        current_db_config = {
            "host": host,
            "port": port,
            "dbname": dbname,
            "user": user,
            "password": password
        }
        
        # 获取schema信息
        current_schema_hint = get_schema_hint(current_db_config)
        
        return "✅ 数据库配置已更新，schema已加载！", current_schema_hint
        
    except Exception as e:
        return f"❌ 配置更新失败：{str(e)}", ""

def natural_language_query(user_input):
    global current_db_config, current_schema_hint
    
    if not current_db_config:
        return "❌ 请先配置数据库连接！", None
        
    try:
        sql = generate_sql(user_input, schema_hint=current_schema_hint)
        if any(kw in sql.lower() for kw in ["drop", "delete", "truncate", "update", "insert"]):
            return f"⚠️ 检测到敏感 SQL，已阻止执行：\n{sql}", None
        result = execute_sql(sql, current_db_config)
        if isinstance(result, str):
            return result, None
        else:
            df = pd.DataFrame(result)
            return f"✅ SQL 已执行：\n{sql}", df
    except Exception as e:
        return f"❌ 出错：{e}", None

# 创建数据库配置界面
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 NaturalSQL MCP 数据库查询系统")
    
    with gr.Tab("数据库配置"):
        gr.Markdown("### 配置PostgreSQL数据库连接")
        
        with gr.Row():
            host_input = gr.Textbox(label="主机地址", value="localhost", placeholder="localhost")
            port_input = gr.Textbox(label="端口", value="5432", placeholder="5432")
        
        with gr.Row():
            dbname_input = gr.Textbox(label="数据库名", placeholder="your_database")
            user_input = gr.Textbox(label="用户名", placeholder="postgres")
        
        password_input = gr.Textbox(label="密码", type="password", placeholder="your_password")
        
        with gr.Row():
            test_btn = gr.Button("测试连接", variant="secondary")
            update_btn = gr.Button("更新配置", variant="primary")
        
        connection_status = gr.Textbox(label="连接状态", interactive=False)
        schema_display = gr.Textbox(label="数据库Schema", interactive=False, lines=10)
        
        # 绑定按钮事件
        test_btn.click(
            fn=test_db_connection,
            inputs=[host_input, port_input, dbname_input, user_input, password_input],
            outputs=connection_status
        )
        
        update_btn.click(
            fn=update_db_config,
            inputs=[host_input, port_input, dbname_input, user_input, password_input],
            outputs=[connection_status, schema_display]
        )
    
    with gr.Tab("SQL查询"):
        gr.Markdown("### 自然语言SQL查询")
        gr.Markdown("请先在'数据库配置'标签页中配置数据库连接，然后在此处进行查询。")
        
        query_input = gr.Textbox(
            lines=3, 
            label="自然语言查询", 
            placeholder="例如：查看2023年7月的订单总金额"
        )
        
        query_btn = gr.Button("执行查询", variant="primary")
        
        result_text = gr.Textbox(label="执行结果", interactive=False)
        result_table = gr.Dataframe(label="查询数据")
        
        # 绑定查询按钮事件
        query_btn.click(
            fn=natural_language_query,
            inputs=query_input,
            outputs=[result_text, result_table]
        )

demo.launch(mcp_server=True)