Update scripts/restore_from_dataset.py

scripts/restore_from_dataset.py

@@ -2,87 +2,114 @@
 import os
 import sys
 import shutil
+import json
+import sqlite3
+import tarfile
+import tempfile
 import datetime
 from huggingface_hub import HfApi, login, hf_hub_download
 
-# ===== 配置 =====
-DB_TARGET = "/app/server/data/freeapi.db"          # 本地数据库路径
-DATASET_REPO = "lydgs/freellm-backup"              # 你的数据集名称
-BACKUP_PREFIX = "freeapi_backup"                   # 备份文件前缀
-# ================
+DB_TARGET = "/app/server/data/freeapi.db"
+DATASET_REPO = "lydgs/freellm-backup"
+CONFIG_PREFIX = "config_export"
 
-def upload_backup(db_path, dataset_repo, backup_prefix):
-    """上传数据库作为备份，返回是否成功"""
-    if not os.path.exists(db_path):
-        print(f"⚠️ 本地数据库不存在: {db_path}，无法创建备份")
+def import_json_to_db(json_dir, db_path):
+    """从包含 JSON 文件的目录重建数据库"""
+    if not os.path.exists(json_dir):
         return False
+    # 确保数据库所在目录存在
+    os.makedirs(os.path.dirname(db_path), exist_ok=True)
+    # 如果已有数据库，先删除（避免表冲突）
+    if os.path.exists(db_path):
+        os.remove(db_path)
     
-    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()
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
     
-    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
+    for json_file in os.listdir(json_dir):
+        if not json_file.endswith('.json'):
+            continue
+        table_name = json_file[:-5]
+        with open(os.path.join(json_dir, json_file), 'r', encoding='utf-8') as f:
+            data = json.load(f)
+        if not data:
+            # 空表：至少需要创建表结构。我们假设有列信息？简化处理：跳过或根据第一个 row 创建。
+            continue
+        # 根据第一行数据创建表（如果表不存在）
+        first_row = data[0]
+        columns = ', '.join(f'"{col}"' for col in first_row.keys())
+        placeholders = ', '.join('?' for _ in first_row)
+        # 删除旧表（如果有）
+        cursor.execute(f"DROP TABLE IF EXISTS {table_name}")
+        cursor.execute(f"CREATE TABLE {table_name} ({columns})")
+        # 插入所有行
+        for row in data:
+            values = list(row.values())
+            sql = f"INSERT INTO {table_name} ({', '.join(row.keys())}) VALUES ({placeholders})"
+            cursor.execute(sql, values)
+    conn.commit()
+    conn.close()
+    return True
 
 def restore_latest_backup():
     token = os.getenv("HF_TOKEN")
     if not token:
-        print("⚠️ 未设置 HF_TOKEN 环境变量，跳过数据库恢复")
+        print("⚠️ 未设置 HF_TOKEN，跳过恢复")
         return
-
     login(token=token)
     api = HfApi()
+    
+    # 1. 优先尝试使用 JSON 配置包重建数据库
+    try:
+        files = api.list_repo_files(repo_id=DATASET_REPO, repo_type="dataset")
+        config_files = [f for f in files if f.startswith(CONFIG_PREFIX) and f.endswith('.tar.gz')]
+        if config_files:
+            config_files.sort(reverse=True)  # 最新的排前面
+            latest_config = config_files[0]
+            print(f"🔄 发现配置包: {latest_config}，将用 JSON 重建数据库...")
+            # 下载到临时目录
+            with tempfile.TemporaryDirectory() as tmpdir:
+                downloaded = hf_hub_download(
+                    repo_id=DATASET_REPO,
+                    filename=latest_config,
+                    repo_type="dataset",
+                    local_dir=tmpdir
+                )
+                # 解压
+                extract_dir = os.path.join(tmpdir, "extract")
+                os.makedirs(extract_dir, exist_ok=True)
+                with tarfile.open(downloaded, "r:gz") as tar:
+                    tar.extractall(extract_dir)
+                # 导入 JSON 到数据库
+                if import_json_to_db(extract_dir, DB_TARGET):
+                    print(f"✅ 数据库从 JSON 配置重建成功: {DB_TARGET}")
+                    return
+                else:
+                    print("⚠️ JSON 导入失败，尝试回退到二进制备份...")
+        else:
+            print("ℹ️ 未找到 JSON 配置包，尝试二进制备份...")
+    except Exception as e:
+        print(f"⚠️ 处理 JSON 配置包时出错: {e}，尝试二进制备份...")
 
+    # 2. 回退：使用二进制 DB 备份
     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(BACKUP_PREFIX)]
-        
+        db_files = [f for f in files if f.endswith(".db") and f.startswith("freeapi_backup_")]
         if not db_files:
-            print("ℹ️ 数据集中没有找到备份文件")
-            # 如果没有备份，但本地数据库已存在，则主动创建一份备份
-            if os.path.exists(DB_TARGET):
-                print("🔄 检测到本地数据库存在，将创建初始备份...")
-                upload_backup(DB_TARGET, DATASET_REPO, BACKUP_PREFIX)
-            else:
-                print("ℹ️ 本地数据库也不存在，跳过恢复")
+            print("ℹ️ 未找到任何备份，跳过恢复")
             return
-
-        # 有备份文件：取最新的进行恢复
         db_files.sort()
-        latest = db_files[-1]
-        print(f"🔄 发现最新备份: {latest}，开始下载...")
-
+        latest_db = db_files[-1]
+        print(f"🔄 使用二进制备份恢复: {latest_db}")
         downloaded = hf_hub_download(
             repo_id=DATASET_REPO,
-            filename=latest,
+            filename=latest_db,
             repo_type="dataset",
             local_dir="/tmp"
         )
-
-        # 确保目标目录存在
         os.makedirs(os.path.dirname(DB_TARGET), exist_ok=True)
-
-        # 复制到目标位置
         shutil.copy(downloaded, DB_TARGET)
-        print(f"✅ 数据库恢复成功: {DB_TARGET} (来自备份 {latest})")
+        print(f"✅ 数据库从二进制备份恢复成功: {DB_TARGET}")
         os.remove(downloaded)
     except Exception as e:
         print(f"❌ 数据库恢复失败: {e}")