在亚马逊云科技环境上基于 Dify Agent 快速部署 text2SQL 智能数据分析助手

aws rds create-db-instance \
--db-instance-identifier mydb \
--db-instance-class db.m7g.large \
--engine mysql \
--engine-version 8.0.40 \
--allocated-storage 20 \
--master-username admin \
--master-user-password xxxxxxxx \
--port 3306 \
--db-name mydatabase \
--storage-type gp3 \
--publicly-accessible \
--backup-retention-period 7 \
--multi-az \
--tags Key=Environment,Value=Development

from starlette.applications import Starlette
from starlette.responses import JSONResponse
from starlette.routing import Route
from sqlalchemy.engine.url import URL
import pymysql
import asyncio
from sqlalchemy import create_engine, text
from typing import List, Dict, Any, Optional, Tuple
from pydantic import BaseModel, validator
import databases
from contextlib import contextmanager

# Data model for database structure information
class DatabaseStructure(BaseModel):
    """Model representing database table structure"""
    name: str 
    fields: List[Tuple[str, str]]
    preview_rows: Optional[List[Tuple]] = None


# Database connection handler
@contextmanager
def create_db_connection(server: str, username: str, pwd: str, database: str):
    """
    Creates and manages database connection
    
    Args:
        server: Database server address
        username: Database login name
        pwd: Database password
        database: Target database name
    """
    try:
        # Construct connection string
        connection_string = f"mysql+pymysql://{username}:{pwd}@{server}:3306/{database}"
        engine = create_engine(connection_string)
        connection = engine.connect()
        yield connection
    except Exception as error:
        raise RuntimeError(f"Database connection error: {str(error)}")
    finally:
        if 'connection' in locals():
            connection.close()


# Retrieve database schema information
async def fetch_schema(request):
    """Endpoint for retrieving database structure information"""
    params = request.query_params
    server = params.get('server')
    username = params.get('username')
    pwd = params.get('pwd')
    database = params.get('database')
    
    if not all([server, username, pwd, database]):
        return JSONResponse({"error": "Missing connection parameters"}, status_code=400)
    
    try:
        with create_db_connection(server, username, pwd, database) as conn:
            # Get all tables
            tables_query = text("SHOW TABLES")
            tables_result = conn.execute(tables_query)
            tables = [row[0] for row in tables_result]
            
            schema_info = []
            
            for table in tables:
                # Get columns information
                columns_query = text(f"DESCRIBE {table}")
                columns_result = conn.execute(columns_query)
                columns = [(row[0], row[1]) for row in columns_result]
                
                # Get sample data
                sample_query = text(f"SELECT * FROM {table} LIMIT 3")
                sample_result = conn.execute(sample_query)
                samples = [tuple(row) for row in sample_result]
                
                table_info = DatabaseStructure(
                    name=table,
                    fields=columns,
                    preview_rows=samples
                )
                schema_info.append(table_info.dict())
                
            return JSONResponse({"database_structure": schema_info})
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)


# Execute SQL query
async def run_sql_query(request):
    """Endpoint for running SQL queries against the database"""
    data = await request.json()
    sql_statement = data.get('sql_statement')
    server = data.get('server')
    username = data.get('username')
    pwd = data.get('pwd')
    database = data.get('database')
    
    if not all([sql_statement, server, username, pwd, database]):
        return JSONResponse({"error": "Missing required parameters"}, status_code=400)
    
    try:
        with create_db_connection(server, username, pwd, database) as conn:
            result = conn.execute(text(sql_statement))
            
            if result.returns_rows:
                columns = result.keys()
                data = [dict(zip(columns, row)) for row in result]
                return JSONResponse({"result": data})
            else:
                return JSONResponse({"message": "Query executed successfully"})
    except Exception:
        return JSONResponse({"error": "Query execution failed. Please check your SQL syntax."}, status_code=400)


# Routes configuration
routes = [
    Route('/database-schema', endpoint=fetch_schema, methods=["GET"]),
    Route('/query', endpoint=run_sql_query, methods=["POST"]),
]

# Initialize application
api = Starlette(routes=routes)

# Server startup
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(api, host="0.0.0.0", port=8080)

{
  "openapi": "3.0.0",
  "info": {
    "title": "SQL Query Execution API",
    "description": "API for executing SQL queries against databases",
    "version": "1.0.0"
  },
  "servers": [
    {
      "url": "https://api.example.com/v1",
      "description": "Main production API server"
    },
    {
      "url": "/",
      "description": "Relative server path"
    }
  ],
  "paths": {
    "/query": {
      "post": {
        "summary": "Execute SQL query",
        "description": "Executes a SQL statement against the specified database and returns results",
        "operationId": "executeSqlQuery",
        "requestBody": {
          "description": "Query parameters and database connection information",
          "required": true,
          "content": {
            "application/json": {
              "schema": {
                "type": "object",
                "required": ["sql_statement", "server", "username", "pwd", "database"],
                "properties": {
                  "sql_statement": {
                    "type": "string",
                    "description": "SQL query to execute"
                  },
                  "server": {
                    "type": "string",
                    "description": "Database server address"
                  },
                  "username": {
                    "type": "string",
                    "description": "Database username"
                  },
                  "pwd": {
                    "type": "string",
                    "description": "Database password"
                  },
                  "database": {
                    "type": "string",
                    "description": "Target database name"
                  }
                }
              },
              "examples": {
                "select": {
                  "summary": "Select query example",
                  "value": {
                    "sql_statement": "SELECT * FROM users LIMIT 10",
                    "server": "localhost",
                    "username": "dbuser",
                    "pwd": "password123",
                    "database": "mydb"
                  }
                }
              }
            }
          }
        },
        "responses": {
          "200": {
            "description": "Successful operation",
            "content": {
              "application/json": {
                "schema": {
                  "type": "object",
                  "oneOf": [
                    {
                      "properties": {
                        "result": {
                          "type": "array",
                          "description": "Query results (for SELECT statements)",
                          "items": {
                            "type": "object",
                            "additionalProperties": true
                          }
                        }
                      }
                    },
                    {
                      "properties": {
                        "message": {
                          "type": "string",
                          "description": "Success message (for non-SELECT statements)"
                        }
                      }
                    }
                  ]
                }
              }
            }
          },
          "400": {
            "description": "Bad request - invalid query or missing parameters",
            "content": {
              "application/json": {
                "schema": {
                  "type": "object",
                  "properties": {
                    "error": {
                      "type": "string"
                    }
                  }
                }
              }
            }
          }
        }
      }
    }
  }
}