utils/db.py

"""
数据库工具模块 - 处理SQLite数据库连接和初始化
"""
import sqlite3
import os
from config import DATABASE_PATH

def get_db_connection():
    """
    获取数据库连接
    
    Returns:
        sqlite3.Connection: 数据库连接对象
    """
    conn = sqlite3.connect(DATABASE_PATH)
    conn.row_factory = sqlite3.Row  # 使查询结果可以通过列名访问
    return conn

def init_db():
    """
    初始化数据库，创建必要的表
    """
    # 确保数据目录存在
    os.makedirs(os.path.dirname(DATABASE_PATH), exist_ok=True)
    
    conn = get_db_connection()
    try:
        cursor = conn.cursor()
        
        # 创建API密钥表
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS api_keys (
            id TEXT PRIMARY KEY,
            platform TEXT NOT NULL,
            name TEXT NOT NULL,
            key TEXT NOT NULL,
            created_at TEXT NOT NULL,
            updated_at TEXT NOT NULL,
            success INTEGER DEFAULT 0,
            return_message TEXT DEFAULT '等待测试',
            states TEXT DEFAULT '',
            balance REAL DEFAULT 0
        )
        ''')
        
        conn.commit()
    except Exception as e:
        print(f"初始化数据库时出错: {str(e)}")
    finally:
        conn.close()

# 确保模块被导入时初始化数据库
init_db()