qinglong/mysql_db.py

import os
import mysql
from mysql.connector import Error
from contextlib import contextmanager


class DBManager:
    def __init__(self):
        # 从环境变量获取数据库配置
        self.host = os.getenv('MYSQL_HOST', '100.102.238.76')
        self.port = int(os.getenv('MYSQL_PORT', '3307'))
        self.database = os.getenv('MYSQL_DATABASE', 'qinglong')
        self.user = os.getenv('MYSQL_USER', 'root')
        self.password = os.getenv('MYSQL_PASSWORD', 'root')
        self.connection = None

    def connect(self):
        """建立数据库连接"""
        try:
            self.connection = mysql.connector.connect(
                host=self.host,
                port=self.port,
                database=self.database,
                user=self.user,
                password=self.password,
                autocommit=False  # 手动控制事务
            )
            if self.connection.is_connected():
                print("成功连接到MySQL数据库")
        except Error as e:
            print(f"连接数据库时出错: {e}")
            self.connection = None

    @contextmanager
    def managed_cursor(self):
        """上下文管理器，自动管理连接和游标"""
        try:
            if self.connection is None or not self.connection.is_connected():
                self.connect()

            if self.connection is None:
                raise Exception("数据库连接失败")

            cursor = self.connection.cursor()
            yield cursor  # 在此处返回游标给调用方
        except Error as e:
            print(f"数据库操作失败: {e}")
            # 尝试重新连接
            self.connect()
            if self.connection and self.connection.is_connected():
                cursor = self.connection.cursor()
                yield cursor
            else:
                raise
        finally:
            try:
                if 'cursor' in locals() and cursor:
                    cursor.close()
            except:
                pass  # 忽略关闭游标时的异常