use rusqlite::Connection; use anyhow::{Result, anyhow}; use std::collections::HashMap; /// 数据库迁移版本号 pub type MigrationVersion = u32; /// 单个迁移脚本 #[derive(Debug, Clone)] pub struct Migration { /// 迁移版本号 pub version: MigrationVersion, /// 迁移描述 pub description: String, /// 向前迁移SQL pub up_sql: String, /// 回滚迁移SQL(可选) pub down_sql: Option, } /// 迁移管理器 pub struct MigrationManager { migrations: HashMap, } impl MigrationManager { /// 创建新的迁移管理器 pub fn new() -> Self { let mut manager = Self { migrations: HashMap::new(), }; // 注册所有迁移 manager.register_migrations(); manager } /// 注册所有迁移脚本 fn register_migrations(&mut self) { // 迁移 1: 初始化数据库表结构 self.add_migration(Migration { version: 1, description: "初始化数据库表结构".to_string(), up_sql: include_str!("migrations/001_initial_schema.sql").to_string(), down_sql: Some(include_str!("migrations/001_initial_schema_down.sql").to_string()), }); // 迁移 2: 添加模特ID字段到素材表 self.add_migration(Migration { version: 2, description: "添加模特ID字段到素材表".to_string(), up_sql: include_str!("migrations/002_add_model_id_to_materials.sql").to_string(), down_sql: Some(include_str!("migrations/002_add_model_id_to_materials_down.sql").to_string()), }); // 迁移 3: 添加文件存在状态字段到模板素材表 self.add_migration(Migration { version: 3, description: "添加文件存在状态字段到模板素材表".to_string(), up_sql: include_str!("migrations/003_add_file_status_to_template_materials.sql").to_string(), down_sql: Some(include_str!("migrations/003_add_file_status_to_template_materials_down.sql").to_string()), }); // 迁移 4: 移除模板表的project_id字段 self.add_migration(Migration { version: 4, description: "移除模板表的project_id字段".to_string(), up_sql: include_str!("migrations/004_remove_project_id_from_templates.sql").to_string(), down_sql: None, // 这个迁移不支持回滚,因为涉及数据丢失 }); // 迁移 5: 修复AI分类表的is_active字段类型 self.add_migration(Migration { version: 5, description: "修复AI分类表的is_active字段类型".to_string(), up_sql: include_str!("migrations/005_fix_ai_classifications_is_active.sql").to_string(), down_sql: None, // 类型修复不支持回滚 }); // 迁移 6: 添加导出状态字段到模板匹配结果表 self.add_migration(Migration { version: 6, description: "添加导出状态字段到模板匹配结果表".to_string(), up_sql: include_str!("migrations/006_add_export_status_to_template_matching.sql").to_string(), down_sql: Some(include_str!("migrations/006_add_export_status_to_template_matching_down.sql").to_string()), }); // 迁移 7: 修复项目表的is_active字段数据 self.add_migration(Migration { version: 7, description: "修复项目表的is_active字段数据".to_string(), up_sql: include_str!("migrations/007_fix_projects_is_active_data.sql").to_string(), down_sql: None, // 数据修复不支持回滚 }); // 迁移 8: 添加缩略图和使用跟踪字段 self.add_migration(Migration { version: 8, description: "添加缩略图和使用跟踪字段".to_string(), up_sql: include_str!("migrations/008_add_thumbnail_and_usage_tracking.sql").to_string(), down_sql: Some(include_str!("migrations/008_add_thumbnail_and_usage_tracking_down.sql").to_string()), }); // 迁移 9: 修复template_materials表的file_size字段,允许为NULL self.add_migration(Migration { version: 9, description: "修复template_materials表的file_size字段,允许为NULL".to_string(), up_sql: include_str!("migrations/009_fix_template_materials_file_size_nullable.sql").to_string(), down_sql: Some(include_str!("migrations/009_fix_template_materials_file_size_nullable_down.sql").to_string()), }); // 迁移 10: 修复video_classification_records表结构,添加缺失字段 self.add_migration(Migration { version: 10, description: "修复video_classification_records表结构,添加缺失字段".to_string(), up_sql: include_str!("migrations/010_fix_video_classification_records_schema.sql").to_string(), down_sql: Some(include_str!("migrations/010_fix_video_classification_records_schema_down.sql").to_string()), }); // 迁移 11: 修复template_matching_results表结构,添加缺失字段 self.add_migration(Migration { version: 11, description: "修复template_matching_results表结构,添加缺失字段".to_string(), up_sql: include_str!("migrations/011_fix_template_matching_results_schema.sql").to_string(), down_sql: Some(include_str!("migrations/011_fix_template_matching_results_schema_down.sql").to_string()), }); // 迁移 12: 创建video_classification_tasks表 self.add_migration(Migration { version: 12, description: "创建video_classification_tasks表".to_string(), up_sql: include_str!("migrations/012_create_video_classification_tasks_table.sql").to_string(), down_sql: Some(include_str!("migrations/012_create_video_classification_tasks_table_down.sql").to_string()), }); // 迁移 13: 创建export_records表 self.add_migration(Migration { version: 13, description: "创建export_records表".to_string(), up_sql: include_str!("migrations/013_create_export_records_table.sql").to_string(), down_sql: Some(include_str!("migrations/013_create_export_records_table_down.sql").to_string()), }); // 迁移 14: 创建project_template_bindings表 self.add_migration(Migration { version: 14, description: "创建project_template_bindings表".to_string(), up_sql: include_str!("migrations/014_create_project_template_bindings_table.sql").to_string(), down_sql: Some(include_str!("migrations/014_create_project_template_bindings_table_down.sql").to_string()), }); // 迁移 15: 添加默认的AI分类数据 self.add_migration(Migration { version: 15, description: "添加默认的AI分类数据".to_string(), up_sql: include_str!("migrations/015_add_default_ai_classifications.sql").to_string(), down_sql: Some(include_str!("migrations/015_add_default_ai_classifications_down.sql").to_string()), }); // 迁移 16: 修复AI分类表中的日期时间格式 self.add_migration(Migration { version: 16, description: "修复AI分类表中的日期时间格式".to_string(), up_sql: include_str!("migrations/016_fix_ai_classifications_datetime_format.sql").to_string(), down_sql: Some(include_str!("migrations/016_fix_ai_classifications_datetime_format_down.sql").to_string()), }); // 迁移 17: 创建matching_segment_results和matching_failed_segment_results表 self.add_migration(Migration { version: 17, description: "创建matching_segment_results和matching_failed_segment_results表".to_string(), up_sql: include_str!("migrations/017_create_matching_segment_results_tables.sql").to_string(), down_sql: Some(include_str!("migrations/017_create_matching_segment_results_tables_down.sql").to_string()), }); // 迁移 18: 创建material_usage_records表 self.add_migration(Migration { version: 18, description: "创建material_usage_records表".to_string(), up_sql: include_str!("migrations/018_create_material_usage_records_table.sql").to_string(), down_sql: Some(include_str!("migrations/018_create_material_usage_records_table_down.sql").to_string()), }); // 迁移 19: 为ai_classifications表添加weight字段 self.add_migration(Migration { version: 19, description: "为ai_classifications表添加weight字段".to_string(), up_sql: include_str!("migrations/019_add_weight_to_ai_classifications.sql").to_string(), down_sql: Some(include_str!("migrations/019_add_weight_to_ai_classifications_down.sql").to_string()), }); // 迁移 20: 创建模板片段权重配置表 self.add_migration(Migration { version: 20, description: "创建模板片段权重配置表".to_string(), up_sql: include_str!("migrations/020_create_template_segment_weights_table.sql").to_string(), down_sql: Some(include_str!("migrations/020_create_template_segment_weights_table_down.sql").to_string()), }); // 迁移 21: 创建穿搭方案收藏表 self.add_migration(Migration { version: 21, description: "创建穿搭方案收藏表".to_string(), up_sql: include_str!("migrations/021_outfit_favorites.sql").to_string(), down_sql: Some(include_str!("migrations/021_outfit_favorites_down.sql").to_string()), }); // 迁移 22: 创建图片生成记录表 self.add_migration(Migration { version: 22, description: "创建图片生成记录表".to_string(), up_sql: include_str!("migrations/022_create_image_generation_records_table.sql").to_string(), down_sql: Some(include_str!("migrations/022_create_image_generation_records_table_down.sql").to_string()), }); // 迁移 23: 创建穿搭图片相关表 self.add_migration(Migration { version: 23, description: "创建穿搭图片相关表".to_string(), up_sql: include_str!("migrations/023_create_outfit_images_tables.sql").to_string(), down_sql: Some(include_str!("migrations/023_create_outfit_images_tables_down.sql").to_string()), }); // 迁移 23: 创建声音克隆记录表 self.add_migration(Migration { version: 23, description: "创建声音克隆记录表".to_string(), up_sql: include_str!("migrations/023_create_voice_clone_table.sql").to_string(), down_sql: Some(include_str!("migrations/023_create_voice_clone_table_down.sql").to_string()), }); // 迁移 24: 创建语音生成记录表 self.add_migration(Migration { version: 24, description: "创建语音生成记录表".to_string(), up_sql: include_str!("migrations/024_create_speech_generation_records_table.sql").to_string(), down_sql: None, // 暂时不提供回滚脚本 }); } /// 添加迁移 fn add_migration(&mut self, migration: Migration) { self.migrations.insert(migration.version, migration); } /// 获取当前数据库版本 pub fn get_current_version(&self, conn: &Connection) -> Result { // 确保迁移历史表存在 self.ensure_migration_table(conn)?; let version = conn.query_row( "SELECT MAX(version) FROM schema_migrations WHERE success = 1", [], |row| row.get::<_, Option>(0) ).unwrap_or(Some(0)).unwrap_or(0); Ok(version) } /// 获取最新版本号 pub fn get_latest_version(&self) -> MigrationVersion { self.migrations.keys().max().copied().unwrap_or(0) } /// 确保迁移历史表存在 fn ensure_migration_table(&self, conn: &Connection) -> Result<()> { conn.execute( "CREATE TABLE IF NOT EXISTS schema_migrations ( version INTEGER PRIMARY KEY, description TEXT NOT NULL, applied_at DATETIME DEFAULT CURRENT_TIMESTAMP, success INTEGER NOT NULL DEFAULT 0, error_message TEXT )", [], )?; Ok(()) } /// 运行所有待执行的迁移 pub fn migrate(&self, conn: &Connection) -> Result<()> { let current_version = self.get_current_version(conn)?; let latest_version = self.get_latest_version(); if current_version >= latest_version { println!("数据库已是最新版本 v{}", current_version); return Ok(()); } println!("开始数据库迁移: v{} -> v{}", current_version, latest_version); // 按版本号排序执行迁移 let mut versions: Vec<_> = self.migrations.keys().collect(); versions.sort(); for &version in versions { if version > current_version { self.apply_migration(conn, version)?; } } println!("数据库迁移完成"); Ok(()) } /// 应用单个迁移 fn apply_migration(&self, conn: &Connection, version: MigrationVersion) -> Result<()> { let migration = self.migrations.get(&version) .ok_or_else(|| anyhow!("迁移版本 {} 不存在", version))?; println!("应用迁移 v{}: {}", version, migration.description); // 开始事务 let tx = conn.unchecked_transaction()?; // 记录迁移开始 tx.execute( "INSERT OR REPLACE INTO schema_migrations (version, description, success) VALUES (?1, ?2, 0)", [&version.to_string(), &migration.description], )?; // 执行迁移SQL match self.execute_migration_sql(&tx, &migration.up_sql) { Ok(_) => { // 标记迁移成功 tx.execute( "UPDATE schema_migrations SET success = 1, applied_at = CURRENT_TIMESTAMP WHERE version = ?1", [&version.to_string()], )?; tx.commit()?; println!("迁移 v{} 应用成功", version); Ok(()) } Err(e) => { // 记录错误信息 let _ = tx.execute( "UPDATE schema_migrations SET error_message = ?1 WHERE version = ?2", [&e.to_string(), &version.to_string()], ); tx.rollback()?; Err(anyhow!("迁移 v{} 失败: {}", version, e)) } } } /// 执行迁移SQL(支持多条语句) fn execute_migration_sql(&self, conn: &Connection, sql: &str) -> Result<()> { // 使用SQLite的execute_batch方法来执行多条SQL语句 // 这个方法能正确处理SQL语句的分割,包括注释和字符串中的分号 conn.execute_batch(sql)?; Ok(()) } /// 回滚到指定版本 pub fn rollback_to(&self, conn: &Connection, target_version: MigrationVersion) -> Result<()> { let current_version = self.get_current_version(conn)?; if target_version >= current_version { return Err(anyhow!("目标版本 {} 不能大于等于当前版本 {}", target_version, current_version)); } println!("开始回滚数据库: v{} -> v{}", current_version, target_version); // 按版本号倒序回滚 let mut versions: Vec<_> = self.migrations.keys().collect(); versions.sort_by(|a, b| b.cmp(a)); // 倒序 for &version in versions { if version > target_version && version <= current_version { self.rollback_migration(conn, version)?; } } println!("数据库回滚完成"); Ok(()) } /// 回滚单个迁移 fn rollback_migration(&self, conn: &Connection, version: MigrationVersion) -> Result<()> { let migration = self.migrations.get(&version) .ok_or_else(|| anyhow!("迁移版本 {} 不存在", version))?; let down_sql = migration.down_sql.as_ref() .ok_or_else(|| anyhow!("迁移 v{} 不支持回滚", version))?; println!("回滚迁移 v{}: {}", version, migration.description); // 开始事务 let tx = conn.unchecked_transaction()?; // 执行回滚SQL match self.execute_migration_sql(&tx, down_sql) { Ok(_) => { // 删除迁移记录 tx.execute( "DELETE FROM schema_migrations WHERE version = ?1", [&version.to_string()], )?; tx.commit()?; println!("迁移 v{} 回滚成功", version); Ok(()) } Err(e) => { tx.rollback()?; Err(anyhow!("迁移 v{} 回滚失败: {}", version, e)) } } } /// 获取迁移历史 pub fn get_migration_history(&self, conn: &Connection) -> Result> { self.ensure_migration_table(conn)?; let mut stmt = conn.prepare( "SELECT version, description, applied_at, success FROM schema_migrations ORDER BY version" )?; let rows = stmt.query_map([], |row| { Ok(( row.get::<_, u32>(0)?, row.get::<_, String>(1)?, row.get::<_, String>(2)?, row.get::<_, i32>(3)? == 1, )) })?; let mut history = Vec::new(); for row in rows { history.push(row?); } Ok(history) } } #[cfg(test)] mod tests { use super::*; use rusqlite::Connection; use tempfile::NamedTempFile; fn create_test_connection() -> Connection { let temp_file = NamedTempFile::new().unwrap(); let conn = Connection::open(temp_file.path()).unwrap(); // 配置数据库设置 conn.pragma_update(None, "foreign_keys", "ON").unwrap(); conn.pragma_update(None, "journal_mode", "WAL").unwrap(); conn } #[test] fn test_migration_manager_creation() { let manager = MigrationManager::new(); let latest_version = manager.get_latest_version(); // 应该有至少8个迁移版本 assert!(latest_version >= 8, "应该有至少8个迁移版本,实际: {}", latest_version); } #[test] fn test_initial_version_is_zero() { let conn = create_test_connection(); let manager = MigrationManager::new(); let current_version = manager.get_current_version(&conn).unwrap(); assert_eq!(current_version, 0, "新数据库的版本应该是0"); } #[test] fn test_migration_table_creation() { let conn = create_test_connection(); let manager = MigrationManager::new(); // 确保迁移表存在 manager.ensure_migration_table(&conn).unwrap(); // 验证表是否存在 let table_exists: bool = conn.query_row( "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='schema_migrations'", [], |row| row.get::<_, i32>(0) ).unwrap() > 0; assert!(table_exists, "schema_migrations表应该存在"); } #[test] fn test_full_migration() { let conn = create_test_connection(); let manager = MigrationManager::new(); // 执行所有迁移 manager.migrate(&conn).unwrap(); // 验证版本已更新 let current_version = manager.get_current_version(&conn).unwrap(); let latest_version = manager.get_latest_version(); assert_eq!(current_version, latest_version, "迁移后版本应该是最新版本"); // 验证一些关键表是否存在 let tables = vec![ "projects", "materials", "material_segments", "models", "ai_classifications", "templates", "template_materials", "tracks", "track_segments", "video_classification_records" ]; for table in tables { let table_exists: bool = conn.query_row( "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name=?1", [table], |row| row.get::<_, i32>(0) ).unwrap() > 0; assert!(table_exists, "表 {} 应该存在", table); } } #[test] fn test_migration_history() { let conn = create_test_connection(); let manager = MigrationManager::new(); // 执行迁移 manager.migrate(&conn).unwrap(); // 获取迁移历史 let history = manager.get_migration_history(&conn).unwrap(); // 应该有迁移记录 assert!(!history.is_empty(), "应该有迁移历史记录"); // 所有迁移都应该成功 for (version, description, _applied_at, success) in history { assert!(success, "迁移 v{}: {} 应该成功", version, description); } } #[test] fn test_idempotent_migration() { let conn = create_test_connection(); let manager = MigrationManager::new(); // 第一次迁移 manager.migrate(&conn).unwrap(); let version_after_first = manager.get_current_version(&conn).unwrap(); // 第二次迁移(应该是幂等的) manager.migrate(&conn).unwrap(); let version_after_second = manager.get_current_version(&conn).unwrap(); assert_eq!(version_after_first, version_after_second, "重复迁移应该是幂等的"); } #[test] fn test_specific_migration_features() { let conn = create_test_connection(); let manager = MigrationManager::new(); // 执行迁移 manager.migrate(&conn).unwrap(); // 测试特定的迁移功能 // 1. 检查materials表是否有model_id字段(迁移2) let has_model_id = conn.prepare("SELECT model_id FROM materials LIMIT 1").is_ok(); assert!(has_model_id, "materials表应该有model_id字段"); // 2. 检查template_materials表是否有file_exists字段(迁移3) let has_file_exists = conn.prepare("SELECT file_exists FROM template_materials LIMIT 1").is_ok(); assert!(has_file_exists, "template_materials表应该有file_exists字段"); // 3. 检查materials表是否有thumbnail_path字段(迁移8) let has_materials_thumbnail_path = conn.prepare("SELECT thumbnail_path FROM materials LIMIT 1").is_ok(); assert!(has_materials_thumbnail_path, "materials表应该有thumbnail_path字段"); // 4. 检查material_segments表是否有usage_count字段(迁移8) let has_usage_count = conn.prepare("SELECT usage_count FROM material_segments LIMIT 1").is_ok(); assert!(has_usage_count, "material_segments表应该有usage_count字段"); // 5. 检查template_matching_results表是否有export_count字段(迁移8) let has_export_count = conn.prepare("SELECT export_count FROM template_matching_results LIMIT 1").is_ok(); assert!(has_export_count, "template_matching_results表应该有export_count字段"); // 6. 验证初始化时就存在的字段 let has_matching_rule = conn.prepare("SELECT matching_rule FROM track_segments LIMIT 1").is_ok(); assert!(has_matching_rule, "track_segments表应该有matching_rule字段(初始化时创建)"); let has_segments_thumbnail_path = conn.prepare("SELECT thumbnail_path FROM material_segments LIMIT 1").is_ok(); assert!(has_segments_thumbnail_path, "material_segments表应该有thumbnail_path字段(初始化时创建)"); } }