Implement RBAC for celery

Author: Ben Elam

orchestrator/services/celery.py

@@ -18,6 +18,7 @@
 import structlog
 from celery.result import AsyncResult
 from kombu.exceptions import ConnectionError, OperationalError
+from oauth2_lib.fastapi import OIDCUserModel
 
 from orchestrator import app_settings
 from orchestrator.api.error_handling import raise_status
@@ -42,7 +43,11 @@ def _block_when_testing(task_result: AsyncResult) -> None:
 
 
 def _celery_start_process(
-    workflow_key: str, user_inputs: list[State] | None, user: str = SYSTEM_USER, **kwargs: Any
+    workflow_key: str,
+    user_inputs: list[State] | None,
+    user: str = SYSTEM_USER,
+    user_model: OIDCUserModel | None = None,
+    **kwargs: Any
 ) -> UUID:
     """Client side call of Celery."""
     from orchestrator.services.tasks import NEW_TASK, NEW_WORKFLOW, get_celery_task
@@ -53,7 +58,7 @@ def _celery_start_process(
 
     task_name = NEW_TASK if workflow.target == Target.SYSTEM else NEW_WORKFLOW
     trigger_task = get_celery_task(task_name)
-    pstat = create_process(workflow_key, user_inputs, user)
+    pstat = create_process(workflow_key, user_inputs=user_inputs, user=user, user_model=user_model)
     try:
         result = trigger_task.delay(pstat.process_id, workflow_key, user)
         _block_when_testing(result)
@@ -70,6 +75,7 @@ def _celery_resume_process(
     *,
     user_inputs: list[State] | None,
     user: str | None,
+    user_model: OIDCUserModel | None = None,
     **kwargs: Any,
 ) -> UUID:
     """Client side call of Celery."""
@@ -87,7 +93,7 @@ def _celery_resume_process(
     store_input_state(pstat.process_id, user_inputs, "user_input")
     try:
         _celery_set_process_status_resumed(process)
-        result = trigger_task.delay(pstat.process_id, user)
+        result = trigger_task.delay(pstat.process_id, user, user_model=user_model)
         _block_when_testing(result)
 
         return pstat.process_id

orchestrator/services/processes.py

@@ -628,7 +628,7 @@ def resume_process(
         raise
 
     resume_func = get_execution_context()["resume"]
-    return resume_func(process, user_inputs=user_inputs, user=user, broadcast_func=broadcast_func)
+    return resume_func(process, user_inputs=user_inputs, user=user, user_model=user_model, broadcast_func=broadcast_func)
 
 
 def ensure_correct_callback_token(pstat: ProcessStat, *, token: str) -> None:

orchestrator/services/tasks.py

@@ -32,6 +32,7 @@
     thread_resume_process,
 )
 from orchestrator.types import BroadcastFunc
+from orchestrator.utils.auth import Authorizer
 from orchestrator.utils.json import json_dumps, json_loads
 from orchestrator.workflow import ProcessStat, ProcessStatus, Success, runwf
 from orchestrator.workflows import get_workflow
@@ -103,12 +104,12 @@ def start_process(process_id: UUID, workflow_key: str, state: dict[str, Any], us
         else:
             return process_id
 
-    def resume_process(process_id: UUID, user_inputs: list[State] | None, user: str) -> UUID | None:
+    def resume_process(process_id: UUID, user_inputs: list[State] | None, user: str, user_model: Authorizer | None = None) -> UUID | None:
         try:
             process = _get_process(process_id)
             ensure_correct_process_status(process_id, ProcessStatus.RESUMED)
             process_id = thread_resume_process(
-                process, user_inputs=user_inputs, user=user, broadcast_func=process_broadcast_fn
+                process, user_inputs=user_inputs, user=user, user_model=user_model, broadcast_func=process_broadcast_fn
             )
         except Exception as exc:
             local_logger.error("Worker failed to resume workflow", process_id=process_id, details=str(exc))
@@ -131,16 +132,16 @@ def new_workflow(process_id, workflow_key: str, user: str) -> UUID | None:
         return start_process(process_id, workflow_key, state=state, user=user)
 
     @celery_task(name=RESUME_TASK)  # type: ignore
-    def resume_task(process_id: UUID, user: str) -> UUID | None:
+    def resume_task(process_id: UUID, user: str, user_model: Authorizer | None = None) -> UUID | None:
         local_logger.info("Resume task", process_id=process_id)
         state = retrieve_input_state(process_id, "user_input").input_state
-        return resume_process(process_id, user_inputs=state, user=user)
+        return resume_process(process_id, user_inputs=state, user=user, user_model=user_model)
 
     @celery_task(name=RESUME_WORKFLOW)  # type: ignore
-    def resume_workflow(process_id: UUID, user: str) -> UUID | None:
+    def resume_workflow(process_id: UUID, user: str, user_model: Authorizer | None = None) -> UUID | None:
         local_logger.info("Resume workflow", process_id=process_id)
         state = retrieve_input_state(process_id, "user_input").input_state
-        return resume_process(process_id, user_inputs=state, user=user)
+        return resume_process(process_id, user_inputs=state, user=user, user_model=user_model)
 
 
 class CeleryJobWorkerStatus(WorkerStatus):