playwrite/db/sqlserver_connection.py

"""
SQL Server 数据库连接组件

提供 SQL Server 数据库连接和查询接口
"""

import pyodbc
from typing import List, Dict, Any, Optional
from db.base_connection import BaseDatabaseConnection


class SQLServerConnection(BaseDatabaseConnection):
    """SQL Server 数据库连接类"""

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """
        初始化数据库连接

        Args:
            config: 数据库配置字典
                - server: 服务器地址
                - database: 数据库名称
                - username: 用户名
                - password: 密码
                - driver: ODBC 驱动名称
                - trust_server_certificate: 是否信任服务器证书
        """
        super().__init__(config)

    def connect(self) -> pyodbc.Connection:
        """
        建立数据库连接

        Returns:
            pyodbc.Connection: 数据库连接对象
        """
        if self.connection is not None:
            return self.connection

        # 构建连接字符串
        driver = self.config.get('driver', 'ODBC Driver 18 for SQL Server')
        conn_str = (
            f"DRIVER={{{driver}}};"
            f"SERVER={self.config['server']};"
            f"DATABASE={self.config['database']};"
            f"UID={self.config['username']};"
            f"PWD={self.config['password']};"
            f"TrustServerCertificate={self.config.get('trust_server_certificate', 'yes')};"
        )

        try:
            self.connection = pyodbc.connect(conn_str)
            print(
                f"成功连接到 SQL Server 数据库: {self.config['server']}/{self.config['database']}"
            )
            return self.connection
        except pyodbc.Error as e:
            print(f"SQL Server 数据库连接失败: {e}")
            raise

    def disconnect(self):
        """关闭数据库连接"""
        if self.connection:
            self.connection.close()
            self.connection = None
            print("SQL Server 数据库连接已关闭")

    def execute_query(
        self, sql: str, params: Optional[tuple] = None
    ) -> List[Dict[str, Any]]:
        """
        执行查询语句并返回结果

        Args:
            sql: SQL 查询语句
            params: 查询参数（可选）

        Returns:
            List[Dict[str, Any]]: 查询结果列表，每个元素为一行数据的字典
        """
        if not self.connection:
            self.connect()

        cursor = self.connection.cursor()

        try:
            if params:
                cursor.execute(sql, params)
            else:
                cursor.execute(sql)

            # 获取列名
            columns = [column[0] for column in cursor.description]

            # 将结果转换为字典列表
            results = []
            for row in cursor.fetchall():
                results.append(dict(zip(columns, row)))

            return results

        except pyodbc.Error as e:
            print(f"查询执行失败: {e}")
            raise
        finally:
            cursor.close()

    def execute_update(self, sql: str, params: Optional[tuple] = None) -> int:
        """
        执行更新/插入/删除语句

        Args:
            sql: SQL 语句
            params: 参数（可选）

        Returns:
            int: 受影响的行数
        """
        if not self.connection:
            self.connect()

        cursor = self.connection.cursor()

        try:
            if params:
                cursor.execute(sql, params)
            else:
                cursor.execute(sql)

            self.connection.commit()
            return cursor.rowcount

        except pyodbc.Error as e:
            self.connection.rollback()
            print(f"执行失败，已回滚: {e}")
            raise
        finally:
            cursor.close()

    def get_placeholder(self) -> str:
        """
        获取参数占位符

        Returns:
            SQL Server 使用 "?" 作为参数占位符
        """
        return "?"