代码示例

import json

import boto3

import pymysql

import logging

 

# 配置日志

logger = logging.getLogger()

logger.setLevel(logging.INFO)

 

# 初始化 Bedrock 客户端

bedrock_runtime = boto3.client('bedrock-runtime', region_name='us-east-1') # 请替换为你的区域

 

# 初始化 RDS 连接参数（从环境变量中获取是安全最佳实践）

db_host = 'your-rds-endpoint.rds.amazonaws.com'

db_name = 'your_database_name'

db_user = 'your_username'

db_password = 'your_password' # 强烈建议使用AWS Secrets Manager来管理密码！

 

def lambda_handler(event, context):

    # 1. 从请求中获取自然语言问题

    # 假设通过API Gateway调用，event['body']包含JSON

    body = json.loads(event.get('body', '{}'))

    natural_language_query = body.get('query', '')

    

    if not natural_language_query:

        return {

            'statusCode': 400,

            'body': json.dumps({'error': 'Missing "query" in request body.'})

        }

 

    logger.info(f"Natural Language Query: {natural_language_query}")

 

    # 2. 构建给Bedrock的提示词（Prompt），这是最关键的一步！

    # 你需要清晰地告诉模型你的数据库结构（Schema）和任务

    prompt = f"""

Human: 你是一名高级SQL专家。请根据以下数据库表结构，将用户的自然语言问题转换为一条安全、准确、高效的MySQL查询语句。

 

### 数据库表结构 (Schema):

* Table: `employees`

    * `id` (INT, PRIMARY KEY)

    * `first_name` (VARCHAR)

    * `last_name` (VARCHAR)

    * `department` (VARCHAR)

    * `salary` (DECIMAL)

    * `hire_date` (DATE)

 

### 用户问题:

{natural_language_query}

 

### 要求:

1. 只输出SQL语句，不要输出任何其他解释或文本。

2. 必须使用合法的MySQL语法。

3. 只能查询（SELECT）数据，不能修改（INSERT/UPDATE/DELETE）数据。

4. 确保语句效率，不需要关联时不要使用JOIN。

 

Assistant: 以下是生成的SQL查询语句：

"""

    try:

        # 3. 调用Bedrock模型（例如Claude 3 Sonnet）

        model_id = 'anthropic.claude-3-sonnet-20240229-v1:0' # 选择一个模型

        

        request_body = {

            "anthropic_version": "bedrock-2023-05-31",

            "max_tokens": 1024,

            "messages": [{"role": "user", "content": prompt}]

        }

        

        response = bedrock_runtime.invoke_model(

            modelId=model_id,

            body=json.dumps(request_body)

        )

        

        response_body = json.loads(response['body'].read().decode('utf-8'))

        generated_sql = response_body['content'][0]['text'].strip()

        

        # 清理SQL输出，移除可能存在的代码标记

        if generated_sql.startswith('```sql'):

            generated_sql = generated_sql.split('```sql')[1].split('```')[0].strip()

        elif generated_sql.startswith('```'):

            generated_sql = generated_sql.split('```')[1].split('```')[0].strip()

        

        logger.info(f"Generated SQL: {generated_sql}")

 

        # 4. (关键安全步骤) 进行SQL校验

        # 这里可以添加自定义规则，例如禁止某些危险操作

        if not generated_sql.upper().startswith('SELECT'):

            raise Exception("Generated SQL is not a SELECT statement. Operation not permitted.")

        # 可以添加更多黑名单关键字检查，如 'DROP', 'DELETE', 'INSERT', 'UPDATE' 等

 

        # 5. 连接RDS并执行生成的SQL

        connection = pymysql.connect(

            host=db_host,

            user=db_user,

            password=db_password,

            database=db_name,

            connect_timeout=5

        )

        

        with connection.cursor(pymysql.cursors.DictCursor) as cursor:

            cursor.execute(generated_sql)

            results = cursor.fetchall()

        

        connection.close()

 

        # 6. 返回查询结果给用户

        return {

            'statusCode': 200,

            'headers': {

                'Content-Type': 'application/json'

            },

            'body': json.dumps({

                'generated_sql': generated_sql, # 返回生成的SQL用于调试和透明化

                'results': results

            })

        }

        

    except pymysql.MySQLError as e:

        logger.error(f"Database error: {e}")

        return {

            'statusCode': 500,

            'body': json.dumps({'error': 'Database error occurred.', 'details': str(e)})

        }

    except Exception as e:

        logger.error(f"Unexpected error: {e}")

        return {

            'statusCode': 500,

            'body': json.dumps({'error': 'Failed to process query.', 'details': str(e)})

        }