abot/db/base.py

# -*- coding: utf-8 -*-
import time

from loguru import logger
from typing import List, Dict, Any, Optional, Tuple, Union

from db.connection import DBConnectionManager


class BaseDBOperator:
    """基础数据库操作类"""

    def __init__(self, db_manager: DBConnectionManager):
        self.db_manager = db_manager
        self.LOG = logger

    @staticmethod
    def _compact_sql(sql: str) -> str:
        """把 SQL 压成单行，便于日志里快速定位问题。"""
        return " ".join(str(sql or "").split())

    @classmethod
    def _truncate_text(cls, value, max_length: int = 240) -> str:
        """截断长文本，避免日志被超长 SQL 或参数刷屏。"""
        text = str(value or "")
        if len(text) <= max_length:
            return text
        return f"{text[:max_length]}..."

    def _log_sql_timing(self, operation: str, sql: str, params, elapsed_ms: float, affected_rows: Optional[int] = None) -> None:
        """记录慢 SQL 日志。

        设计说明：
        1. 只在超过阈值时输出 warning，避免日常日志噪声过大；
        2. 统一输出压缩后的 SQL 与截断参数，便于线上排查具体慢点；
        3. 查询/更新/批量/事务都走同一入口，后续如果要接后台审计也更容易扩展。
        """
        if not self.db_manager.is_slow_query_log_enabled():
            return

        threshold_ms = self.db_manager.get_slow_query_threshold_ms()
        if elapsed_ms < threshold_ms:
            return

        affected_text = ""
        if affected_rows is not None:
            affected_text = f" affected_rows={affected_rows}"
        self.LOG.warning(
            f"检测到慢SQL operation={operation} cost_ms={round(elapsed_ms, 2)} threshold_ms={threshold_ms}"
            f"{affected_text} sql={self._truncate_text(self._compact_sql(sql), 400)} "
            f"params={self._truncate_text(params, 240)}"
        )

    def execute_query(self, sql: str, params: Optional[tuple] = None, fetch_one: bool = False) -> Union[
        List[Dict], Dict, None]:
        """执行查询SQL"""
        conn = self.db_manager.get_mysql_connection()
        started_at = time.perf_counter()
        try:
            with conn.cursor(dictionary=True) as cursor:
                cursor.execute(sql, params or ())
                elapsed_ms = (time.perf_counter() - started_at) * 1000
                if fetch_one:
                    result = cursor.fetchone()
                    self._log_sql_timing("query_one", sql, params, elapsed_ms, 1 if result else 0)
                    return result
                result = cursor.fetchall()
                self._log_sql_timing("query", sql, params, elapsed_ms, len(result or []))
                return result
        except Exception as e:
            self.LOG.error(
                f"执行查询SQL出错: {e}, SQL: {sql}, 参数: {str(params)[:200] + '...' if len(str(params)) > 200 else params}"
            )
            return None
        finally:
            conn.close()

    def execute_update(self, sql: str, params: Optional[tuple] = None) -> bool:
        """执行更新SQL"""
        conn = self.db_manager.get_mysql_connection()
        started_at = time.perf_counter()
        try:
            with conn.cursor() as cursor:
                cursor.execute(sql, params or ())
                affected_rows = cursor.rowcount
            conn.commit()
            self._log_sql_timing("update", sql, params, (time.perf_counter() - started_at) * 1000, affected_rows)
            return True
        except Exception as e:
            self.LOG.error(
                f"执行更新SQL出错: {e}, SQL: {sql}, 参数: {str(params)[:200] + '...' if len(str(params)) > 200 else params}"
            )
            conn.rollback()
            return False
        finally:
            conn.close()

    def execute_batch(self, sql: str, params_list: List[tuple]) -> bool:
        """批量执行SQL"""
        if not params_list:
            return True

        conn = self.db_manager.get_mysql_connection()
        started_at = time.perf_counter()
        try:
            with conn.cursor() as cursor:
                cursor.executemany(sql, params_list)
                affected_rows = cursor.rowcount
            conn.commit()
            self._log_sql_timing(
                "batch_update",
                sql,
                f"params_count={len(params_list)}",
                (time.perf_counter() - started_at) * 1000,
                affected_rows,
            )
            return True
        except Exception as e:
            self.LOG.error(f"批量执行SQL出错: {e}, SQL: {sql}, 参数数量: {len(params_list)}")
            conn.rollback()
            return False
        finally:
            conn.close()

    def execute_transaction(self, operations: List[Tuple[str, tuple]]) -> bool:
        """执行事务"""
        if not operations:
            return True

        conn = self.db_manager.get_mysql_connection()
        started_at = time.perf_counter()
        try:
            with conn.cursor() as cursor:
                for sql, params in operations:
                    cursor.execute(sql, params)
            conn.commit()
            self._log_sql_timing(
                "transaction",
                f"{len(operations)} statements",
                f"operations={len(operations)}",
                (time.perf_counter() - started_at) * 1000,
            )
            return True
        except Exception as e:
            self.LOG.error(f"执行事务出错: {e}, 操作数量: {len(operations)}")
            conn.rollback()
            return False
        finally:
            conn.close()