crossborder/crossborder/utils/db_helper.py

import pandas as pd
from sqlalchemy import create_engine, text
from crossborder.utils.crypto_utils import AESCryptor

from crossborder.utils.log import log

DB_CONFIG = {
    'host': '10.130.75.149',
    'port': 3307,
    'user': 'yto_crm',
    'password': 'ENC(Fl9g4899OmVYddM42Rt2fA==:sDy1QG/7bmx/iHo4xEOBGQ==)',
    'database': 'crm_uat',
    'charset': 'utf8mb4'
}

cryptor = AESCryptor("uat_ff419620e7047a3c372e2513c5a2b9a5")


def get_decrypted_password():
    encrypted_pass = DB_CONFIG['password']
    if encrypted_pass.startswith("ENC("):
        try:
            return cryptor.decrypt(encrypted_pass)
        except Exception as e:
            log.error(f"密码解密失败: {str(e)}")
            raise
    return encrypted_pass

class DBHelper:
    def __init__(self):
        db_config = DB_CONFIG.copy()
        db_config['password'] = get_decrypted_password()
        self.engine = create_engine(
            f'mysql+pymysql://{db_config["user"]}:{db_config["password"]}@{db_config["host"]}:{db_config["port"]}/{db_config["database"]}?charset={db_config["charset"]}',
            pool_size=5,
            max_overflow=10
        )

    def get_commodity_id(self, name):
        """获取商品编码对应的分类ID[1,3](@ref)"""
        with self.engine.connect() as conn:
            result = conn.execute(
                text("SELECT id FROM t_yujin_crossborder_prov_commodity_category WHERE commodity_name = :name"),
                {'name': name}
            ).fetchone()
            return result[0] if result else None

    def bulk_insert(self, df, table_name, conflict_columns=None, update_columns=None):
        """
        增强版批量插入（支持覆盖更新）
        :param df: 要插入的DataFrame
        :param table_name: 目标表名
        :param conflict_columns: 冲突检测字段列表
        :param update_columns: 需要更新的字段列表
        """
        if df.empty:
            log.info("空数据集，跳过插入")
            return

        # 生成带参数的SQL模板
        columns = ', '.join(df.columns)
        placeholders = ', '.join([f":{col}" for col in df.columns])
        sql = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"

        # 添加ON DUPLICATE KEY UPDATE（MySQL语法）
        if conflict_columns and update_columns:
            # 1. 处理用户指定的更新字段
            update_clauses = [f"{col}=VALUES({col})" for col in update_columns]

            # 2. 强制添加create_time=NOW()
            update_clauses.append("create_time = NOW()")  # 新增

            # 3. 合并所有更新条件
            update_set = ', '.join(update_clauses)
            sql += f" ON DUPLICATE KEY UPDATE {update_set}"

        # 转换数据为字典列表格式
        data = df.to_dict(orient='records')
        # print("data:", data)
        try:
            with self.engine.connect() as conn:
                # 显式开启事务
                with conn.begin():
                    # 使用text()包装SQL语句
                    stmt = text(sql)

                    # 批量执行
                    conn.execute(stmt, data)

                    log.info(f"成功插入/更新 {len(df)} 行到 {table_name}")
        except Exception as e:
            log.error(f"数据库操作失败: {str(e)}")
            raise

    def update_january_yoy(self, prov_name='福建省'):
        """
        更新指定省份1月份同比数据
        :param prov_name: 省份名称，默认为福建省
        """
        update_sql = text("""
                          UPDATE t_yujin_crossborder_prov_region_trade AS curr
                              INNER JOIN t_yujin_crossborder_prov_region_trade AS prev
                          ON curr.city_code = prev.city_code
                              AND prev.crossborder_year_month = DATE_FORMAT(
                              DATE_SUB(
                              STR_TO_DATE(CONCAT(curr.crossborder_year_month, '-01'), '%Y-%m-%d'),
                              INTERVAL 1 YEAR
                              ),
                              '%Y-01'
                              )
                              SET
                                curr.yoy_import_export = COALESCE (
                                    TRUNCATE((curr.monthly_total - prev.monthly_total) / NULLIF (prev.monthly_total, 0) * 100, 4),
                                    0.0000
                                ),
                                curr.yoy_import = COALESCE (
                                    TRUNCATE((curr.monthly_import - prev.monthly_import) / NULLIF (prev.monthly_import, 0) * 100, 4),
                                    0.0000
                                ),
                                curr.yoy_export = COALESCE (
                                    TRUNCATE((curr.monthly_export - prev.monthly_export) / NULLIF (prev.monthly_export, 0) * 100, 4),
                                    0.0000
                                )
                          WHERE
                              curr.prov_name = :prov_name
                            AND curr.crossborder_year_month LIKE '%-01'
                            AND curr.crossborder_year_month
                              > '2023-01'
                          """)

        try:
            with self.engine.begin() as conn:
                result = conn.execute(update_sql, {'prov_name': prov_name})
                log.info(f"Updated {result.rowcount} rows for {prov_name}")
                return result.rowcount

        except Exception as e:
            log.error(f"Update failed: {str(e)}")
            raise RuntimeError(f"同比数据更新失败: {str(e)}") from e

    def update_prov_yoy(self, prov_name):
        """
        完整更新山东省同比数据（包含新旧数据处理）
        """
        try:
            # 步骤1：清理旧数据
            cleared = self.clear_old_prov_yoy(prov_name)

            # 步骤2：计算新数据
            updated = self._update_prov_new_yoy(prov_name)

            log.info(f"{prov_name}同比处理完成 | 清零:{cleared} 更新:{updated}")
            return {'cleared': cleared, 'updated': updated}
        except Exception as e:
            log.error(f"{prov_name}数据处理失败", exc_info=True)
            raise

    def clear_old_prov_yoy(self, prov_name):
        """
        清理指定省份2024年前数据的同比指标
        """
        clear_sql = text("""
                         UPDATE t_yujin_crossborder_prov_region_trade
                         SET yoy_import_export = null,
                             yoy_export        = null,
                             yoy_import        = null
                         WHERE prov_name = :prov_name
                           AND crossborder_year_month < '2024-01'
                           AND (yoy_import_export != 0 
                   OR yoy_export != 0 
                   OR yoy_import != 0) -- 优化：仅更新非零记录
                         """)

        try:
            with self.engine.begin() as conn:
                result = conn.execute(clear_sql, {'prov_name': prov_name})
                log.info(f"{prov_name}旧数据清零记录数: {result.rowcount}")
                return result.rowcount
        except Exception as e:
            log.error(f"旧数据清零失败: {str(e)}")
            raise


    def _update_prov_new_yoy(self,prov_name):
        """
        更新2024年及之后的省份城市同比数据
        """
        update_sql = text("""
                          UPDATE t_yujin_crossborder_prov_region_trade AS curr
                              INNER JOIN t_yujin_crossborder_prov_region_trade AS prev
                          ON curr.city_code = prev.city_code
                              AND prev.crossborder_year_month = DATE_FORMAT(
                              DATE_SUB(
                              STR_TO_DATE(CONCAT(curr.crossborder_year_month, '-01'), '%Y-%m-%d'),
                              INTERVAL 1 YEAR
                              ),
                              '%Y-%m'
                              )
                              SET
                                curr.yoy_import_export = COALESCE (
                                    TRUNCATE((curr.monthly_total - prev.monthly_total) / NULLIF (prev.monthly_total, 0) * 100, 4),
                                    0.0000
                                ),
                                curr.yoy_import = COALESCE (
                                    TRUNCATE((curr.monthly_import - prev.monthly_import) / NULLIF (prev.monthly_import, 0) * 100, 4),
                                    0.0000
                                ),
                                curr.yoy_export = COALESCE (
                                    TRUNCATE((curr.monthly_export - prev.monthly_export) / NULLIF (prev.monthly_export, 0) * 100, 4),
                                    0.0000
                                )
                          WHERE
                              curr.prov_name = :prov_name
                            AND curr.crossborder_year_month >= '2024-01'
                            AND prev.monthly_total IS NOT NULL
                          """)

        with self.engine.begin() as conn:
            result = conn.execute(update_sql, {'prov_name': prov_name})
            log.info(f"{prov_name}新数据更新数: {result.rowcount}")
            return result.rowcount


    def query(self, sql, params=None, return_df=True):
        """
        执行带参数的SQL语句（支持批量插入/更新）

        :param sql: 参数化的SQL语句（如含%s、%s等）
        :param params_list: 参数列表，每个元素是一个tuple或dict（根据SQL风格而定）
        :return: 受影响行数
        """
        try:
            with self.engine.connect() as conn:
                if return_df:
                    # 使用pandas直接读取为DataFrame
                    result = pd.read_sql(sql, conn, params=params)
                    log.info(f"查询成功，返回 {len(result)} 条记录")
                    return result
                else:
                    # 返回原始结果
                    result = conn.execute(sql, params or {}).fetchall()
                    log.info(f"查询成功，返回 {len(result)} 条记录")
                    return result
        except Exception as e:
            log.error(f"查询失败: {str(e)}")
            raise


    def execute_sql_with_params(self, sql: str, params_list: list):
        """
        执行带参数的SQL语句（支持批量插入/更新）

        :param sql: 参数化的SQL语句（如含%s、%s等）
        :param params_list: 参数列表，每个元素是一个tuple或dict（根据SQL风格而定）
        :return: 受影响行数
        """
        try:
            with self.engine.connect() as conn:
                with conn.begin():
                    # 使用text()包装原始SQL
                    stmt = text(sql)

                    # 判断是否为多组参数（批量插入）
                    if isinstance(params_list[0], (list, tuple)):
                        result = conn.execute(stmt, params_list)
                    else:
                        result = conn.execute(stmt, params_list)

                    affected_rows = result.rowcount
                    log.info(f"成功执行SQL，受影响行数：{affected_rows}")
                    return affected_rows
        except Exception as e:
            log.error(f"SQL执行失败: {str(e)}")
            raise