feat(ts-gen): display progress bar via websockets (#2194)

antarest/core/filesystem_blueprint.py

@@ -23,10 +23,11 @@
 
 import typing_extensions as te
 from fastapi import APIRouter, Depends, HTTPException
-from pydantic import BaseModel, Field
+from pydantic import Field
 from starlette.responses import PlainTextResponse, StreamingResponse
 
 from antarest.core.config import Config
+from antarest.core.serialization import AntaresBaseModel
 from antarest.core.utils.web import APITag
 from antarest.login.auth import Auth
 
@@ -35,7 +36,7 @@
 
 
 class FilesystemDTO(
-    BaseModel,
+    AntaresBaseModel,
     extra="forbid",
     json_schema_extra={
         "example": {
@@ -61,7 +62,7 @@ class FilesystemDTO(
 
 
 class MountPointDTO(
-    BaseModel,
+    AntaresBaseModel,
     extra="forbid",
     json_schema_extra={
         "example": {
@@ -109,7 +110,7 @@ async def from_path(cls, name: str, path: Path) -> "MountPointDTO":
 
 
 class FileInfoDTO(
-    BaseModel,
+    AntaresBaseModel,
     extra="forbid",
     json_schema_extra={
         "example": {

antarest/core/interfaces/eventbus.py

@@ -46,7 +46,7 @@ class EventType(StrEnum):
     WORKER_TASK_STARTED = "WORKER_TASK_STARTED"
     WORKER_TASK_ENDED = "WORKER_TASK_ENDED"
     LAUNCH_PROGRESS = "LAUNCH_PROGRESS"
-    TS_GENERATION_PROGRESS = "TS_GENERATION_PROGRESS"
+    TASK_PROGRESS = "TASK_PROGRESS"
 
 
 class EventChannelDirectory:
@@ -137,6 +137,9 @@ def start(self, threaded: bool = True) -> None:
 
 
 class DummyEventBusService(IEventBus):
+    def __init__(self) -> None:
+        self.events: List[Event] = []
+
     def queue(self, event: Event, queue: str) -> None:
         # Noop
         pass
@@ -150,7 +153,7 @@ def remove_queue_consumer(self, listener_id: str) -> None:
 
     def push(self, event: Event) -> None:
         # Noop
-        pass
+        self.events.append(event)
 
     def add_listener(
         self,

antarest/core/tasks/service.py

@@ -42,13 +42,23 @@
 
 logger = logging.getLogger(__name__)
 
-TaskUpdateNotifier = t.Callable[[str], None]
-Task = t.Callable[[TaskUpdateNotifier], TaskResult]
-
 DEFAULT_AWAIT_MAX_TIMEOUT = 172800  # 48 hours
 """Default timeout for `await_task` in seconds."""
 
 
+class ITaskNotifier(ABC):
+    @abstractmethod
+    def notify_message(self, message: str) -> None:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def notify_progress(self, progress: int) -> None:
+        raise NotImplementedError()
+
+
+Task = t.Callable[[ITaskNotifier], TaskResult]
+
+
 class ITaskService(ABC):
     @abstractmethod
     def add_worker_task(
@@ -94,11 +104,17 @@ def await_task(self, task_id: str, timeout_sec: int = DEFAULT_AWAIT_MAX_TIMEOUT)
 
 
 # noinspection PyUnusedLocal
-def noop_notifier(message: str) -> None:
-    """This function is used in tasks when no notification is required."""
+class NoopNotifier(ITaskNotifier):
+    """This class is used in tasks when no notification is required."""
+
+    def notify_message(self, message: str) -> None:
+        return
 
+    def notify_progress(self, progress: int) -> None:
+        return
 
-class TaskJobLogRecorder:
+
+class TaskLogAndProgressRecorder(ITaskNotifier):
     """
     Callback used to register log messages in the TaskJob table.
 
@@ -107,15 +123,32 @@ class TaskJobLogRecorder:
         session: The database session created in the same thread as the task thread.
     """
 
-    def __init__(self, task_id: str, session: Session):
+    def __init__(self, task_id: str, session: Session, event_bus: IEventBus) -> None:
         self.session = session
         self.task_id = task_id
+        self.event_bus = event_bus
 
-    def __call__(self, message: str) -> None:
+    def notify_message(self, message: str) -> None:
         task = self.session.query(TaskJob).get(self.task_id)
         if task:
             task.logs.append(TaskJobLog(message=message, task_id=self.task_id))
-            db.session.commit()
+            self.session.commit()
+
+    def notify_progress(self, progress: int) -> None:
+        self.session.query(TaskJob).filter(TaskJob.id == self.task_id).update({TaskJob.progress: progress})
+        self.session.commit()
+
+        self.event_bus.push(
+            Event(
+                type=EventType.TASK_PROGRESS,
+                payload={
+                    "task_id": self.task_id,
+                    "progress": progress,
+                },
+                permissions=PermissionInfo(public_mode=PublicMode.READ),
+                channel=EventChannelDirectory.TASK + self.task_id,
+            )
+        )
 
 
 class TaskJobService(ITaskService):
@@ -138,7 +171,7 @@ def _create_worker_task(
         task_id: str,
         task_type: str,
         task_args: t.Dict[str, t.Union[int, float, bool, str]],
-    ) -> t.Callable[[TaskUpdateNotifier], TaskResult]:
+    ) -> Task:
         task_result_wrapper: t.List[TaskResult] = []
 
         def _create_awaiter(
@@ -152,7 +185,7 @@ async def _await_task_end(event: Event) -> None:
             return _await_task_end
 
         # noinspection PyUnusedLocal
-        def _send_worker_task(logger_: TaskUpdateNotifier) -> TaskResult:
+        def _send_worker_task(logger_: ITaskNotifier) -> TaskResult:
             listener_id = self.event_bus.add_listener(
                 _create_awaiter(task_result_wrapper),
                 [EventType.WORKER_TASK_ENDED],
@@ -380,7 +413,7 @@ def _run_task(
         try:
             with db():
                 # We must use the DB session attached to the current thread
-                result = callback(TaskJobLogRecorder(task_id, session=db.session))
+                result = callback(TaskLogAndProgressRecorder(task_id, db.session, self.event_bus))
 
             status = TaskStatus.COMPLETED if result.success else TaskStatus.FAILED
             logger.info(f"Task {task_id} ended with status {status}")

antarest/launcher/service.py

@@ -33,7 +33,7 @@
 from antarest.core.model import PermissionInfo, PublicMode, StudyPermissionType
 from antarest.core.requests import RequestParameters, UserHasNotPermissionError
 from antarest.core.tasks.model import TaskResult, TaskType
-from antarest.core.tasks.service import ITaskService, TaskUpdateNotifier
+from antarest.core.tasks.service import ITaskNotifier, ITaskService
 from antarest.core.utils.archives import ArchiveFormat, archive_dir, is_zip, read_in_zip
 from antarest.core.utils.fastapi_sqlalchemy import db
 from antarest.core.utils.utils import StopWatch, concat_files, concat_files_to_str
@@ -598,7 +598,7 @@ def _download_fallback_output(self, job_id: str, params: RequestParameters) -> F
             export_path = Path(export_file_download.path)
             export_id = export_file_download.id
 
-            def export_task(_: TaskUpdateNotifier) -> TaskResult:
+            def export_task(_: ITaskNotifier) -> TaskResult:
                 try:
                     #
                     archive_dir(output_path, export_path, archive_format=ArchiveFormat.ZIP)

antarest/matrixstore/service.py

@@ -32,7 +32,7 @@
 from antarest.core.requests import RequestParameters, UserHasNotPermissionError
 from antarest.core.serialization import from_json
 from antarest.core.tasks.model import TaskResult, TaskType
-from antarest.core.tasks.service import ITaskService, TaskUpdateNotifier
+from antarest.core.tasks.service import ITaskNotifier, ITaskService
 from antarest.core.utils.archives import ArchiveFormat, archive_dir
 from antarest.core.utils.fastapi_sqlalchemy import db
 from antarest.core.utils.utils import StopWatch
@@ -510,7 +510,7 @@ def download_matrix_list(
         export_path = Path(export_file_download.path)
         export_id = export_file_download.id
 
-        def export_task(notifier: TaskUpdateNotifier) -> TaskResult:
+        def export_task(notifier: ITaskNotifier) -> TaskResult:
             try:
                 self.create_matrix_files(matrix_ids=matrix_list, export_path=export_path)
                 self.file_transfer_manager.set_ready(export_id)

antarest/study/business/utils.py

@@ -13,18 +13,19 @@
 import typing as t
 
 from antares.study.version import StudyVersion
-from pydantic import BaseModel
 
 from antarest.core.exceptions import CommandApplicationError
 from antarest.core.jwt import DEFAULT_ADMIN_USER
 from antarest.core.requests import RequestParameters
+from antarest.core.serialization import AntaresBaseModel
 from antarest.study.business.all_optional_meta import camel_case_model
 from antarest.study.model import RawStudy, Study
 from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy
 from antarest.study.storage.storage_service import StudyStorageService
 from antarest.study.storage.utils import is_managed
 from antarest.study.storage.variantstudy.business.utils import transform_command_to_dto
 from antarest.study.storage.variantstudy.model.command.icommand import ICommand
+from antarest.study.storage.variantstudy.model.command_listener.command_listener import ICommandListener
 
 # noinspection SpellCheckingInspection
 GENERAL_DATA_PATH = "settings/generaldata"
@@ -35,11 +36,12 @@ def execute_or_add_commands(
     file_study: FileStudy,
     commands: t.Sequence[ICommand],
     storage_service: StudyStorageService,
+    listener: t.Optional[ICommandListener] = None,
 ) -> None:
     if isinstance(study, RawStudy):
         executed_commands: t.MutableSequence[ICommand] = []
         for command in commands:
-            result = command.apply(file_study)
+            result = command.apply(file_study, listener)
             if not result.status:
                 raise CommandApplicationError(result.message)
             executed_commands.append(command)
@@ -72,7 +74,7 @@ def execute_or_add_commands(
 
 @camel_case_model
 class FormFieldsBaseModel(
-    BaseModel,
+    AntaresBaseModel,
     extra="forbid",
     validate_assignment=True,
     populate_by_name=True,