enterprise-kb-rag/3.3 基于Mysql库的问答系统/integrated_qa_system/mysql_qa/db/mysql_client.py

# 导入 MySQL 连接库
import pymysql
# 导入 pandas
import pandas as pd
import sys
from pathlib import Path

# 直接运行本文件时，把项目根目录加入模块搜索路径。
if __package__ in (None, ""):
    project_root = Path(__file__).resolve().parents[2]
    sys.path.insert(0, str(project_root))

# 导入配置和日志
from base import Config, logger


class MySQLClient:
    def __init__(self):
        # 初始化日志
        self.logger = logger
        try:
            # 连接 MySQL 数据库
            self.connection = pymysql.connect(
                host=Config().MYSQL_HOST,
                user=Config().MYSQL_USER,
                password=Config().MYSQL_PASSWORD,
                database=Config().MYSQL_DATABASE,
                charset="utf8mb4",
            )
            # 创建游标
            self.cursor = self.connection.cursor()
            # 记录连接成功
            self.logger.info("MySQL 连接成功")
        except pymysql.MySQLError as exc:
            # 记录连接失败
            self.logger.error(f"MySQL 连接失败: {exc}")
            raise

    def create_table(self):
        create_table_query = """
        CREATE TABLE IF NOT EXISTS jpkb (
            id INT AUTO_INCREMENT PRIMARY KEY,
            subject_name VARCHAR(20),
            question VARCHAR(1000),
            answer VARCHAR(1000)
        )
        """
        try:
            self.cursor.execute(create_table_query)
            self.connection.commit()
            self.logger.info("表创建成功")
        except pymysql.MySQLError as exc:
            self.logger.error(f"表创建失败: {exc}")
            raise

    def insert_data(self, csv_path):
        try:
            data = pd.read_csv(csv_path)
            records = [
                (row["知识域"], row["问题"], row["答案"])
                for _, row in data.iterrows()
            ]
            insert_query = (
                "INSERT INTO jpkb (subject_name, question, answer) "
                "VALUES (%s, %s, %s)"
            )

            # 删除和重新插入处于同一事务中，失败时可恢复原有数据。
            self.cursor.execute("DELETE FROM jpkb")
            self.cursor.executemany(insert_query, records)
            self.connection.commit()
            self.logger.info(f"数据刷新成功，共插入 {len(records)} 条")
        except Exception as exc:
            self.logger.error(f"数据插入失败: {exc}")
            self.connection.rollback()
            raise

    def fetch_questions(self):
        # 获取所有问题
        try:
            self.cursor.execute("SELECT question FROM jpkb")
            results = self.cursor.fetchall()
            self.logger.info("成功获取问题")
            return results
        except pymysql.MySQLError as exc:
            self.logger.error(f"查询失败: {exc}")
            return []

    def fetch_answer(self, question):
        # 获取指定问题的答案
        try:
            self.cursor.execute(
                "SELECT answer FROM jpkb WHERE question=%s", (question,)
            )
            result = self.cursor.fetchone()
            return result[0] if result else None
        except pymysql.MySQLError as exc:
            self.logger.error(f"答案获取失败: {exc}")
            return None

    def close(self):
        # 关闭数据库连接
        try:
            self.cursor.close()
            self.connection.close()
            self.logger.info("MySQL 连接已关闭")
        except pymysql.MySQLError as exc:
            self.logger.error(f"关闭连接失败: {exc}")


if __name__ == "__main__":
    mysql_client = MySQLClient()
    try:
        mysql_client.create_table()
        data_file = Path(__file__).resolve().parents[1] / "data" / "JP企业知识库问答.csv"
        mysql_client.insert_data(data_file)
    finally:
        mysql_client.close()