From 8c1411a31d78da94369d8248224378cb9ee6c24b Mon Sep 17 00:00:00 2001
From: imeepos <imeepos@sker.com>
Date: Wed, 20 Aug 2025 16:40:37 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8DComfyUI=E4=BB=BB?=
 =?UTF-8?q?=E5=8A=A1=E6=95=B0=E6=8D=AE=E5=9C=A8=E9=87=8D=E5=90=AF=E5=90=8E?=
 =?UTF-8?q?=E4=B8=A2=E5=A4=B1=E7=9A=84=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 修复数据库锁检查逻辑，使用query_row替代execute避免误判
- 改进WAL文件处理机制，使用安全的checkpoint合并数据
- 将UniComfyUI相关表添加到健康检查列表中
- 防止WAL文件被过度清理导致数据丢失

解决了ComfyUI任务管理中心页面数据在应用重启后被清空的问题
---
 .../src-tauri/src/infrastructure/database.rs  | 43 ++++++++++++++++---
 1 file changed, 36 insertions(+), 7 deletions(-)

diff --git a/apps/desktop/src-tauri/src/infrastructure/database.rs b/apps/desktop/src-tauri/src/infrastructure/database.rs
index 6b293ca..c4ab7e5 100644
--- a/apps/desktop/src-tauri/src/infrastructure/database.rs
+++ b/apps/desktop/src-tauri/src/infrastructure/database.rs
@@ -133,8 +133,8 @@ impl Database {
         // 尝试打开数据库进行快速检查
         match Connection::open(db_path) {
             Ok(conn) => {
-                // 尝试执行一个简单的查询
-                match conn.execute("SELECT 1", []) {
+                // 尝试执行一个简单的查询（使用query而不是execute）
+                match conn.query_row("SELECT 1", [], |_| Ok(())) {
                     Ok(_) => {
                         println!("数据库连接检查通过");
                         Ok(())
@@ -146,19 +146,45 @@ impl Database {
         }
     }
 
-    /// 清理数据库锁文件
+    /// 安全清理数据库锁文件
+    /// 只有在确认数据库无法正常访问时才清理WAL文件
     fn cleanup_database_locks(db_path: &PathBuf) -> Result<()> {
         let wal_path = db_path.with_extension("db-wal");
         let shm_path = db_path.with_extension("db-shm");
 
+        // 首先尝试通过正常的SQLite机制来处理WAL文件
         if wal_path.exists() {
-            println!("清理 WAL 文件: {:?}", wal_path);
-            std::fs::remove_file(&wal_path).ok(); // 忽略错误
+            println!("⚠️  检测到WAL文件，尝试安全恢复数据");
+
+            // 尝试打开数据库并执行checkpoint来合并WAL文件
+            match Connection::open(db_path) {
+                Ok(conn) => {
+                    // 尝试执行WAL checkpoint来合并数据
+                    match conn.pragma_update(None, "wal_checkpoint", "TRUNCATE") {
+                        Ok(_) => {
+                            println!("✅ WAL文件已安全合并到主数据库");
+                            return Ok(());
+                        },
+                        Err(e) => {
+                            println!("⚠️  WAL checkpoint失败: {}, 将保留WAL文件", e);
+                            // 不删除WAL文件，让SQLite自己处理
+                            return Ok(());
+                        }
+                    }
+                },
+                Err(e) => {
+                    println!("⚠️  无法打开数据库进行WAL恢复: {}", e);
+                    // 只有在完全无法访问数据库时才考虑删除WAL文件
+                    println!("🚨 作为最后手段，删除WAL文件（可能丢失数据）");
+                    std::fs::remove_file(&wal_path).ok();
+                }
+            }
         }
 
+        // SHM文件可以安全删除，它只是共享内存映射
         if shm_path.exists() {
             println!("清理 SHM 文件: {:?}", shm_path);
-            std::fs::remove_file(&shm_path).ok(); // 忽略错误
+            std::fs::remove_file(&shm_path).ok();
         }
 
         Ok(())
@@ -477,7 +503,10 @@ impl Database {
                 "projects", "materials", "material_segments", "models",
                 "ai_classifications", "templates", "template_materials",
                 "tracks", "track_segments", "video_classification_records",
-                "schema_migrations"
+                "schema_migrations",
+                // UniComfyUI 相关表
+                "uni_comfyui_task", "uni_comfyui_batch_task",
+                "uni_comfyui_batch_item", "uni_comfyui_workflow_cache"
             ];
 
             let mut missing_tables = Vec::new();