Merge branch 'enh/conversation-read' of github.com:odeimaiz/osparc-simcore into enh/conversation-read

Author: odeimaiz

api/specs/web-server/_long_running_tasks.py

@@ -15,6 +15,10 @@
 from simcore_service_webserver.tasks._controller._rest_exceptions import (
     _TO_HTTP_ERROR_MAP,
 )
+from simcore_service_webserver.tasks._controller._rest_schemas import (
+    TaskStreamQueryParams,
+    TaskStreamResponse,
+)
 
 router = APIRouter(
     prefix=f"/{API_VTAG}",
@@ -63,3 +67,14 @@ def get_async_job_result(
     _path_params: Annotated[_PathParam, Depends()],
 ):
     """Retrieves the result of a task"""
+
+
+@router.get(
+    "/tasks/{task_id}/stream",
+    response_model=Envelope[TaskStreamResponse],
+)
+def get_async_job_stream(
+    _path_params: Annotated[_PathParam, Depends()],
+    _query_params: Annotated[TaskStreamQueryParams, Depends()],
+):
+    """Retrieves the stream of a task"""

api/specs/web-server/_storage.py

@@ -4,12 +4,10 @@
 # pylint: disable=too-many-arguments
 
 
-from typing import Annotated, TypeAlias
+from typing import Annotated, Any, Final, TypeAlias
 
 from fastapi import APIRouter, Depends, Query, status
-from models_library.api_schemas_long_running_tasks.tasks import (
-    TaskGet,
-)
+from models_library.api_schemas_long_running_tasks.tasks import TaskGet
 from models_library.api_schemas_storage.storage_schemas import (
     FileLocation,
     FileMetaDataGet,
@@ -25,6 +23,7 @@
     BatchDeletePathsBodyParams,
     DataExportPost,
     ListPathsQueryParams,
+    SearchBodyParams,
     StorageLocationPathParams,
     StoragePathComputeSizeParams,
 )
@@ -220,14 +219,33 @@ async def is_completed_upload_file(
     """Returns state of upload completion"""
 
 
+_RESPONSES: Final[dict[int | str, dict[str, Any]]] = {
+    i.status_code: {"model": EnvelopedError} for i in _TO_HTTP_ERROR_MAP.values()
+}
+
+
 @router.post(
-    "/storage/locations/{location_id}/export-data",
+    "/storage/locations/{location_id}:export-data",
+    status_code=status.HTTP_202_ACCEPTED,
     response_model=Envelope[TaskGet],
     name="export_data",
     description="Export data",
-    responses={
-        i.status_code: {"model": EnvelopedError} for i in _TO_HTTP_ERROR_MAP.values()
-    },
+    responses=_RESPONSES,
 )
 async def export_data(export_data: DataExportPost, location_id: LocationID):
     """Trigger data export. Returns async job id for getting status and results"""
+
+
+@router.post(
+    "/storage/locations/{location_id}:search",
+    status_code=status.HTTP_202_ACCEPTED,
+    response_model=Envelope[TaskGet],
+    name="search",
+    description="Starts a files/folders search",
+    responses=_RESPONSES,
+)
+async def search(
+    _path: Annotated[StorageLocationPathParams, Depends()],
+    _body: SearchBodyParams,
+):
+    """Trigger search. Returns async job id for getting status and results"""

packages/celery-library/src/celery_library/backends/redis.py

@@ -1,7 +1,7 @@
 import contextlib
 import logging
 from dataclasses import dataclass
-from datetime import timedelta
+from datetime import UTC, datetime, timedelta
 from typing import TYPE_CHECKING, Final
 
 from models_library.progress_bar import ProgressReport
@@ -13,21 +13,38 @@
     Task,
     TaskKey,
     TaskStore,
+    TaskStreamItem,
 )
 from servicelib.redis import RedisClientSDK, handle_redis_returns_union_types
 
+_CELERY_TASK_DELIMTATOR: Final[str] = ":"
+
 _CELERY_TASK_PREFIX: Final[str] = "celery-task-"
 _CELERY_TASK_ID_KEY_ENCODING = "utf-8"
 _CELERY_TASK_SCAN_COUNT_PER_BATCH: Final[int] = 1000
-_CELERY_TASK_METADATA_KEY: Final[str] = "metadata"
+_CELERY_TASK_EXEC_METADATA_KEY: Final[str] = "exec-meta"
 _CELERY_TASK_PROGRESS_KEY: Final[str] = "progress"
 
+### Redis list to store streamed results
+_CELERY_TASK_STREAM_PREFIX: Final[str] = "celery-task-stream-"
+_CELERY_TASK_STREAM_EXPIRY: Final[timedelta] = timedelta(minutes=3)
+_CELERY_TASK_STREAM_METADATA: Final[str] = "meta"
+_CELERY_TASK_STREAM_DONE_KEY: Final[str] = "done"
+_CELERY_TASK_STREAM_LAST_UPDATE_KEY: Final[str] = "last_update"
 
 _logger = logging.getLogger(__name__)
 
 
 def _build_redis_task_key(task_key: TaskKey) -> str:
-    return _CELERY_TASK_PREFIX + task_key
+    return f"{_CELERY_TASK_PREFIX}{task_key}"
+
+
+def _build_redis_stream_key(task_key: TaskKey) -> str:
+    return f"{_CELERY_TASK_STREAM_PREFIX}{task_key}"
+
+
+def _build_redis_stream_meta_key(task_key: TaskKey) -> str:
+    return f"{_build_redis_stream_key(task_key)}{_CELERY_TASK_DELIMTATOR}{_CELERY_TASK_STREAM_METADATA}"
 
 
 @dataclass(frozen=True)
@@ -44,7 +61,7 @@ async def create_task(
         await handle_redis_returns_union_types(
             self._redis_client_sdk.redis.hset(
                 name=redis_key,
-                key=_CELERY_TASK_METADATA_KEY,
+                key=_CELERY_TASK_EXEC_METADATA_KEY,
                 value=execution_metadata.model_dump_json(),
             )
         )
@@ -57,7 +74,7 @@ async def get_task_metadata(self, task_key: TaskKey) -> ExecutionMetadata | None
         raw_result = await handle_redis_returns_union_types(
             self._redis_client_sdk.redis.hget(
                 _build_redis_task_key(task_key),
-                _CELERY_TASK_METADATA_KEY,
+                _CELERY_TASK_EXEC_METADATA_KEY,
             )
         )
         if not raw_result:
@@ -99,7 +116,7 @@ async def list_tasks(self, owner_metadata: OwnerMetadata) -> list[Task]:
         )
 
         keys: list[str] = []
-        pipeline = self._redis_client_sdk.redis.pipeline()
+        pipe = self._redis_client_sdk.redis.pipeline()
         async for key in self._redis_client_sdk.redis.scan_iter(
             match=search_key, count=_CELERY_TASK_SCAN_COUNT_PER_BATCH
         ):
@@ -110,9 +127,9 @@ async def list_tasks(self, owner_metadata: OwnerMetadata) -> list[Task]:
                 else key
             )
             keys.append(_key)
-            pipeline.hget(_key, _CELERY_TASK_METADATA_KEY)
+            pipe.hget(_key, _CELERY_TASK_EXEC_METADATA_KEY)
 
-        results = await pipeline.execute()
+        results = await pipe.execute()
 
         tasks = []
         for key, raw_metadata in zip(keys, results, strict=True):
@@ -153,6 +170,62 @@ async def task_exists(self, task_key: TaskKey) -> bool:
         assert isinstance(n, int)  # nosec
         return n > 0
 
+    async def push_task_stream_items(
+        self, task_key: TaskKey, *result: TaskStreamItem
+    ) -> None:
+        stream_key = _build_redis_stream_key(task_key)
+        stream_meta_key = _build_redis_stream_meta_key(task_key)
+
+        pipe = self._redis_client_sdk.redis.pipeline()
+        pipe.rpush(stream_key, *(r.model_dump_json(by_alias=True) for r in result))
+        pipe.hset(
+            stream_meta_key, mapping={"last_update": datetime.now(UTC).isoformat()}
+        )
+        pipe.expire(stream_key, _CELERY_TASK_STREAM_EXPIRY)
+        pipe.expire(stream_meta_key, _CELERY_TASK_STREAM_EXPIRY)
+        await pipe.execute()
+
+    async def set_task_stream_done(self, task_key: TaskKey) -> None:
+        stream_meta_key = _build_redis_stream_meta_key(task_key)
+        await handle_redis_returns_union_types(
+            self._redis_client_sdk.redis.hset(
+                name=stream_meta_key,
+                key=_CELERY_TASK_STREAM_DONE_KEY,
+                value="1",
+            )
+        )
+
+    async def pull_task_stream_items(
+        self, task_key: TaskKey, limit: int = 20
+    ) -> tuple[list[TaskStreamItem], bool, datetime | None]:
+        stream_key = _build_redis_stream_key(task_key)
+        meta_key = _build_redis_stream_meta_key(task_key)
+
+        async with self._redis_client_sdk.redis.pipeline(transaction=True) as pipe:
+            pipe.lpop(stream_key, limit)
+            pipe.hget(meta_key, _CELERY_TASK_STREAM_DONE_KEY)
+            pipe.hget(meta_key, _CELERY_TASK_STREAM_LAST_UPDATE_KEY)
+            raw_items, done, last_update = await pipe.execute()
+
+        stream_items = (
+            [TaskStreamItem.model_validate_json(item) for item in raw_items]
+            if raw_items
+            else []
+        )
+
+        empty = (
+            await handle_redis_returns_union_types(
+                self._redis_client_sdk.redis.llen(stream_key)
+            )
+            == 0
+        )
+
+        return (
+            stream_items,
+            done == "1" and empty,
+            datetime.fromisoformat(last_update) if last_update else None,
+        )
+
 
 if TYPE_CHECKING:
     _: type[TaskStore] = RedisTaskStore

packages/celery-library/src/celery_library/rpc/_async_jobs.py

@@ -2,7 +2,6 @@
 
 import logging
 
-from celery.exceptions import CeleryError  # type: ignore[import-untyped]
 from models_library.api_schemas_rpc_async_jobs.async_jobs import (
     AsyncJobGet,
     AsyncJobId,
@@ -22,6 +21,7 @@
 from servicelib.rabbitmq import RPCRouter
 
 from ..errors import (
+    TaskManagerError,
     TaskNotFoundError,
     TransferrableCeleryError,
     decode_celery_transferrable_error,
@@ -44,7 +44,7 @@ async def cancel(
         )
     except TaskNotFoundError as exc:
         raise JobMissingError(job_id=job_id) from exc
-    except CeleryError as exc:
+    except TaskManagerError as exc:
         raise JobSchedulerError(exc=f"{exc}") from exc
 
 
@@ -62,7 +62,7 @@ async def status(
         )
     except TaskNotFoundError as exc:
         raise JobMissingError(job_id=job_id) from exc
-    except CeleryError as exc:
+    except TaskManagerError as exc:
         raise JobSchedulerError(exc=f"{exc}") from exc
 
     return AsyncJobStatus(
@@ -101,7 +101,7 @@ async def result(
         )
     except TaskNotFoundError as exc:
         raise JobMissingError(job_id=job_id) from exc
-    except CeleryError as exc:
+    except TaskManagerError as exc:
         raise JobSchedulerError(exc=f"{exc}") from exc
 
     if _status.task_state == TaskState.FAILURE:
@@ -136,7 +136,7 @@ async def list_jobs(
         tasks = await task_manager.list_tasks(
             owner_metadata=owner_metadata,
         )
-    except CeleryError as exc:
+    except TaskManagerError as exc:
         raise JobSchedulerError(exc=f"{exc}") from exc
 
     return [