tools.py

import pyodbc
from dotenv import load_dotenv
from openai import OpenAI
import os, json
import subprocess
import time
import pyautogui
import pygetwindow
from datetime import datetime
from pydantic_ai import Tool
import re
from typing import Union
import pyperclip
image_counter=1

load_dotenv()

# 获取数据库表结构的函数
import os
import pyodbc

def get_db_structure(db_name:str) -> dict:
    """
    获取当前连接的 SQL Server 数据库名称和结构信息。

    Returns:
        dict: {
            "database": str,       # 当前数据库名称
            "structure": str,      # 表及字段结构信息
            "views": str           # 视图及字段结构信息
        }
    """
    try:
        # 从环境变量中读取连接配置
        driver = os.getenv("DB_DRIVER", "ODBC Driver 17 for SQL Server")
        server = os.getenv("DB_SERVER")
        # database = os.getenv("DB_NAME")
        trusted = os.getenv("DB_TRUSTED_CONNECTION", "yes")

        # 根据传入的 db_name 动态构建连接字符串
        conn_str = f"Driver={driver};Server={server};Database={db_name};Trusted_Connection={trusted};"

        conn = pyodbc.connect(conn_str)
        cursor = conn.cursor()

        # 初始化表结构和视图结构
        structure = ""
        views = ""

        # 查询所有用户表
        cursor.execute("""
            SELECT TABLE_SCHEMA, TABLE_NAME
            FROM INFORMATION_SCHEMA.TABLES
            WHERE TABLE_TYPE = 'BASE TABLE'
            ORDER BY TABLE_SCHEMA, TABLE_NAME
        """)
        tables = cursor.fetchall()

        if not tables:
            return {
                "dbname": db_name,
                "structure": "当前数据库中未找到任何用户表。"
            }

        if not tables:
            structure = "当前数据库中未找到任何用户表。\n"

            # 遍历所有表，获取字段信息
        for schema, table_name in tables:
            structure += f"表名：{schema}.{table_name}\n字段："

            cursor.execute("""
                        SELECT COLUMN_NAME, DATA_TYPE
                        FROM INFORMATION_SCHEMA.COLUMNS
                        WHERE TABLE_NAME = ? AND TABLE_SCHEMA = ?
                    """, (table_name, schema))

            columns = cursor.fetchall()
            if not columns:
                structure += "（无字段信息）\n\n"
                continue

            fields = [f"{col.COLUMN_NAME} ({col.DATA_TYPE})" for col in columns]
            structure += "、".join(fields) + "\n\n"

            # 获取所有视图
        cursor.execute("""
                    SELECT TABLE_SCHEMA, TABLE_NAME
                    FROM INFORMATION_SCHEMA.VIEWS
                    ORDER BY TABLE_SCHEMA, TABLE_NAME
                """)
        views_data = cursor.fetchall()

        if not views_data:
            views = "当前数据库中未找到任何视图。\n"

        # 遍历所有视图，获取字段信息
        for schema, view_name in views_data:
            views += f"视图名：{schema}.{view_name}\n字段："

            cursor.execute("""
                        SELECT COLUMN_NAME, DATA_TYPE
                        FROM INFORMATION_SCHEMA.COLUMNS
                        WHERE TABLE_NAME = ? AND TABLE_SCHEMA = ?
                    """, (view_name, schema))

            columns = cursor.fetchall()
            if not columns:
                views += "（无字段信息）\n\n"
                continue

            fields = [f"{col.COLUMN_NAME} ({col.DATA_TYPE})" for col in columns]
            views += "、".join(fields) + "\n\n"

        cursor.close()
        conn.close()

        return {
            "db_name": db_name,
            "structure": structure.strip(),
            "views": views.strip()
        }

    except Exception as e:
        return {
            "db_name": None,
            "structure": f"获取数据库结构失败：{str(e)}",
            "views": ""
        }

log_file_path = "log.txt"

def log_print(message: str):
    # 获取当前时间戳
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    log_message = f"[{timestamp}] {message}\n"


    # 写入文件
    with open(log_file_path, "a", encoding="utf-8") as log_file:
        log_file.write(log_message)

@Tool
def execute_sql_tool(sql_json: dict, db_name: str, num: int) -> str:
    """
    执行 SQL Server 语句（支持执行前/后截图），并返回统一 JSON 结果。
    本函数用于执行由 LLM 生成的 SQL Server 操作任务。

    每次执行任务可包含：
        - before_img：执行前截图 SQL 列表；
        - sql：主执行 SQL（INSERT、UPDATE、DELETE、SELECT 等）；
        - after_img：执行后截图 SQL 列表。

    ---
    Args:
        sql_json (dict): SQL 执行任务描述对象。
            结构示例如下：
            {
                "before_img": [
                    "SELECT * FROM Stu_View;",
                    "SELECT * FROM Stu_View_1;"
                ],
                "sql": "INSERT INTO Stu_View VALUES (702030, '严正', '男', '1991-02-13');",
                "after_img": [
                    "SELECT * FROM Stu_View;",
                    "SELECT * FROM Stu_View_1;"
                ],
            }

        db_name (str): 数据库名称。
            - 确保 SQL 在正确数据库上下文中运行。
        num (int):用于命名截图的数字
    ---
    Returns:
        str: JSON 字符串（UTF-8 编码），格式如下：
        {
            "before_img": [
                "image/1_before_Stu_View.png",
                "image/2_before_Stu_View_1.png"
            ],
            "sql": "USE SchoolDB; INSERT INTO Stu_View VALUES (...);",
            "after_img": [
                "image/3_after_Stu_View.png",
                "image/4_after_Stu_View_1.png"
            ],
            "error": ""
        }

        字段说明：
        - sql：最终执行的完整 SQL 文本；
        - before_img / after_img：截图路径列表（执行前/后）；
        - error：错误信息（成功时为空字符串）。
    """

    log_print(f'执行sql的函数\n传进来的参数\nsql_json{sql_json}\ndb_name {db_name}\nnum {num}\n')
    server = os.getenv("DB_SERVER")
    driver = os.getenv("DB_DRIVER", "ODBC Driver 17 for SQL Server")
    trusted = os.getenv("DB_TRUSTED_CONNECTION", "yes")

    # conn_str = (
    #     f"Driver={{{driver}}};"
    #     f"Server={server};"
    #     f"Database={db_name};"
    #     f"Trusted_Connection={trusted};"
    # )
    # 根据传入的 db_name 动态构建连接字符串
    conn_str = f"Driver={driver};Server={server};Database={db_name};Trusted_Connection={trusted};"

    before_img_paths, after_img_paths = [], []
    main_sql = sql_json.get("sql", "").strip()
    before_sqls = sql_json.get("before_img", []) or []
    after_sqls = sql_json.get("after_img", []) or []
    db_name = db_name.upper()

    def run_and_maybe_screenshot(sql_list, prefix):
        log_print('INFO:进入截图了')
        log_print(f"in_{prefix}")

        """执行一组查询语句并截图"""
        results = []
        for i, q in enumerate(sql_list, 1):
            try:
                name=f"{prefix}_{i}"
                screenshot = execute_sql_tool_with_screenshot(f"USE {db_name};\n{q}", db_name,name,num)
                path = screenshot.get("screenshot_path", f"image/{num}_{prefix}_{i}.png")
                results.append(path)
            except Exception as e:
                log_print(f"[WARN] 截图执行失败: {q} - {e}")
        return results



    # before 部分截图
    if before_sqls:
        before_img_paths = run_and_maybe_screenshot(before_sqls, "before")

    # 如果 SQL 已经包含 USE db_name，就去掉 USE 部分
    if f"USE {db_name}" in main_sql.upper():
        main_sql = main_sql.replace(f"USE {db_name}", "").strip()

    # 首先执行 `USE db_name`
    use_sql = f"USE {db_name};"

    try:
        conn = pyodbc.connect(conn_str)
        cursor = conn.cursor()

        # 首先执行 `USE db_name`
        cursor.execute(f"USE {db_name};")

        # 执行主 SQL 语句
        cursor.execute(main_sql)

        # 提交操作（如果是 INSERT/UPDATE/DELETE 等）
        if not re.match(r"^\s*(SELECT|SHOW|DESC|PRAGMA)", main_sql, re.I):
            conn.commit()

        # 成功执行后，输出消息
        success = True
        error_msg = ""

    except Exception as e:
        success = False
        error_msg = str(e)
        log_print(f"[ERROR] 执行失败: {e}")
    finally:
        # 确保关闭数据库连接
        if cursor:
            cursor.close()
        if conn:
            conn.close()
    # after 部分截图
    if after_sqls:
        after_img_paths = run_and_maybe_screenshot(after_sqls, "after")

    result = {
        "before_img": before_img_paths,
        "sql": main_sql,
        "after_img": after_img_paths,
        "error": "" if success else f"SQL 执行失败（{error_msg}）"
    }
    log_print(f'执行sql的函数\nreturn 的结果\n{result}\n')

    return json.dumps(result, ensure_ascii=False)


def execute_sql_tool_with_screenshot(sql: str, db_name: str,file_name:str,num:int):

    global image_counter
    temp_sql_file = "temp_query.sql"
    base_title = "Microsoft SQL Server Management Studio"
    image_folder = "image"
    if not os.path.exists(image_folder):
        os.makedirs(image_folder)

    try:
        with open(temp_sql_file, "w", encoding="utf-8") as f:
            f.write(sql)

        all_windows = pygetwindow.getWindowsWithTitle(base_title)
        if not all_windows:
            return {"status": "error", "message": "未找到 SSMS 窗口，请确认已打开。"}

        ssms_window = all_windows[0]
        if ssms_window.isMinimized:
            ssms_window.restore()
        ssms_window.activate()
        ssms_window.maximize()
        time.sleep(1.2)

        pyautogui.hotkey('ctrl', 'a')
        time.sleep(0.3)
        pyperclip.copy(sql)
        pyautogui.hotkey('ctrl', 'v')
        time.sleep(0.5)
        pyautogui.press('f5')
        time.sleep(5)

        # x y 代表的是 截图的左上顶点 width 指横向距离 从屏幕左到右
        x, y, width, height = 400, 800, 1200, 800
        screenshot_filename = os.path.join(image_folder, f"{num}_{file_name}.png")
        screenshot = pyautogui.screenshot(region=(x, y, width, height))
        screenshot.save(screenshot_filename)
        image_counter += 1

        return {"screenshot_path": screenshot_filename, "before_img": ""}

    except Exception as e:
        return {"screenshot_path": "", "before_img": "", "error": str(e)}


if  __name__=="__main__":
    pass
    # execute_sql_tool_with_screenshot("SELECT * FROM student WHERE 专业名 = '计算机科学与技术';","studentcouse",'now')
    # log_print(get_db_structure('studentcourse'))
    # log_log_print(get_db_structure('goodsorder'))