import json import os from pathlib import Path from typing import Dict, Any, Optional from datetime import datetime import threading class DataStore: """数据持久化存储服务""" def __init__(self, storage_dir: str = "data"): self.storage_dir = Path(storage_dir) self.storage_dir.mkdir(exist_ok=True) # 数据文件路径 self.data_file = self.storage_dir / "reports_data.json" # 内存缓存 self._cache: Dict[str, Any] = {} # 线程锁,防止并发写入冲突 self._lock = threading.Lock() # 启动时加载数据 self._load_data() def _load_data(self): """从文件加载数据""" if self.data_file.exists(): try: with open(self.data_file, 'r', encoding='utf-8') as f: self._cache = json.load(f) print(f"✓ 成功加载 {len(self._cache)} 份报告数据") except Exception as e: print(f"⚠ 加载数据失败: {e},将使用空数据") self._cache = {} else: print("✓ 数据文件不存在,将创建新文件") self._cache = {} def _save_data(self): """保存数据到文件""" try: with self._lock: # 创建临时文件,避免写入过程中断导致数据损坏 temp_file = self.data_file.with_suffix('.json.tmp') with open(temp_file, 'w', encoding='utf-8') as f: json.dump(self._cache, f, ensure_ascii=False, indent=2) # 原子性替换 temp_file.replace(self.data_file) except Exception as e: print(f"⚠ 保存数据失败: {e}") def get_all(self) -> Dict[str, Any]: """获取所有报告数据""" return self._cache.copy() def get(self, file_id: str) -> Optional[Dict[str, Any]]: """获取单个报告数据""" return self._cache.get(file_id) def set(self, file_id: str, data: Dict[str, Any]) -> None: """设置/更新报告数据""" self._cache[file_id] = data self._save_data() def update(self, file_id: str, updates: Dict[str, Any]) -> None: """更新报告数据的部分字段""" if file_id in self._cache: self._cache[file_id].update(updates) self._save_data() else: raise KeyError(f"报告 {file_id} 不存在") def delete(self, file_id: str) -> None: """删除报告数据""" if file_id in self._cache: del self._cache[file_id] self._save_data() def exists(self, file_id: str) -> bool: """检查报告是否存在""" return file_id in self._cache def count(self) -> int: """获取报告总数""" return len(self._cache) def cleanup_orphaned_files(self, upload_dir: Path) -> int: """清理孤立的文件(数据库中有记录但文件不存在)""" cleaned = 0 orphaned_ids = [] for file_id, report in self._cache.items(): filepath = report.get('filepath') if filepath and not os.path.exists(filepath): orphaned_ids.append(file_id) for file_id in orphaned_ids: self.delete(file_id) cleaned += 1 if cleaned > 0: print(f"✓ 清理了 {cleaned} 条孤立记录") return cleaned def export_backup(self, backup_path: Optional[str] = None) -> str: """导出备份""" if backup_path is None: timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") backup_path = self.storage_dir / f"backup_{timestamp}.json" else: backup_path = Path(backup_path) with open(backup_path, 'w', encoding='utf-8') as f: json.dump(self._cache, f, ensure_ascii=False, indent=2) print(f"✓ 数据已备份到: {backup_path}") return str(backup_path) def import_backup(self, backup_path: str) -> None: """从备份恢复数据""" backup_path = Path(backup_path) if not backup_path.exists(): raise FileNotFoundError(f"备份文件不存在: {backup_path}") with open(backup_path, 'r', encoding='utf-8') as f: imported_data = json.load(f) self._cache.update(imported_data) self._save_data() print(f"✓ 已从备份恢复 {len(imported_data)} 份报告")