""" SQLite 数据库适配器 基于 SQLite + aiosqlite 实现的数据库适配器，适用于 macOS 本地训练场景。 """ import json import sqlite3 import uuid from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Optional import aiosqlite from project_config import settings from ..base import DatabaseAdapter from ...models.domain import Task, TaskStatus # 阶段类型列表 STAGE_TYPES = [ "audio_slice", "asr", "text_feature", "hubert_feature", "semantic_token", "sovits_train", "gpt_train", ] class SQLiteAdapter(DatabaseAdapter): """ SQLite 数据库适配器 特点： 1. 使用 aiosqlite 实现异步数据库操作 2. 支持 Task (Quick Mode) 和 Experiment (Advanced Mode) 管理 3. 自动初始化数据库表结构 表结构： - tasks: Quick Mode 任务 - experiments: Advanced Mode 实验 - stages: 实验阶段状态 - files: 文件记录 Example: >>> adapter = SQLiteAdapter() >>> task = Task(id="task-123", exp_name="my_voice", config={}) >>> await adapter.create_task(task) >>> task = await adapter.get_task("task-123") """ def __init__(self, db_path: Optional[str] = None): """ 初始化 SQLite 适配器 Args: db_path: 数据库文件路径，默认使用 settings.SQLITE_PATH """ if db_path: self.db_path = db_path else: self.db_path = str(settings.SQLITE_PATH) # 确保目录存在 Path(self.db_path).parent.mkdir(parents=True, exist_ok=True) # 同步初始化数据库 self._init_db_sync() def _init_db_sync(self) -> None: """同步初始化数据库表结构""" with sqlite3.connect(self.db_path) as conn: # Tasks 表 (Quick Mode) conn.execute(''' CREATE TABLE IF NOT EXISTS tasks ( id TEXT PRIMARY KEY, job_id TEXT, exp_name TEXT NOT NULL, status TEXT NOT NULL DEFAULT 'queued', config TEXT, current_stage TEXT, progress REAL DEFAULT 0, stage_progress REAL DEFAULT 0, message TEXT, error_message TEXT, created_at TEXT NOT NULL, started_at TEXT, completed_at TEXT ) ''') conn.execute('CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status)') conn.execute('CREATE INDEX IF NOT EXISTS idx_tasks_created ON tasks(created_at)') # Experiments 表 (Advanced Mode) conn.execute(''' CREATE TABLE IF NOT EXISTS experiments ( id TEXT PRIMARY KEY, exp_name TEXT NOT NULL, version TEXT NOT NULL DEFAULT 'v2', exp_root TEXT DEFAULT 'logs', gpu_numbers TEXT DEFAULT '0', is_half INTEGER DEFAULT 1, audio_file_id TEXT, status TEXT NOT NULL DEFAULT 'created', created_at TEXT NOT NULL, updated_at TEXT ) ''') conn.execute('CREATE INDEX IF NOT EXISTS idx_experiments_status ON experiments(status)') conn.execute('CREATE INDEX IF NOT EXISTS idx_experiments_created ON experiments(created_at)') # Stages 表 (Advanced Mode 阶段状态) conn.execute(''' CREATE TABLE IF NOT EXISTS stages ( id TEXT PRIMARY KEY, experiment_id TEXT NOT NULL, stage_type TEXT NOT NULL, status TEXT DEFAULT 'pending', progress REAL DEFAULT 0, message TEXT, job_id TEXT, config TEXT, outputs TEXT, started_at TEXT, completed_at TEXT, error_message TEXT, FOREIGN KEY (experiment_id) REFERENCES experiments(id) ON DELETE CASCADE, UNIQUE (experiment_id, stage_type) ) ''') conn.execute('CREATE INDEX IF NOT EXISTS idx_stages_experiment ON stages(experiment_id)') conn.execute('CREATE INDEX IF NOT EXISTS idx_stages_status ON stages(status)') # Files 表 (文件记录) conn.execute(''' CREATE TABLE IF NOT EXISTS files ( id TEXT PRIMARY KEY, filename TEXT NOT NULL, content_type TEXT, size_bytes INTEGER DEFAULT 0, purpose TEXT DEFAULT 'training', duration_seconds REAL, sample_rate INTEGER, storage_path TEXT, uploaded_at TEXT NOT NULL ) ''') conn.execute('CREATE INDEX IF NOT EXISTS idx_files_purpose ON files(purpose)') conn.execute('CREATE INDEX IF NOT EXISTS idx_files_uploaded ON files(uploaded_at)') conn.commit() # ============================================================ # Task CRUD (Quick Mode) # ============================================================ async def create_task(self, task: Task) -> Task: """创建任务""" async with aiosqlite.connect(self.db_path) as db: await db.execute( '''INSERT INTO tasks (id, job_id, exp_name, status, config, current_stage, progress, stage_progress, message, error_message, created_at, started_at, completed_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)''', ( task.id, task.job_id, task.exp_name, task.status.value if isinstance(task.status, TaskStatus) else task.status, json.dumps(task.config, ensure_ascii=False) if task.config else None, task.current_stage, task.progress, task.stage_progress, task.message, task.error_message, task.created_at.isoformat() if task.created_at else datetime.utcnow().isoformat(), task.started_at.isoformat() if task.started_at else None, task.completed_at.isoformat() if task.completed_at else None, ) ) await db.commit() return task async def get_task(self, task_id: str) -> Optional[Task]: """获取任务""" async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( "SELECT * FROM tasks WHERE id = ?", (task_id,) ) as cursor: row = await cursor.fetchone() if row: return self._row_to_task(dict(row)) return None async def update_task(self, task_id: str, updates: Dict[str, Any]) -> Optional[Task]: """更新任务""" if not updates: return await self.get_task(task_id) # 处理特殊字段 processed = {} for key, value in updates.items(): if key == "status" and isinstance(value, TaskStatus): processed[key] = value.value elif key == "config" and isinstance(value, dict): processed[key] = json.dumps(value, ensure_ascii=False) elif key in ("created_at", "started_at", "completed_at") and isinstance(value, datetime): processed[key] = value.isoformat() else: processed[key] = value async with aiosqlite.connect(self.db_path) as db: set_clause = ", ".join(f"{k} = ?" for k in processed.keys()) values = list(processed.values()) + [task_id] await db.execute( f"UPDATE tasks SET {set_clause} WHERE id = ?", values ) await db.commit() return await self.get_task(task_id) async def list_tasks( self, status: Optional[str] = None, limit: int = 50, offset: int = 0 ) -> List[Task]: """查询任务列表""" async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row if status: query = """ SELECT * FROM tasks WHERE status = ? ORDER BY created_at DESC LIMIT ? OFFSET ? """ params = (status, limit, offset) else: query = """ SELECT * FROM tasks ORDER BY created_at DESC LIMIT ? OFFSET ? """ params = (limit, offset) async with db.execute(query, params) as cursor: rows = await cursor.fetchall() return [self._row_to_task(dict(row)) for row in rows] async def delete_task(self, task_id: str) -> bool: """删除任务""" async with aiosqlite.connect(self.db_path) as db: cursor = await db.execute( "DELETE FROM tasks WHERE id = ?", (task_id,) ) await db.commit() return cursor.rowcount > 0 async def count_tasks(self, status: Optional[str] = None) -> int: """统计任务数量""" async with aiosqlite.connect(self.db_path) as db: if status: async with db.execute( "SELECT COUNT(*) FROM tasks WHERE status = ?", (status,) ) as cursor: row = await cursor.fetchone() else: async with db.execute("SELECT COUNT(*) FROM tasks") as cursor: row = await cursor.fetchone() return row[0] if row else 0 async def get_task_by_exp_name(self, exp_name: str) -> Optional[Task]: """根据实验名称获取任务""" async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( "SELECT * FROM tasks WHERE exp_name = ? LIMIT 1", (exp_name,) ) as cursor: row = await cursor.fetchone() if row: return self._row_to_task(dict(row)) return None def _row_to_task(self, row: Dict[str, Any]) -> Task: """将数据库行转换为 Task 对象""" # 解析 config JSON config = row.get("config") if config and isinstance(config, str): try: config = json.loads(config) except json.JSONDecodeError: config = {} return Task.from_dict({ "id": row["id"], "job_id": row.get("job_id"), "exp_name": row["exp_name"], "status": row.get("status", "queued"), "config": config or {}, "current_stage": row.get("current_stage"), "progress": row.get("progress", 0.0), "stage_progress": row.get("stage_progress", 0.0), "message": row.get("message"), "error_message": row.get("error_message"), "created_at": row.get("created_at"), "started_at": row.get("started_at"), "completed_at": row.get("completed_at"), }) # ============================================================ # Experiment CRUD (Advanced Mode) # ============================================================ async def create_experiment(self, experiment: Dict[str, Any]) -> Dict[str, Any]: """创建实验""" exp_id = experiment.get("id") or f"exp-{uuid.uuid4().hex[:8]}" now = datetime.utcnow().isoformat() async with aiosqlite.connect(self.db_path) as db: # 创建实验记录 await db.execute( '''INSERT INTO experiments (id, exp_name, version, exp_root, gpu_numbers, is_half, audio_file_id, status, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)''', ( exp_id, experiment["exp_name"], experiment.get("version", "v2"), experiment.get("exp_root", "logs"), experiment.get("gpu_numbers", "0"), 1 if experiment.get("is_half", True) else 0, experiment.get("audio_file_id"), experiment.get("status", "created"), now, now, ) ) # 创建所有阶段的初始状态 for stage_type in STAGE_TYPES: stage_id = f"{exp_id}-{stage_type}" await db.execute( '''INSERT INTO stages (id, experiment_id, stage_type, status) VALUES (?, ?, ?, 'pending')''', (stage_id, exp_id, stage_type) ) await db.commit() return await self.get_experiment(exp_id) async def get_experiment(self, exp_id: str) -> Optional[Dict[str, Any]]: """获取实验""" async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row # 获取实验基本信息 async with db.execute( "SELECT * FROM experiments WHERE id = ?", (exp_id,) ) as cursor: row = await cursor.fetchone() if not row: return None experiment = dict(row) experiment["is_half"] = bool(experiment.get("is_half", 1)) # 获取所有阶段状态 stages = {} async with db.execute( "SELECT * FROM stages WHERE experiment_id = ?", (exp_id,) ) as cursor: stage_rows = await cursor.fetchall() for stage_row in stage_rows: stage = dict(stage_row) stage_type = stage["stage_type"] # 解析 JSON 字段 for json_field in ("config", "outputs"): if stage.get(json_field) and isinstance(stage[json_field], str): try: stage[json_field] = json.loads(stage[json_field]) except json.JSONDecodeError: stage[json_field] = None stages[stage_type] = stage experiment["stages"] = stages return experiment async def update_experiment( self, exp_id: str, updates: Dict[str, Any] ) -> Optional[Dict[str, Any]]: """更新实验""" if not updates: return await self.get_experiment(exp_id) # 处理 is_half 布尔值 processed = {} for key, value in updates.items(): if key == "is_half": processed[key] = 1 if value else 0 elif key == "updated_at" and isinstance(value, datetime): processed[key] = value.isoformat() elif key != "stages": # stages 单独处理 processed[key] = value # 添加更新时间 if "updated_at" not in processed: processed["updated_at"] = datetime.utcnow().isoformat() async with aiosqlite.connect(self.db_path) as db: if processed: set_clause = ", ".join(f"{k} = ?" for k in processed.keys()) values = list(processed.values()) + [exp_id] await db.execute( f"UPDATE experiments SET {set_clause} WHERE id = ?", values ) await db.commit() return await self.get_experiment(exp_id) async def list_experiments( self, status: Optional[str] = None, limit: int = 50, offset: int = 0 ) -> List[Dict[str, Any]]: """查询实验列表""" async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row if status: query = """ SELECT * FROM experiments WHERE status = ? ORDER BY created_at DESC LIMIT ? OFFSET ? """ params = (status, limit, offset) else: query = """ SELECT * FROM experiments ORDER BY created_at DESC LIMIT ? OFFSET ? """ params = (limit, offset) async with db.execute(query, params) as cursor: rows = await cursor.fetchall() results = [] for row in rows: exp = dict(row) exp["is_half"] = bool(exp.get("is_half", 1)) # 简化列表，不包含完整的 stages results.append(exp) return results async def delete_experiment(self, exp_id: str) -> bool: """删除实验及其阶段""" async with aiosqlite.connect(self.db_path) as db: # 先删除阶段 await db.execute( "DELETE FROM stages WHERE experiment_id = ?", (exp_id,) ) # 再删除实验 cursor = await db.execute( "DELETE FROM experiments WHERE id = ?", (exp_id,) ) await db.commit() return cursor.rowcount > 0 # ============================================================ # Stage 操作 (Advanced Mode) # ============================================================ async def update_stage( self, exp_id: str, stage_type: str, updates: Dict[str, Any] ) -> Optional[Dict[str, Any]]: """更新阶段状态""" if not updates: return await self.get_stage(exp_id, stage_type) # 处理 JSON 字段 processed = {} for key, value in updates.items(): if key in ("config", "outputs") and isinstance(value, dict): processed[key] = json.dumps(value, ensure_ascii=False) elif key in ("started_at", "completed_at") and isinstance(value, datetime): processed[key] = value.isoformat() else: processed[key] = value async with aiosqlite.connect(self.db_path) as db: set_clause = ", ".join(f"{k} = ?" for k in processed.keys()) values = list(processed.values()) + [exp_id, stage_type] await db.execute( f"UPDATE stages SET {set_clause} WHERE experiment_id = ? AND stage_type = ?", values ) await db.commit() # 同时更新实验的 updated_at await self.update_experiment(exp_id, {}) return await self.get_stage(exp_id, stage_type) async def get_stage( self, exp_id: str, stage_type: str ) -> Optional[Dict[str, Any]]: """获取阶段状态""" async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( "SELECT * FROM stages WHERE experiment_id = ? AND stage_type = ?", (exp_id, stage_type) ) as cursor: row = await cursor.fetchone() if not row: return None stage = dict(row) # 解析 JSON 字段 for json_field in ("config", "outputs"): if stage.get(json_field) and isinstance(stage[json_field], str): try: stage[json_field] = json.loads(stage[json_field]) except json.JSONDecodeError: stage[json_field] = None return stage async def get_all_stages(self, exp_id: str) -> List[Dict[str, Any]]: """获取实验的所有阶段状态""" async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( "SELECT * FROM stages WHERE experiment_id = ? ORDER BY id", (exp_id,) ) as cursor: rows = await cursor.fetchall() results = [] for row in rows: stage = dict(row) # 解析 JSON 字段 for json_field in ("config", "outputs"): if stage.get(json_field) and isinstance(stage[json_field], str): try: stage[json_field] = json.loads(stage[json_field]) except json.JSONDecodeError: stage[json_field] = None results.append(stage) return results # ============================================================ # File 记录 # ============================================================ async def create_file_record(self, file_data: Dict[str, Any]) -> Dict[str, Any]: """创建文件记录""" file_id = file_data.get("id") or str(uuid.uuid4()) now = datetime.utcnow().isoformat() async with aiosqlite.connect(self.db_path) as db: await db.execute( '''INSERT INTO files (id, filename, content_type, size_bytes, purpose, duration_seconds, sample_rate, storage_path, uploaded_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)''', ( file_id, file_data["filename"], file_data.get("content_type"), file_data.get("size_bytes", 0), file_data.get("purpose", "training"), file_data.get("duration_seconds"), file_data.get("sample_rate"), file_data.get("storage_path"), file_data.get("uploaded_at", now), ) ) await db.commit() return await self.get_file_record(file_id) async def get_file_record(self, file_id: str) -> Optional[Dict[str, Any]]: """获取文件记录""" async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( "SELECT * FROM files WHERE id = ?", (file_id,) ) as cursor: row = await cursor.fetchone() if row: return dict(row) return None async def delete_file_record(self, file_id: str) -> bool: """删除文件记录""" async with aiosqlite.connect(self.db_path) as db: cursor = await db.execute( "DELETE FROM files WHERE id = ?", (file_id,) ) await db.commit() return cursor.rowcount > 0 async def list_file_records( self, purpose: Optional[str] = None, limit: int = 50, offset: int = 0 ) -> List[Dict[str, Any]]: """查询文件记录列表""" async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row if purpose: query = """ SELECT * FROM files WHERE purpose = ? ORDER BY uploaded_at DESC LIMIT ? OFFSET ? """ params = (purpose, limit, offset) else: query = """ SELECT * FROM files ORDER BY uploaded_at DESC LIMIT ? OFFSET ? """ params = (limit, offset) async with db.execute(query, params) as cursor: rows = await cursor.fetchall() return [dict(row) for row in rows] async def count_file_records(self, purpose: Optional[str] = None) -> int: """统计文件记录数量""" async with aiosqlite.connect(self.db_path) as db: if purpose: async with db.execute( "SELECT COUNT(*) FROM files WHERE purpose = ?", (purpose,) ) as cursor: row = await cursor.fetchone() else: async with db.execute("SELECT COUNT(*) FROM files") as cursor: row = await cursor.fetchone() return row[0] if row else 0