Coverage for src/lexigram/admin/services/background_jobs.py: 0%

170 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-24 23:18 +0800

1"""Background job queue for admin import/export operations. 

2 

3Wraps import and export services to run large operations asynchronously with 

4real-time progress tracking via :class:`ProgressTrackerProtocol`. 

5 

6Usage:: 

7 

8 from lexigram.admin.services.background_jobs import BackgroundJobService 

9 from lexigram.tasks.progress import InMemoryProgressTracker 

10 

11 tracker = InMemoryProgressTracker() 

12 job_svc = BackgroundJobService( 

13 progress_tracker=tracker, 

14 import_service=importer, 

15 export_service=exporter, 

16 ) 

17 

18 # Enqueue a CSV import — returns immediately with a job_id 

19 job_id = await job_svc.enqueue_import( 

20 resource_type="user", 

21 file_content=csv_bytes, 

22 file_format="csv", 

23 actor_id="admin-user", 

24 ) 

25 

26 # Poll progress 

27 status = await job_svc.get_status(job_id) 

28 # → {"job_id": "...", "percent": 42, "status": "running", ...} 

29 

30 # Enqueue CSV/JSON export 

31 job_id = await job_svc.enqueue_export( 

32 resource_type="user", 

33 export_format="csv", 

34 filters={"active": True}, 

35 actor_id="admin-user", 

36 ) 

37""" 

38 

39from __future__ import annotations 

40 

41import asyncio 

42from dataclasses import dataclass, field 

43from datetime import UTC, datetime 

44from typing import Any 

45 

46from lexigram.admin.services.export import ExportFormat 

47from lexigram.contracts.infra.tasks.progress import ProgressTrackerProtocol 

48from lexigram.logging import get_logger 

49from lexigram.serialization import dumps_str 

50 

51logger = get_logger(__name__) 

52 

53 

54class JobNotFoundError(KeyError): 

55 """Raised when a job_id is not found in the queue.""" 

56 

57 _code: str = "LEX_ERR_ADMIN_028" 

58 

59 

60@dataclass 

61class BackgroundJob: 

62 """Runtime record of a background operation. 

63 

64 Attributes: 

65 job_id: Unique identifier. 

66 job_type: ``"import"`` or ``"export"``. 

67 resource_type: Admin resource name. 

68 actor_id: Who triggered the job. 

69 status: ``"pending"``, ``"running"``, ``"completed"``, or ``"failed"``. 

70 percent: 0-100 completion percentage. 

71 created_at: UTC timestamp. 

72 completed_at: UTC timestamp (when finished). 

73 result: Final result data (on completion). 

74 error: Error message (on failure). 

75 metadata: Extra context (format, filename, etc.). 

76 """ 

77 

78 job_id: str 

79 job_type: str 

80 resource_type: str 

81 actor_id: str 

82 status: str = "pending" 

83 percent: int = 0 

84 created_at: datetime = field(default_factory=lambda: datetime.now(UTC)) 

85 completed_at: datetime | None = None 

86 result: dict[str, Any] = field(default_factory=dict) 

87 error: str = "" 

88 metadata: dict[str, Any] = field(default_factory=dict) 

89 

90 def to_dict(self) -> dict[str, Any]: 

91 """Serialise to a JSON-safe dict for SSE/polling endpoints.""" 

92 return { 

93 "job_id": self.job_id, 

94 "job_type": self.job_type, 

95 "resource_type": self.resource_type, 

96 "actor_id": self.actor_id, 

97 "status": self.status, 

98 "percent": self.percent, 

99 "created_at": self.created_at.isoformat(), 

100 "completed_at": self.completed_at.isoformat() 

101 if self.completed_at 

102 else None, 

103 "result": self.result, 

104 "error": self.error, 

105 "metadata": self.metadata, 

106 } 

107 

108 

109class BackgroundJobService: 

110 """Manages background import/export jobs with progress tracking. 

111 

112 All heavy work runs in a background :func:`asyncio.Task`. Progress is 

113 broadcast via :class:`ProgressTrackerProtocol` so subscribers can receive 

114 live updates and callers can poll :meth:`get_status` without blocking. 

115 

116 Args: 

117 progress_tracker: Optional :class:`ProgressTrackerProtocol` instance. 

118 When provided, jobs broadcast live updates so the admin progress 

119 UI works automatically. 

120 import_service: Optional ``AdminImportService`` instance. 

121 export_service: Optional ``ExportService`` instance. 

122 export_data_source: Optional ``IExportDataSource`` instance used by 

123 ``export_service`` to read rows for a queued export. 

124 max_retained_jobs: Maximum number of completed/failed jobs to keep 

125 in memory before oldest are evicted. 

126 """ 

127 

128 def __init__( 

129 self, 

130 progress_tracker: ProgressTrackerProtocol | None = None, 

131 import_service: Any = None, 

132 export_service: Any = None, 

133 export_data_source: Any | None = None, 

134 max_retained_jobs: int = 200, 

135 ) -> None: 

136 self._progress = progress_tracker 

137 self._importer = import_service 

138 self._exporter = export_service 

139 self._export_data_source = export_data_source 

140 self._max_retained = max_retained_jobs 

141 self._jobs: dict[str, BackgroundJob] = {} 

142 self._tasks: dict[str, asyncio.Task[None]] = {} 

143 self._background_tasks: set[asyncio.Task[None]] = set() 

144 self._counter = 0 

145 

146 # ------------------------------------------------------------------ 

147 # Internal helpers 

148 # ------------------------------------------------------------------ 

149 

150 def _new_id(self, prefix: str) -> str: 

151 self._counter += 1 

152 ts = datetime.now(UTC).strftime("%Y%m%d%H%M%S%f") 

153 return f"{prefix}-{ts}-{self._counter}" 

154 

155 def _register(self, job: BackgroundJob) -> None: 

156 self._jobs[job.job_id] = job 

157 self._evict_old() 

158 

159 def _evict_old(self) -> None: 

160 done = [j for j in self._jobs.values() if j.status in ("completed", "failed")] 

161 if len(done) > self._max_retained: 

162 for old in sorted(done, key=lambda j: j.created_at)[ 

163 : len(done) - self._max_retained 

164 ]: 

165 del self._jobs[old.job_id] 

166 

167 async def _update( 

168 self, job: BackgroundJob, percent: int, status: str = "running" 

169 ) -> None: 

170 job.percent = percent 

171 job.status = status 

172 if self._progress: 

173 try: 

174 await self._progress.update( 

175 job.job_id, percent, 100, f"{status} ({percent}%)" 

176 ) 

177 except (TimeoutError, OSError, RuntimeError): 

178 pass 

179 

180 async def _finish(self, job: BackgroundJob, result: dict[str, Any]) -> None: 

181 job.status = "completed" 

182 job.percent = 100 

183 job.result = result 

184 job.completed_at = datetime.now(UTC) 

185 if self._progress: 

186 try: 

187 await self._progress.complete(job.job_id, dumps_str(result)) 

188 except (TimeoutError, OSError, RuntimeError): 

189 pass 

190 

191 async def _fail(self, job: BackgroundJob, error: str) -> None: 

192 job.status = "failed" 

193 job.error = error 

194 job.completed_at = datetime.now(UTC) 

195 if self._progress: 

196 try: 

197 await self._progress.fail(job.job_id, error) 

198 except (TimeoutError, OSError, RuntimeError): 

199 pass 

200 

201 # ------------------------------------------------------------------ 

202 # Enqueue import 

203 # ------------------------------------------------------------------ 

204 

205 async def enqueue_import( 

206 self, 

207 resource_type: str, 

208 file_content: bytes, 

209 *, 

210 file_format: str = "csv", 

211 actor_id: str = "system", 

212 filename: str = "", 

213 options: dict[str, Any] | None = None, 

214 ) -> str: 

215 """Enqueue a background import job. 

216 

217 Args: 

218 resource_type: Resource name (e.g. ``"user"``). 

219 file_content: Raw file bytes. 

220 file_format: ``"csv"`` or ``"json"``. 

221 actor_id: Principal who triggered the import. 

222 filename: Optional original filename (stored in metadata). 

223 options: Extra options forwarded to the import service. 

224 

225 Returns: 

226 JobProtocol ID for polling via :meth:`get_status`. 

227 """ 

228 job_id = self._new_id("import") 

229 job = BackgroundJob( 

230 job_id=job_id, 

231 job_type="import", 

232 resource_type=resource_type, 

233 actor_id=actor_id, 

234 metadata={"file_format": file_format, "filename": filename}, 

235 ) 

236 self._register(job) 

237 

238 task = asyncio.create_task( 

239 self._run_import(job, file_content, file_format, options or {}) 

240 ) 

241 self._background_tasks.add(task) 

242 task.add_done_callback(self._background_tasks.discard) 

243 self._tasks[job_id] = task 

244 task.add_done_callback(lambda _t: self._tasks.pop(job_id, None)) 

245 logger.info("Enqueued import job %s for resource '%s'", job_id, resource_type) 

246 return job_id 

247 

248 async def _run_import( 

249 self, 

250 job: BackgroundJob, 

251 file_content: bytes, 

252 file_format: str, 

253 options: dict[str, Any], 

254 ) -> None: 

255 try: 

256 await self._update(job, 5) 

257 

258 if self._importer is None: 

259 # Simulate when no service wired (useful for testing) 

260 await asyncio.sleep(0) 

261 await self._finish( 

262 job, {"rows_imported": 0, "note": "no import service configured"} 

263 ) 

264 return 

265 

266 await self._update(job, 20) 

267 filename = job.metadata.get("filename") or f"upload.{file_format}" 

268 import_job = await self._importer.parse( 

269 file_content, filename=filename, **options 

270 ) 

271 if import_job.is_err(): 

272 await self._fail(job, str(import_job.unwrap_err())) 

273 return 

274 await self._update(job, 60) 

275 

276 result = await self._importer.commit(import_job.unwrap()) 

277 await self._update(job, 95) 

278 

279 if result.is_ok(): 

280 r = result.unwrap() 

281 await self._finish( 

282 job, 

283 { 

284 "rows_imported": getattr(r, "created", 0), 

285 "rows_failed": getattr(r, "failed", 0), 

286 }, 

287 ) 

288 else: 

289 err = result.unwrap_err() 

290 await self._fail(job, str(err)) 

291 

292 except Exception as exc: # noqa: BLE001 — top-level task runner must capture all failures to mark job as failed 

293 logger.exception("Import job %s failed: %s", job.job_id, exc) 

294 await self._fail(job, str(exc)) 

295 

296 # ------------------------------------------------------------------ 

297 # Enqueue export 

298 # ------------------------------------------------------------------ 

299 

300 async def enqueue_export( 

301 self, 

302 resource_type: str, 

303 *, 

304 export_format: str = "csv", 

305 actor_id: str = "system", 

306 filters: dict[str, Any] | None = None, 

307 columns: list[str] | None = None, 

308 ) -> str: 

309 """Enqueue a background export job. 

310 

311 Args: 

312 resource_type: Resource name. 

313 export_format: ``"csv"``, ``"json"``, or ``"xlsx"``. 

314 actor_id: Principal who triggered the export. 

315 filters: Optional filters to apply to the query. 

316 columns: Optional list of columns to include. 

317 

318 Returns: 

319 JobProtocol ID for polling via :meth:`get_status`. 

320 """ 

321 job_id = self._new_id("export") 

322 job = BackgroundJob( 

323 job_id=job_id, 

324 job_type="export", 

325 resource_type=resource_type, 

326 actor_id=actor_id, 

327 metadata={"export_format": export_format, "columns": columns or []}, 

328 ) 

329 self._register(job) 

330 

331 task = asyncio.create_task( 

332 self._run_export(job, export_format, filters or {}, columns or []) 

333 ) 

334 self._background_tasks.add(task) 

335 task.add_done_callback(self._background_tasks.discard) 

336 self._tasks[job_id] = task 

337 task.add_done_callback(lambda _t: self._tasks.pop(job_id, None)) 

338 logger.info("Enqueued export job %s for resource '%s'", job_id, resource_type) 

339 return job_id 

340 

341 async def _run_export( 

342 self, 

343 job: BackgroundJob, 

344 export_format: str, 

345 filters: dict[str, Any], 

346 columns: list[str], 

347 ) -> None: 

348 try: 

349 await self._update(job, 10) 

350 

351 if self._exporter is None: 

352 await asyncio.sleep(0) 

353 await self._finish( 

354 job, {"rows_exported": 0, "note": "no export service configured"} 

355 ) 

356 return 

357 

358 if self._export_data_source is None: 

359 await self._fail( 

360 job, 

361 "Export requires a data source; pass export_data_source to " 

362 "BackgroundJobService.", 

363 ) 

364 return 

365 

366 await self._update(job, 30) 

367 job_id = self._exporter.create_job( 

368 resource_name=job.resource_type, 

369 file_format=ExportFormat(export_format), 

370 filters=filters, 

371 columns=columns, 

372 user_id=job.actor_id, 

373 ) 

374 result = await self._exporter.execute_export( 

375 job_id, self._export_data_source 

376 ) 

377 await self._update(job, 90) 

378 

379 if result.is_ok(): 

380 output = result.unwrap() 

381 await self._finish( 

382 job, 

383 { 

384 "rows_exported": getattr(output, "total_records", 0), 

385 "download_url": getattr(output, "download_url", "") or "", 

386 "file_path": getattr(output, "file_path", "") or "", 

387 "export_format": export_format, 

388 }, 

389 ) 

390 else: 

391 await self._fail(job, str(result.unwrap_err())) 

392 

393 except Exception as exc: # noqa: BLE001 — top-level task runner must capture all failures to mark job as failed 

394 logger.exception("Export job %s failed: %s", job.job_id, exc) 

395 await self._fail(job, str(exc)) 

396 

397 # ------------------------------------------------------------------ 

398 # Status / query 

399 # ------------------------------------------------------------------ 

400 

401 async def get_status(self, job_id: str) -> dict[str, Any] | None: 

402 """Return current job status as a JSON-safe dict. 

403 

404 Args: 

405 job_id: JobProtocol identifier returned by :meth:`enqueue_import` / 

406 :meth:`enqueue_export`. 

407 

408 Returns: 

409 Status dict or ``None`` if job_id not found. 

410 """ 

411 job = self._jobs.get(job_id) 

412 return job.to_dict() if job else None 

413 

414 async def get_job(self, job_id: str) -> BackgroundJob | None: 

415 """Return the :class:`BackgroundJob` for *job_id*, or ``None``.""" 

416 return self._jobs.get(job_id) 

417 

418 def list_jobs( 

419 self, 

420 *, 

421 resource_type: str | None = None, 

422 job_type: str | None = None, 

423 status: str | None = None, 

424 ) -> list[BackgroundJob]: 

425 """Return jobs matching the given filters. 

426 

427 Args: 

428 resource_type: Filter by resource type. 

429 job_type: ``"import"`` or ``"export"``. 

430 status: Filter by status string. 

431 """ 

432 jobs = list(self._jobs.values()) 

433 if resource_type: 

434 jobs = [j for j in jobs if j.resource_type == resource_type] 

435 if job_type: 

436 jobs = [j for j in jobs if j.job_type == job_type] 

437 if status: 

438 jobs = [j for j in jobs if j.status == status] 

439 return sorted(jobs, key=lambda j: j.created_at, reverse=True) 

440 

441 async def cancel(self, job_id: str) -> bool: 

442 """Cancel a pending or running job. 

443 

444 Args: 

445 job_id: JobProtocol to cancel. 

446 

447 Returns: 

448 ``True`` if the job was cancelled, ``False`` if not found or 

449 already finished. 

450 """ 

451 job = self._jobs.get(job_id) 

452 if job is None or job.status in ("completed", "failed"): 

453 return False 

454 

455 task = self._tasks.get(job_id) 

456 if task and not task.done(): 

457 task.cancel() 

458 

459 job.status = "failed" 

460 job.error = "Cancelled by user" 

461 job.completed_at = datetime.now(UTC) 

462 return True 

463 

464 

465__all__ = [ 

466 "BackgroundJob", 

467 "BackgroundJobService", 

468 "JobNotFoundError", 

469]