Update scripts/restore_from_dataset.py

scripts/restore_from_dataset.py

@@ -2,10 +2,44 @@
 import os
 import sys
 import shutil
+import datetime
 from huggingface_hub import HfApi, login, hf_hub_download
 
-DB_TARGET = "/app/server/data/freeapi.db"
-DATASET_REPO = "lydgs/freellm-backup"   # 已改为你的用户名
+# ===== 配置 =====
+DB_TARGET = "/app/server/data/freeapi.db"          # 本地数据库路径
+DATASET_REPO = "lydgs/freellm-backup"              # 你的数据集名称
+BACKUP_PREFIX = "freeapi_backup"                   # 备份文件前缀
+# ================
+
+def upload_backup(db_path, dataset_repo, backup_prefix):
+    """上传数据库作为备份，返回是否成功"""
+    if not os.path.exists(db_path):
+        print(f"⚠️ 本地数据库不存在: {db_path}，无法创建备份")
+        return False
+    
+    token = os.getenv("HF_TOKEN")
+    if not token:
+        print("❌ 环境变量 HF_TOKEN 未设置，无法创建备份")
+        return False
+    
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    backup_name = f"{backup_prefix}_{timestamp}.db"
+    
+    login(token=token)
+    api = HfApi()
+    
+    try:
+        api.upload_file(
+            path_or_fileobj=db_path,
+            path_in_repo=backup_name,
+            repo_id=dataset_repo,
+            repo_type="dataset"
+        )
+        print(f"✅ 初始备份创建成功: {backup_name}")
+        return True
+    except Exception as e:
+        print(f"❌ 创建初始备份失败: {e}")
+        return False
 
 def restore_latest_backup():
     token = os.getenv("HF_TOKEN")
@@ -17,12 +51,21 @@ def restore_latest_backup():
     api = HfApi()
 
     try:
+        # 列出数据集中的所有文件
         files = api.list_repo_files(repo_id=DATASET_REPO, repo_type="dataset")
-        db_files = [f for f in files if f.endswith(".db") and f.startswith("freeapi_backup_")]
+        db_files = [f for f in files if f.endswith(".db") and f.startswith(BACKUP_PREFIX)]
+        
         if not db_files:
-            print("ℹ️ 数据集中没有找到备份文件，跳过恢复")
+            print("ℹ️ 数据集中没有找到备份文件")
+            # 如果没有备份，但本地数据库已存在，则主动创建一份备份
+            if os.path.exists(DB_TARGET):
+                print("🔄 检测到本地数据库存在，将创建初始备份...")
+                upload_backup(DB_TARGET, DATASET_REPO, BACKUP_PREFIX)
+            else:
+                print("ℹ️ 本地数据库也不存在，跳过恢复")
             return
 
+        # 有备份文件：取最新的进行恢复
         db_files.sort()
         latest = db_files[-1]
         print(f"🔄 发现最新备份: {latest}，开始下载...")
@@ -34,7 +77,10 @@ def restore_latest_backup():
             local_dir="/tmp"
         )
 
+        # 确保目标目录存在
         os.makedirs(os.path.dirname(DB_TARGET), exist_ok=True)
+
+        # 复制到目标位置
         shutil.copy(downloaded, DB_TARGET)
         print(f"✅ 数据库恢复成功: {DB_TARGET} (来自备份 {latest})")
         os.remove(downloaded)