让lark机器人查询数据库

安装依赖库

pip install fastapi uvicorn pandas openpyxl lark-oapi pymysql sqlalchemy requests
核心代码实现 (main.py)
python 复制代码
import os
import re
import json
import pandas as pd
from fastapi import FastAPI, Request, BackgroundTasks 
from sqlalchemy import create_engine
import lark_oapi as lark
from lark_oapi.api.im.v1 import *
from sqlalchemy import text 
import openpyxl
from openpyxl.styles import numbers
# ================= 配置区 =================
FEISHU_APP_ID = "xxx"       # 替换为你的飞书 App ID
FEISHU_APP_SECRET = "xxx"  # 替换为你的飞书 App Secret
BOT_NAME = "lark机器人" 

DATA_SOURCES = {
    "pre": "mysql+pymysql://lark_bot:passwd@172.31.0.1:4000/mysql",
    "pre-report": "mysql+pymysql://lark_bot:passwd@172.31.0.2:4000/mysql"
}
# ==========================================

app = FastAPI()

# 初始化飞书客户端
client = lark.Client.builder() \
    .app_id(FEISHU_APP_ID) \
    .app_secret(FEISHU_APP_SECRET) \
    .log_level(lark.LogLevel.INFO) \
    .build()

def clean_and_parse_message(text: str):
    """清洗 @ 信息并解析数据源和 SQL"""
    text = text.strip()
    
    # 1. 移除飞书内置的 @ 机器人占位符（例如 @_user_1）以及可能存在的普通 @机器人 文字
    # 支持移除形如 "@_user_1", "@机器人", "<at id=...>" 等开头的字符
    cleaned_text = re.sub(r'^@_user_\d+\s*', '', text)
    cleaned_text = re.sub(r'^@\S+\s*', '', cleaned_text)
    cleaned_text = cleaned_text.strip()
    
    lines = cleaned_text.split('\n')
    if not lines or "数据源:" not in lines[0]:
        raise ValueError("格式错误！请确保首行为: @机器人 数据源: 名字")
    
    # 2. 提取数据源名称
    data_source_name = lines[0].split('数据源:')[1].strip()
    
    # 3. 提取并切分 SQL
    sql_text = " ".join(lines[1:])
    sqls = [sql.strip() for sql in sql_text.split(';') if sql.strip()]
    
    return data_source_name, sqls

def execute_sql_to_excel(data_source: str, sqls: list, session_id: str):
    """执行SQL并生成Excel文件路径列表"""
    if data_source not in DATA_SOURCES:
        raise ValueError(f"找不到数据源: {data_source}")
    
    engine = create_engine(DATA_SOURCES[data_source])
    excel_files = []
    
    for idx, sql in enumerate(sqls):
        try:
            # 执行查询并转为 DataFrame
            df = pd.read_sql(text(sql), engine)
            print(f"SQL执行成功，获取到 {len(df)} 行数据")
            
            # 生成临时 Excel 文件
            filename = f"result_{session_id}_sql{idx+1}.xlsx"
            filepath = os.path.join("/tmp", filename) # Windows环境下可改为 "./"
            #df.to_excel(filepath, index=False, engine='openpyxl')
            # 1. 使用 ExcelWriter 并指定 openpyxl 引擎
            with pd.ExcelWriter(filepath, engine='openpyxl') as writer:
                df.to_excel(writer, index=False, sheet_name='Sheet1')
                
                # 获取 openpyxl 工作表对象
                workbook = writer.book
                worksheet = writer.sheets['Sheet1']
                
                # 2. 遍历数据列，如果发现整型、大数字或 ID 类型的列，在 Excel 单元格中将其数字格式设为文本("@")
                for col_idx, col_name in enumerate(df.columns, start=1):
                    # 获取该列的数据类型
                    col_type = df[col_name].dtype
                    
                    # 针对 int64、object 或其他可能产生长数字的列
                    if col_type == 'int64' or col_type == 'object':
                        # 遍历该列的所有数据行（排除第一行表头，所以从 row=2 开始）
                        for row_idx in range(2, worksheet.max_row + 1):
                            cell = worksheet.cell(row=row_idx, column=col_idx)
                            if cell.value is not None:
                                # 强制转换值为字符串，并设置单元格数字格式为"文本"
                                cell.value = str(cell.value)
                                cell.number_format = '@' 
                                
                    # 针对浮点数，设置数字格式为不使用科学计数法的普通小数形式
                    elif col_type == 'float64':
                        for row_idx in range(2, worksheet.max_row + 1):
                            cell = worksheet.cell(row=row_idx, column=col_idx)
                            if cell.value is not None:
                                # 例如：保留 4 位小数，不使用科学计数法
                                cell.number_format = '0.0000' 
            excel_files.append(filepath)
        except Exception as e:
            print(f"SQL执行失败: {sql}, 错误: {str(e)}")
            # 可以选择将错误信息记录到文本文件传给用户，此处跳过错误SQL
            
    return excel_files

def upload_and_send_files(message_id: str, files: list):
    """将生成的 Excel 上传至飞书并回复给用户"""
    for file_path in files:
        file_name = os.path.basename(file_path)
        
        # 1. 上传文件到飞书 (注意这里直接传递 file object 而不是 f.read() 的 bytes)
        with open(file_path, "rb") as f:
            upload_req = CreateFileRequest.builder() \
                .request_body(CreateFileRequestBody.builder()
                              .file_type("stream")  # stream 适用于 xlsx 等任意格式
                              .file_name(file_name)
                              .file(f)              # <--- 直接传文件流 f 
                              .build()) \
                .build()
            upload_resp = client.im.v1.file.create(upload_req)
        
        if not upload_resp.success():
            print(f"文件上传失败: {upload_resp.msg}")
            # 如果上传失败，最好发条消息告诉用户
            error_req = ReplyMessageRequest.builder() \
                .message_id(message_id) \
                .request_body(ReplyMessageRequestBody.builder()
                              .content(f'{{"text":"文件 {file_name} 上传失败"}}')
                              .msg_type("text")
                              .build()) \
                .build()
            client.im.v1.message.reply(error_req)
            continue
            
        file_key = upload_resp.data.file_key
        
        # 2. 回复文件消息给用户
        reply_req = ReplyMessageRequest.builder() \
            .message_id(message_id) \
            .request_body(ReplyMessageRequestBody.builder()
                          .content(f'{{"file_key":"{file_key}"}}')
                          .msg_type("file") # 确保类型是 file
                          .build()) \
            .build()
        client.im.v1.message.reply(reply_req)
        
        # 3. 清理临时文件
        if os.path.exists(file_path):
            os.remove(file_path)

@app.post("/webhook/feishu")
async def feishu_webhook(request: Request, background_tasks: BackgroundTasks): # 注入
    """接收飞书 Webhook 事件"""

    req_json = await request.json()
    
    # 飞书 URL 验证挑战
    if "challenge" in req_json:
        return {"challenge": req_json["challenge"]}
        
    # 处理消息事件
    if req_json.get("header", {}).get("event_type") == "im.message.receive_v1":
        # 立即把处理逻辑丢进后台任务
        background_tasks.add_task(handle_message_logic, req_json)
        
        # 立即给飞书服务器返回响应，防止它重试
        return {"status": "ok"}        
def handle_message_logic(req_json):
    event = req_json.get("event", {})
    message = event.get("message", {})
    message_id = message.get("message_id")
    chat_type = message.get("chat_type")  # "p2p" (单聊) 或 "group" (群聊)
    mentions = message.get("mentions", []) # 获取被 @ 的人员列表
    
    if chat_type == "group":
        is_bot_mentioned = False
        # 遍历所有被 @ 的人，检查机器人的名字是否在其中
        for mention in mentions:
            if mention.get("name") == BOT_NAME:
                is_bot_mentioned = True
                break
        
        # 如果机器人没有被 @，直接忽略（此时别人互相 @ 将不会触发机器人）
        if not is_bot_mentioned:
            return {"status": "ignored"}    
    content_str = message.get("content", "")
    
    try:
        content_dict = json.loads(content_str)
        text = content_dict.get("text", "")
        
        # 1. 解析消息
        data_source, sqls = clean_and_parse_message(text)
        
        # 2. 执行SQL并生成Excel
        excel_files = execute_sql_to_excel(data_source, sqls, message_id)
        
        # 3. 发送给用户
        if excel_files:
            upload_and_send_files(message_id, excel_files)
        else:
            # 告知用户无结果或执行失败
            reply_req = ReplyMessageRequest.builder() \
                .message_id(message_id) \
                .request_body(ReplyMessageRequestBody.builder()
                              .content('{"text":"SQL执行失败或没有返回任何数据"}')
                              .msg_type("text")
                              .build()) \
                .build()
            client.im.v1.message.reply(reply_req)
            
    except Exception as e:
        # 捕获格式错误等异常并通知用户
        error_req = ReplyMessageRequest.builder() \
            .message_id(message_id) \
            .request_body(ReplyMessageRequestBody.builder()
                          .content(f'{{"text":"处理出错: {str(e)}"}}')
                          .msg_type("text")
                          .build()) \
            .build()
        client.im.v1.message.reply(error_req)
3.运行中间件
uvicorn main:app --host 0.0.0.0 --port 8080 >uvicorn.log 2>&1 &
4.配置机器人 Webhook 订阅事件
将事件发送至开发者服务器的地址：http://公网ip:8080/webhook/feishu
事件订阅页面，点击添加事件，搜索并勾选：
接收消息 (im.message.receive_v1)
测试
@lark机器人数据源: pre
select * from log limit 3;
select * from log limit 9;