ylt_diy/lib/db.py

"""
数据库连接配置
"""
import os
import pymysql
from contextlib import contextmanager
from lib.logger import log_info, log_error, log_warning

def get_db_config():
    """从环境变量或默认值获取数据库配置"""
    return {
        'host': os.environ.get('DB_HOST', 'haoslm2.xicp.net'),
        'port': int(os.environ.get('DB_PORT', 10216)),
        'user': os.environ.get('DB_USER', 'root'),
        'password': os.environ.get('DB_PASSWORD', 'DsideaL147258369'),
        'database': os.environ.get('DB_NAME', 'yltcharge'),
        'charset': 'utf8mb4',
        'cursorclass': pymysql.cursors.DictCursor,
        'autocommit': True,
        'connect_timeout': 30,
        'read_timeout': 120,
        'write_timeout': 120
    }

DB_CONFIG = get_db_config()


@contextmanager
def get_connection():
    """获取数据库连接（上下文管理器）"""
    try:
        conn = pymysql.connect(**DB_CONFIG)
        yield conn
    except pymysql.Error as e:
        log_error(f'数据库连接失败: {e}', 'db', exc_info=True)
        raise
    finally:
        try:
            conn.close()
        except:
            pass


def execute_query(sql, params=None, retry=2):
    """执行查询并返回结果"""
    for attempt in range(retry + 1):
        try:
            with get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(sql, params)
                    result = cursor.fetchall()
                    if attempt > 0:
                        log_info(f'[数据库] 查询重试成功，第{attempt+1}次尝试', 'db')
                    return result
        except pymysql.err.OperationalError as e:
            if attempt < retry and (e.args[0] == 2013 or e.args[0] == 2006):
                log_warning(f'[数据库] 查询连接断开，正在重试（第{attempt+1}次）: {e}', 'db')
                import time
                time.sleep(1)
                continue
            log_error(f'查询失败: {e}\nSQL: {sql}\nParams: {params}', 'db', exc_info=True)
            raise
        except Exception as e:
            log_error(f'查询失败: {e}\nSQL: {sql}\nParams: {params}', 'db', exc_info=True)
            raise


def execute_update(sql, params=None, retry=2):
    """执行更新操作"""
    for attempt in range(retry + 1):
        try:
            with get_connection() as conn:
                with conn.cursor() as cursor:
                    affected = cursor.execute(sql, params)
                    if attempt > 0:
                        log_info(f'[数据库] 更新重试成功，第{attempt+1}次尝试', 'db')
                    return cursor.rowcount
        except pymysql.err.OperationalError as e:
            if attempt < retry and (e.args[0] == 2013 or e.args[0] == 2006):
                log_warning(f'[数据库] 更新连接断开，正在重试（第{attempt+1}次）: {e}', 'db')
                import time
                time.sleep(1)
                continue
            log_error(f'更新失败: {e}\nSQL: {sql}\nParams: {params}', 'db', exc_info=True)
            raise
        except Exception as e:
            log_error(f'更新失败: {e}\nSQL: {sql}\nParams: {params}', 'db', exc_info=True)
            raise


def execute_insert(sql, params=None, retry=2):
    """执行插入操作并返回插入ID"""
    for attempt in range(retry + 1):
        try:
            with get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(sql, params)
                    last_id = cursor.lastrowid
                    if attempt > 0:
                        log_info(f'[数据库] 插入重试成功，第{attempt+1}次尝试', 'db')
                    return last_id
        except pymysql.err.OperationalError as e:
            if attempt < retry and (e.args[0] == 2013 or e.args[0] == 2006):
                log_warning(f'[数据库] 插入连接断开，正在重试（第{attempt+1}次）: {e}', 'db')
                import time
                time.sleep(1)
                continue
            log_error(f'插入失败: {e}\nSQL: {sql}\nParams: {params}', 'db', exc_info=True)
            raise
        except Exception as e:
            log_error(f'插入失败: {e}\nSQL: {sql}\nParams: {params}', 'db', exc_info=True)
            raise