Coverage for src\core\storage.py: 0%

1""" 

2数据存储系统 

3提供统一的存储接口，支持多种存储后端 

4""" 

5 

6import json 

7import sqlite3 

8import aiofiles 

9from pathlib import Path 

10from typing import Dict, List, Optional, Any, Union, Tuple 

11from enum import Enum 

12import asyncio 

13from dataclasses import dataclass, asdict 

14from datetime import datetime 

15import time 

16 

17 

18class StorageType(str, Enum): 

19 """存储类型""" 

20 SQLITE = "sqlite" 

21 JSON_FILE = "json" 

22 MEMORY = "memory" 

23 

24 

25@dataclass 

26class StorageConfig: 

27 """存储配置""" 

28 

29 type: StorageType = StorageType.SQLITE 

30 path: Optional[Path] = None 

31 table_name: str = "memories" 

32 max_connections: int = 5 

33 auto_commit: bool = True 

34 

35 

36class StorageBackend: 

37 """存储后端基类""" 

38 

39 async def connect(self) -> None: 

40 """连接到存储""" 

41 pass 

42 

43 async def disconnect(self) -> None: 

44 """断开连接""" 

45 pass 

46 

47 async def save(self, key: str, data: Dict[str, Any]) -> bool: 

48 """保存数据""" 

49 raise NotImplementedError 

50 

51 async def load(self, key: str) -> Optional[Dict[str, Any]]: 

52 """加载数据""" 

53 raise NotImplementedError 

54 

55 async def delete(self, key: str) -> bool: 

56 """删除数据""" 

57 raise NotImplementedError 

58 

59 async def query(self, conditions: Dict[str, Any], limit: int = 100) -> List[Dict[str, Any]]: 

60 """查询数据""" 

61 raise NotImplementedError 

62 

63 async def exists(self, key: str) -> bool: 

64 """检查数据是否存在""" 

65 raise NotImplementedError 

66 

67 async def list_keys(self, prefix: str = "") -> List[str]: 

68 """列出所有键""" 

69 raise NotImplementedError 

70 

71 

72class SQLiteBackend(StorageBackend): 

73 """SQLite 存储后端""" 

74 

75 def __init__(self, config: StorageConfig): 

76 self.config = config 

77 self.connection: Optional[sqlite3.Connection] = None 

78 

79 # 确保路径存在 

80 if self.config.path: 

81 path = self.config.path if isinstance(self.config.path, Path) else Path(self.config.path) 

82 path.parent.mkdir(parents=True, exist_ok=True) 

83 

84 async def connect(self) -> None: 

85 """连接到 SQLite 数据库""" 

86 if self.config.path is None: 

87 raise ValueError("SQLite 存储需要指定路径") 

88 

89 # SQLite 连接是同步的，但我们在异步环境中使用 

90 self.connection = sqlite3.connect( 

91 str(self.config.path), 

92 check_same_thread=False 

93 ) 

94 

95 # 启用外键和 WAL 模式 

96 self.connection.execute("PRAGMA foreign_keys = ON") 

97 self.connection.execute("PRAGMA journal_mode = WAL") 

98 

99 # 创建表 

100 await self._create_tables() 

101 

102 async def disconnect(self) -> None: 

103 """断开连接""" 

104 if self.connection: 

105 self.connection.close() 

106 self.connection = None 

107 

108 async def save(self, key: str, data: Dict[str, Any]) -> bool: 

109 """保存数据到 SQLite""" 

110 if not self.connection: 

111 await self.connect() 

112 

113 try: 

114 # 序列化数据 

115 data_json = json.dumps(data, ensure_ascii=False) 

116 

117 # 插入或更新 

118 cursor = self.connection.execute( 

119 f""" 

120 INSERT OR REPLACE INTO {self.config.table_name}  

121 (key, data, created_at, updated_at)  

122 VALUES (?, ?, ?, ?) 

123 """, 

124 (key, data_json, time.time(), time.time()) 

125 ) 

126 

127 if self.config.auto_commit: 

128 self.connection.commit() 

129 

130 return cursor.rowcount > 0 

131 

132 except Exception as e: 

133 print(f"SQLite 保存失败: {e}") 

134 return False 

135 

136 async def load(self, key: str) -> Optional[Dict[str, Any]]: 

137 """从 SQLite 加载数据""" 

138 if not self.connection: 

139 await self.connect() 

140 

141 try: 

142 cursor = self.connection.execute( 

143 f"SELECT data FROM {self.config.table_name} WHERE key = ?", 

144 (key,) 

145 ) 

146 

147 row = cursor.fetchone() 

148 if row: 

149 return json.loads(row[0]) 

150 

151 return None 

152 

153 except Exception as e: 

154 print(f"SQLite 加载失败: {e}") 

155 return None 

156 

157 async def delete(self, key: str) -> bool: 

158 """从 SQLite 删除数据""" 

159 if not self.connection: 

160 await self.connect() 

161 

162 try: 

163 cursor = self.connection.execute( 

164 f"DELETE FROM {self.config.table_name} WHERE key = ?", 

165 (key,) 

166 ) 

167 

168 if self.config.auto_commit: 

169 self.connection.commit() 

170 

171 return cursor.rowcount > 0 

172 

173 except Exception as e: 

174 print(f"SQLite 删除失败: {e}") 

175 return False 

176 

177 async def query(self, conditions: Dict[str, Any], limit: int = 100) -> List[Dict[str, Any]]: 

178 """查询 SQLite 数据""" 

179 if not self.connection: 

180 await self.connect() 

181 

182 try: 

183 # 构建查询条件 

184 where_clauses = [] 

185 params = [] 

186 

187 for key, value in conditions.items(): 

188 where_clauses.append(f"json_extract(data, '$.{key}') = ?") 

189 params.append(value) 

190 

191 where_sql = " AND ".join(where_clauses) if where_clauses else "1=1" 

192 

193 # 执行查询 

194 cursor = self.connection.execute( 

195 f""" 

196 SELECT data FROM {self.config.table_name}  

197 WHERE {where_sql}  

198 ORDER BY updated_at DESC  

199 LIMIT ? 

200 """, 

201 (*params, limit) 

202 ) 

203 

204 results = [] 

205 for row in cursor.fetchall(): 

206 results.append(json.loads(row[0])) 

207 

208 return results 

209 

210 except Exception as e: 

211 print(f"SQLite 查询失败: {e}") 

212 return [] 

213 

214 async def exists(self, key: str) -> bool: 

215 """检查键是否存在""" 

216 if not self.connection: 

217 await self.connect() 

218 

219 try: 

220 cursor = self.connection.execute( 

221 f"SELECT 1 FROM {self.config.table_name} WHERE key = ?", 

222 (key,) 

223 ) 

224 

225 return cursor.fetchone() is not None 

226 

227 except Exception as e: 

228 print(f"SQLite 检查存在失败: {e}") 

229 return False 

230 

231 async def list_keys(self, prefix: str = "") -> List[str]: 

232 """列出所有键""" 

233 if not self.connection: 

234 await self.connect() 

235 

236 try: 

237 if prefix: 

238 cursor = self.connection.execute( 

239 f"SELECT key FROM {self.config.table_name} WHERE key LIKE ?", 

240 (f"{prefix}%",) 

241 ) 

242 else: 

243 cursor = self.connection.execute( 

244 f"SELECT key FROM {self.config.table_name}" 

245 ) 

246 

247 return [row[0] for row in cursor.fetchall()] 

248 

249 except Exception as e: 

250 print(f"SQLite 列出键失败: {e}") 

251 return [] 

252 

253 async def _create_tables(self) -> None: 

254 """创建数据库表""" 

255 if not self.connection: 

256 return 

257 

258 self.connection.execute(f""" 

259 CREATE TABLE IF NOT EXISTS {self.config.table_name} ( 

260 key TEXT PRIMARY KEY, 

261 data TEXT NOT NULL, 

262 created_at REAL NOT NULL, 

263 updated_at REAL NOT NULL 

264 ) 

265 """) 

266 

267 # 创建索引 

268 self.connection.execute(f""" 

269 CREATE INDEX IF NOT EXISTS idx_updated_at  

270 ON {self.config.table_name} (updated_at) 

271 """) 

272 

273 self.connection.commit() 

274 

275 

276class JSONFileBackend(StorageBackend): 

277 """JSON 文件存储后端""" 

278 

279 def __init__(self, config: StorageConfig): 

280 self.config = config 

281 self.data: Dict[str, Dict[str, Any]] = {} 

282 

283 # 确保路径存在 

284 if self.config.path: 

285 path = self.config.path if isinstance(self.config.path, Path) else Path(self.config.path) 

286 path.parent.mkdir(parents=True, exist_ok=True) 

287 

288 async def connect(self) -> None: 

289 """连接到 JSON 文件存储""" 

290 if self.config.path and self.config.path.exists(): 

291 await self._load_from_file() 

292 

293 async def disconnect(self) -> None: 

294 """断开连接并保存到文件""" 

295 if self.config.path: 

296 await self._save_to_file() 

297 

298 async def save(self, key: str, data: Dict[str, Any]) -> bool: 

299 """保存数据到 JSON 文件""" 

300 try: 

301 # 添加时间戳 

302 if "created_at" not in data: 

303 data["created_at"] = time.time() 

304 data["updated_at"] = time.time() 

305 

306 # 保存到内存 

307 self.data[key] = data 

308 

309 # 异步保存到文件 

310 asyncio.create_task(self._save_to_file()) 

311 

312 return True 

313 

314 except Exception as e: 

315 print(f"JSON 文件保存失败: {e}") 

316 return False 

317 

318 async def load(self, key: str) -> Optional[Dict[str, Any]]: 

319 """从 JSON 文件加载数据""" 

320 return self.data.get(key) 

321 

322 async def delete(self, key: str) -> bool: 

323 """从 JSON 文件删除数据""" 

324 try: 

325 if key in self.data: 

326 del self.data[key] 

327 

328 # 异步保存到文件 

329 asyncio.create_task(self._save_to_file()) 

330 

331 return True 

332 

333 return False 

334 

335 except Exception as e: 

336 print(f"JSON 文件删除失败: {e}") 

337 return False 

338 

339 async def query(self, conditions: Dict[str, Any], limit: int = 100) -> List[Dict[str, Any]]: 

340 """查询 JSON 文件数据""" 

341 results = [] 

342 

343 for item in self.data.values(): 

344 match = True 

345 

346 for key, value in conditions.items(): 

347 if key not in item or item[key] != value: 

348 match = False 

349 break 

350 

351 if match: 

352 results.append(item) 

353 

354 # 按更新时间排序 

355 results.sort(key=lambda x: x.get("updated_at", 0), reverse=True) 

356 

357 return results[:limit] 

358 

359 async def exists(self, key: str) -> bool: 

360 """检查键是否存在""" 

361 return key in self.data 

362 

363 async def list_keys(self, prefix: str = "") -> List[str]: 

364 """列出所有键""" 

365 if prefix: 

366 return [key for key in self.data.keys() if key.startswith(prefix)] 

367 

368 return list(self.data.keys()) 

369 

370 async def _load_from_file(self) -> None: 

371 """从文件加载数据""" 

372 if not self.config.path or not self.config.path.exists(): 

373 return 

374 

375 try: 

376 async with aiofiles.open(self.config.path, "r", encoding="utf-8") as f: 

377 content = await f.read() 

378 self.data = json.loads(content) 

379 

380 except Exception as e: 

381 print(f"加载 JSON 文件失败: {e}") 

382 self.data = {} 

383 

384 async def _save_to_file(self) -> None: 

385 """保存数据到文件""" 

386 if not self.config.path: 

387 return 

388 

389 try: 

390 async with aiofiles.open(self.config.path, "w", encoding="utf-8") as f: 

391 await f.write(json.dumps(self.data, indent=2, ensure_ascii=False)) 

392 

393 except Exception as e: 

394 print(f"保存 JSON 文件失败: {e}") 

395 

396 

397class MemoryBackend(StorageBackend): 

398 """内存存储后端（用于测试）""" 

399 

400 def __init__(self, config: StorageConfig): 

401 self.config = config 

402 self.data: Dict[str, Dict[str, Any]] = {} 

403 

404 async def save(self, key: str, data: Dict[str, Any]) -> bool: 

405 """保存数据到内存""" 

406 try: 

407 if "created_at" not in data: 

408 data["created_at"] = time.time() 

409 data["updated_at"] = time.time() 

410 

411 self.data[key] = data 

412 return True 

413 

414 except Exception as e: 

415 print(f"内存保存失败: {e}") 

416 return False 

417 

418 async def load(self, key: str) -> Optional[Dict[str, Any]]: 

419 """从内存加载数据""" 

420 return self.data.get(key) 

421 

422 async def delete(self, key: str) -> bool: 

423 """从内存删除数据""" 

424 try: 

425 if key in self.data: 

426 del self.data[key] 

427 return True 

428 

429 return False 

430 

431 except Exception as e: 

432 print(f"内存删除失败: {e}") 

433 return False 

434 

435 async def query(self, conditions: Dict[str, Any], limit: int = 100) -> List[Dict[str, Any]]: 

436 """查询内存数据""" 

437 results = [] 

438 

439 for item in self.data.values(): 

440 match = True 

441 

442 for key, value in conditions.items(): 

443 if key not in item or item[key] != value: 

444 match = False 

445 break 

446 

447 if match: 

448 results.append(item) 

449 

450 # 按更新时间排序 

451 results.sort(key=lambda x: x.get("updated_at", 0), reverse=True) 

452 

453 return results[:limit] 

454 

455 async def exists(self, key: str) -> bool: 

456 """检查键是否存在""" 

457 return key in self.data 

458 

459 async def list_keys(self, prefix: str = "") -> List[str]: 

460 """列出所有键""" 

461 if prefix: 

462 return [key for key in self.data.keys() if key.startswith(prefix)] 

463 

464 return list(self.data.keys()) 

465 

466 

467class StorageManager: 

468 """存储管理器""" 

469 

470 def __init__(self, config: Optional[StorageConfig] = None): 

471 self.config = config or StorageConfig() 

472 self.backend: Optional[StorageBackend] = None 

473 

474 async def initialize(self) -> None: 

475 """初始化存储""" 

476 # 根据配置创建后端 

477 if self.config.type == StorageType.SQLITE: 

478 self.backend = SQLiteBackend(self.config) 

479 

480 elif self.config.type == StorageType.JSON_FILE: 

481 self.backend = JSONFileBackend(self.config) 

482 

483 elif self.config.type == StorageType.MEMORY: 

484 self.backend = MemoryBackend(self.config) 

485 

486 else: 

487 raise ValueError(f"不支持的存储类型: {self.config.type}") 

488 

489 # 连接后端 

490 await self.backend.connect() 

491 

492 async def close(self) -> None: 

493 """关闭存储""" 

494 if self.backend: 

495 await self.backend.disconnect() 

496 

497 async def save(self, key: str, data: Dict[str, Any]) -> bool: 

498 """保存数据""" 

499 if not self.backend: 

500 await self.initialize() 

501 

502 return await self.backend.save(key, data) 

503 

504 async def load(self, key: str) -> Optional[Dict[str, Any]]: 

505 """加载数据""" 

506 if not self.backend: 

507 await self.initialize() 

508 

509 return await self.backend.load(key) 

510 

511 async def delete(self, key: str) -> bool: 

512 """删除数据""" 

513 if not self.backend: 

514 await self.initialize() 

515 

516 return await self.backend.delete(key) 

517 

518 async def query(self, conditions: Dict[str, Any], limit: int = 100) -> List[Dict[str, Any]]: 

519 """查询数据""" 

520 if not self.backend: 

521 await self.initialize() 

522 

523 return await self.backend.query(conditions, limit) 

524 

525 async def exists(self, key: str) -> bool: 

526 """检查数据是否存在""" 

527 if not self.backend: 

528 await self.initialize() 

529 

530 return await self.backend.exists(key) 

531 

532 async def list_keys(self, prefix: str = "") -> List[str]: 

533 """列出所有键""" 

534 if not self.backend: 

535 await self.initialize() 

536 

537 return await self.backend.list_keys(prefix) 

538 

539 async def migrate_data(self, source_backend: StorageBackend, target_backend: StorageBackend) -> int: 

540 """ 

541 迁移数据 

542  

543 Args: 

544 source_backend: 源存储后端 

545 target_backend: 目标存储后端 

546  

547 Returns: 

548 迁移的数据数量 

549 """ 

550 migrated_count = 0 

551 

552 try: 

553 # 获取所有键 

554 keys = await source_backend.list_keys() 

555 

556 for key in keys: 

557 # 加载数据 

558 data = await source_backend.load(key) 

559 if data: 

560 # 保存到目标 

561 if await target_backend.save(key, data): 

562 migrated_count += 1 

563 

564 return migrated_count 

565 

566 except Exception as e: 

567 print(f"数据迁移失败: {e}") 

568 return migrated_count