fix: 施工方案入mysql增加重试连接

src/app/scripts/plan_info_in_database.py

@@ -11,12 +11,12 @@ from pathlib import Path
 from typing import Any, Dict, Optional
 
 from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
 
 from app.config.database import get_async_sessionmaker, get_async_engine
 from app.config.setting import settings
 
 
-
 # ============================================
 # 配置参数区域 - 便于修改
 # ============================================
@@ -192,7 +192,6 @@ async def insert_document_main(
             "kb_method": DOCUMENT_MAIN_DEFAULTS["kb_method"],
             "whether_to_task": DOCUMENT_MAIN_DEFAULTS["whether_to_task"],
         })
-
         return True, None
     except Exception as e:
         error_message = str(e)
@@ -283,7 +282,6 @@ async def insert_plan_base_info(session, doc_data: Dict[str, Any], folder_name:
             "updated_by": PLAN_BASE_INFO_DEFAULTS["updated_by"],
             "updated_time": datetime.now(),
         })
-        
         return True, None
     except Exception as e:
         error_message = str(e)
@@ -370,47 +368,55 @@ async def process_folder(root_folder: str | Path) -> Dict[str, Any]:
             md_url = f"/plan/{doc_id}.md"
             json_url = f"/plan/{doc_id}.json"
 
-            # 插入数据库
-            async with SessionMaker() as session:
-                try:
-                    # 先插入主表
-                    main_ok, main_error = await insert_document_main(session, doc_data, title, file_extension, file_url, md_url, json_url, folder_name)
-                    if not main_ok:
+            # 插入数据库（带连接丢失自动重试）
+            max_retries = 3
+            inserted = False
+            last_error = None
+            for attempt in range(max_retries):
+                async with SessionMaker() as session:
+                    try:
+                        # 先插入主表
+                        main_ok, main_error = await insert_document_main(session, doc_data, title, file_extension, file_url, md_url, json_url, folder_name)
+                        if not main_ok:
+                            await session.rollback()
+                            last_error = main_error
+                            break  # 业务逻辑错误不重试
+
+                        # 再插入施工方案基础信息表
+                        plan_ok, plan_error = await insert_plan_base_info(session, doc_data, folder_name)
+                        if not plan_ok:
+                            await session.rollback()
+                            last_error = plan_error
+                            break  # 业务逻辑错误不重试
+                        
+                        await session.commit()
+                        print(f"📁 {folder_name} ✅")
+                        stats["success"] += 1
+                        inserted = True
+                        break
+                    except OperationalError as e:
                         await session.rollback()
-                        stats["failed_items"].append({
-                            "folder": folder_name,
-                            "error": "插入主表失败",
-                            "reason": main_error,
-                        })
-                        stats["failed"] += 1
-                        continue
-
-                    # 再插入施工方案基础信息表
-                    plan_ok, plan_error = await insert_plan_base_info(session, doc_data, folder_name)
-                    if not plan_ok:
+                        error_message = str(e)
+                        if attempt < max_retries - 1 and ("Lost connection" in error_message or "MySQL server has gone away" in error_message):
+                            wait_time = (attempt + 1) * 2
+                            print(f"📁 {folder_name} ⚠️ 连接丢失，{wait_time}秒后重试 ({attempt + 1}/{max_retries - 1})...")
+                            await asyncio.sleep(wait_time)
+                            continue
+                        last_error = error_message
+                        break
+                    except Exception as e:
                         await session.rollback()
-                        stats["failed_items"].append({
-                            "folder": folder_name,
-                            "error": "插入施工方案基础信息表失败",
-                            "reason": plan_error,
-                        })
-                        stats["failed"] += 1
-                        continue
-                    
-                    await session.commit()
-                    print(f"📁 {folder_name} ✅")
-                    stats["success"] += 1
-                    
-                except Exception as e:
-                    error_message = str(e)
-                    await session.rollback()
-                    print(f"📁 {folder_name} ❌ ({error_message})")
-                    stats["failed_items"].append({
-                        "folder": folder_name,
-                        "error": "未知错误",
-                        "reason": error_message,
-                    })
-                    stats["failed"] += 1
+                        last_error = str(e)
+                        break
+            
+            if not inserted:
+                print(f"📁 {folder_name} ❌ ({last_error})")
+                stats["failed_items"].append({
+                    "folder": folder_name,
+                    "error": "未知错误",
+                    "reason": last_error,
+                })
+                stats["failed"] += 1
                 
         except Exception as e:
             error_message = str(e)