import mysql.connector
from mysql.connector import Error
import config
from flask import g

def get_db_connection():
    """创建数据库连接"""
    try:
        conn = mysql.connector.connect(
            host=config.DB_HOST,
            user=config.DB_USER,
            password=config.DB_PASSWORD,
            database=config.DB_NAME
        )
        return conn
    except Error as e:
        print(f"数据库连接错误: {e}")
        return None

def close_connection(conn):
    """关闭数据库连接"""
    if conn:
        conn.close()

# 添加与routes/classify.py兼容的函数名称
def get_db():
    """获取数据库连接 (与classify.py兼容)"""
    if 'db' not in g:
        g.db = get_db_connection()
    return g.db

def close_db(e=None):
    """关闭数据库连接 (与classify.py兼容)"""
    db = g.pop('db', None)
    if db:
        close_connection(db)

def init_app(app):
    """在Flask应用中注册数据库清理函数"""
    app.teardown_appcontext(close_db)