@mcp.tool()
def execute_sql(sql: str) -> Dict[str, Any]:
    """Execute arbitrary SQL. Returns rows for SELECT else rowCount."""
    if not sql or not isinstance(sql, str):
        return {"error": "Invalid SQL statement"}
    
    try:
        with get_db_connection() as conn:
            with conn.cursor() as cur:
                cur.execute(sql)
                
                # 判断是否为查询语句
                if sql.strip().upper().startswith(('SELECT', 'WITH', 'SHOW', 'DESCRIBE', 'EXPLAIN')):
                    rows = cur.fetchall()
                    return {"rows": rows, "rowCount": len(rows)}
                else:
                    # 对于INSERT/UPDATE/DELETE等语句
                    conn.commit()
                    return {"rowCount": cur.rowcount, "message": "Query executed successfully"}
                    
    except Exception as e:
        logger.error(f"Error executing SQL: {e}")
        return {"error": str(e)}