RAG/core/database_data.py

import mysql.connector
from mysql.connector import Error


# 数据库配置
db_config = {
    'host': 'rm-bp13g3ra2f59q49xs.mysql.rds.aliyuncs.com',  # 数据库主机
    'database': 'ai_knowledge',  # 数据库名称
    'user': 'wqsd',  # 用户名
    'password': 'wqsd@2025'  # 密码
}


class DatabaseHelper:
    def __init__(self, config=None):
        """ 初始化数据库连接配置 """
        if config is None:
            config = db_config
        self.config = config
        self.connection = None
        self.cursor = None
        self._connect()

    def _connect(self):
        """ 建立数据库连接 """
        try:
            self.connection = mysql.connector.connect(**self.config)
            if self.connection.is_connected():
                print("数据库连接成功")
                self.cursor = self.connection.cursor(dictionary=True)
        except Error as e:
            print(f"数据库连接失败: {e}")

    def execute_query(self, query, params=None):
        """
        执行查询SQL语句
        :param query: SQL 查询语句
        :param params: 查询参数 (可选)
        :return: 查询结果
        """
        try:
            self.cursor.execute(query, params)
            return self.cursor.fetchall()
        except Error as e:
            print(f"执行查询失败: {e}")
            return None

    def execute_non_query(self, query, params=None):
        """
        执行没有返回结果的 SQL 语句 (例如 INSERT, UPDATE, DELETE)
        :param query: SQL 执行语句
        :param params: 参数 (可选)
        :return: 受影响的行数
        """
        try:
            self.cursor.execute(query, params)
            self.connection.commit()  # 提交事务
            return self.cursor.rowcount  # 返回受影响的行数
        except Error as e:
            print(f"执行非查询失败: {e}")
            self.connection.rollback()
            return 0

    def close(self):
        """ 关闭数据库连接 """
        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()
            print("数据库连接已关闭")