"""Durable task status management with a test-only in-memory fallback. The Flask application configures a SQLAlchemy session factory at startup. Code that uses TaskManager outside an application (small unit tests and legacy adapters) keeps the old in-memory behavior, but production requests do not. """ from __future__ import annotations import threading import uuid from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from enum import Enum from typing import Any, Dict, Optional, cast from flask import current_app, has_app_context from sqlalchemy import delete, select from ..models.operations import Job, JobStatus from ..utils.locale import t class TaskStatus(str, Enum): PENDING = "pending" PROCESSING = "processing" COMPLETED = "completed" FAILED = "failed" @dataclass class Task: task_id: str task_type: str status: TaskStatus created_at: datetime updated_at: datetime progress: int = 0 message: str = "" result: Optional[Dict] = None error: Optional[str] = None metadata: Dict = field(default_factory=dict) progress_detail: Dict = field(default_factory=dict) def to_dict(self) -> Dict[str, Any]: return { "task_id": self.task_id, "task_type": self.task_type, "status": self.status.value, "created_at": self.created_at.isoformat(), "updated_at": self.updated_at.isoformat(), "progress": self.progress, "message": self.message, "progress_detail": self.progress_detail, "result": self.result, "error": self.error, "metadata": self.metadata, } class TaskManager: """Thread-safe task facade bound to one app/session factory.""" _configured_session_factory = None _config_lock = threading.Lock() _fallback_tasks: Dict[str, Task] = {} _fallback_lock = threading.Lock() def __init__(self, session_factory=None): """Bind this manager to an explicit or current-app session factory. A manager created during a request/app context keeps that app's factory for background work, but it cannot be reused inside a different Flask app. The explicit class configuration remains only for legacy tests and callers that run outside Flask. """ self._bound_app = None if has_app_context(): app = cast(Any, current_app)._get_current_object() current_factory = app.extensions.get("crowdsight_session_factory") if not callable(current_factory): raise RuntimeError("task_session_factory_required") if session_factory is not None and session_factory is not current_factory: raise RuntimeError("task_session_factory_mismatch") bound_factory = current_factory self._bound_app = app elif session_factory is not None: bound_factory = session_factory else: # No Flask app context: use the explicit legacy/test binding. bound_factory = type(self)._configured_session_factory self._session_factory = bound_factory self._tasks = type(self)._fallback_tasks self._task_lock = type(self)._fallback_lock @classmethod def configure(cls, session_factory) -> None: """Set an explicit outside-Flask binding for tests/legacy adapters.""" with cls._config_lock: cls._configured_session_factory = session_factory with cls._fallback_lock: cls._fallback_tasks.clear() def _factory(self): if not has_app_context(): return self._session_factory app = cast(Any, current_app)._get_current_object() current_factory = app.extensions.get("crowdsight_session_factory") if not callable(current_factory): raise RuntimeError("task_session_factory_required") if self._bound_app is not None and self._bound_app is not app: raise RuntimeError("task_app_context_mismatch") if self._session_factory is not None and self._session_factory is not current_factory: raise RuntimeError("task_session_factory_mismatch") if self._session_factory is None: self._session_factory = current_factory self._bound_app = app return self._session_factory @staticmethod def _bounded_text(value: Any, limit: int = 4000) -> str: if value is None: return "" return str(value)[:limit] @staticmethod def _job_status(status: TaskStatus | str | None) -> str | None: if status is None: return None value = status.value if isinstance(status, TaskStatus) else str(status) return { TaskStatus.PENDING.value: JobStatus.QUEUED.value, TaskStatus.PROCESSING.value: JobStatus.RUNNING.value, TaskStatus.COMPLETED.value: JobStatus.SUCCEEDED.value, TaskStatus.FAILED.value: JobStatus.FAILED.value, JobStatus.QUEUED.value: JobStatus.QUEUED.value, JobStatus.RUNNING.value: JobStatus.RUNNING.value, JobStatus.SUCCEEDED.value: JobStatus.SUCCEEDED.value, JobStatus.FAILED.value: JobStatus.FAILED.value, JobStatus.CANCELLED.value: JobStatus.CANCELLED.value, }.get(value) @staticmethod def _task_status(status: str | JobStatus) -> TaskStatus: value = status.value if isinstance(status, JobStatus) else str(status) return { JobStatus.QUEUED.value: TaskStatus.PENDING, JobStatus.RUNNING.value: TaskStatus.PROCESSING, JobStatus.SUCCEEDED.value: TaskStatus.COMPLETED, JobStatus.FAILED.value: TaskStatus.FAILED, JobStatus.CANCELLED.value: TaskStatus.FAILED, }.get(value, TaskStatus.FAILED) @classmethod def _from_job(cls, job: Job) -> Task: metadata = job.job_metadata if isinstance(job.job_metadata, dict) else {} result = job.result if isinstance(job.result, dict) else job.result detail = job.progress_detail if isinstance(job.progress_detail, dict) else {} return Task( task_id=job.id, task_type=job.operation, status=cls._task_status(job.status), created_at=job.created_at, updated_at=job.updated_at, progress=job.progress, message=job.message, result=result, error=job.error_code, metadata=metadata, progress_detail=detail, ) def create_task(self, task_type: str, metadata: Optional[Dict] = None) -> str: metadata = metadata or {} factory = self._factory() if factory is not None: organization_id = metadata.get("organization_id") if not isinstance(organization_id, str) or not organization_id: raise ValueError("task_scope_required") with factory() as session: job = Job( organization_id=organization_id, owner_user_id=metadata.get("owner_user_id"), project_id=metadata.get("project_id"), graph_id=metadata.get("graph_id"), operation=self._bounded_text(task_type, 120), status=JobStatus.QUEUED.value, job_metadata=metadata, progress_detail={}, ) session.add(job) session.commit() return job.id task_id = str(uuid.uuid4()) now = datetime.now(timezone.utc) task = Task( task_id=task_id, task_type=task_type, status=TaskStatus.PENDING, created_at=now, updated_at=now, metadata=metadata, ) with self._task_lock: self._tasks[task_id] = task return task_id def get_task( self, task_id: str, *, organization_id: Optional[str] = None, owner_user_id: Optional[str] = None, ) -> Optional[Task]: factory = self._factory() if factory is not None: with factory() as session: statement = select(Job).where(Job.id == task_id) if organization_id is not None: statement = statement.where(Job.organization_id == organization_id) if owner_user_id is not None: statement = statement.where(Job.owner_user_id == owner_user_id) job = session.scalar(statement) return self._from_job(job) if job is not None else None with self._task_lock: task = self._tasks.get(task_id) if task is None: return None if organization_id is not None and task.metadata.get("organization_id") != organization_id: return None if owner_user_id is not None and task.metadata.get("owner_user_id") != owner_user_id: return None return task def update_task( self, task_id: str, status: Optional[TaskStatus] = None, progress: Optional[int] = None, message: Optional[str] = None, result: Optional[Dict] = None, error: Optional[str] = None, progress_detail: Optional[Dict] = None, ): factory = self._factory() if factory is not None: with factory() as session: job = session.get(Job, task_id) if job is None: return mapped_status = self._job_status(status) if mapped_status is not None: job.status = mapped_status if mapped_status in { JobStatus.SUCCEEDED.value, JobStatus.FAILED.value, JobStatus.CANCELLED.value, }: job.finished_at = datetime.now(timezone.utc) if progress is not None: job.progress = min(max(int(progress), 0), 100) if message is not None: job.message = self._bounded_text(message) if result is not None: job.result = result if error is not None: job.error_code = self._bounded_text(error, 120) if progress_detail is not None: job.progress_detail = progress_detail job.updated_at = datetime.now(timezone.utc) session.commit() return with self._task_lock: task = self._tasks.get(task_id) if task: task.updated_at = datetime.now(timezone.utc) if status is not None: task.status = status if progress is not None: task.progress = min(max(int(progress), 0), 100) if message is not None: task.message = message if result is not None: task.result = result if error is not None: task.error = error if progress_detail is not None: task.progress_detail = progress_detail def complete_task(self, task_id: str, result: Dict): self.update_task( task_id, status=TaskStatus.COMPLETED, progress=100, message=t("progress.taskComplete"), result=result, ) def fail_task(self, task_id: str, error: str): self.update_task( task_id, status=TaskStatus.FAILED, message=t("progress.taskFailed"), error=error, ) def list_tasks( self, task_type: Optional[str] = None, *, organization_id: Optional[str] = None, owner_user_id: Optional[str] = None, ) -> list: factory = self._factory() if factory is not None: with factory() as session: statement = select(Job).order_by(Job.created_at.desc()) if task_type: statement = statement.where(Job.operation == task_type) if organization_id is not None: statement = statement.where(Job.organization_id == organization_id) if owner_user_id is not None: statement = statement.where(Job.owner_user_id == owner_user_id) jobs = session.scalars(statement.limit(100)).all() return [self._from_job(job) for job in jobs] with self._task_lock: tasks = list(self._tasks.values()) if task_type: tasks = [task for task in tasks if task.task_type == task_type] if organization_id is not None: tasks = [task for task in tasks if task.metadata.get("organization_id") == organization_id] if owner_user_id is not None: tasks = [task for task in tasks if task.metadata.get("owner_user_id") == owner_user_id] return [task for task in sorted(tasks, key=lambda item: item.created_at, reverse=True)] def cleanup_old_tasks(self, max_age_hours: int = 24): cutoff = datetime.now(timezone.utc) - timedelta(hours=max_age_hours) factory = self._factory() if factory is not None: with factory() as session: session.execute( delete(Job).where( Job.created_at < cutoff, Job.status.in_([JobStatus.SUCCEEDED.value, JobStatus.FAILED.value]), ) ) session.commit() return with self._task_lock: old_ids = [ task_id for task_id, task in self._tasks.items() if task.created_at < cutoff and task.status in [TaskStatus.COMPLETED, TaskStatus.FAILED] ] for task_id in old_ids: del self._tasks[task_id]