Fix bug and linting issues

Author: Ben Elam

orchestrator/api/api_v1/endpoints/processes.py

@@ -101,27 +101,27 @@ def get_auth_callbacks(pstat: ProcessStat) -> tuple[Authorizer | None, Authorize
     # Default to workflow start callbacks
     auth_resume = pstat.workflow.authorize_callback
     # Retry defaults to the workflow start callback if not otherwise specified
-    auth_retry = pstat.workflow.retry_auth_callback if pstat.workflow.retry_auth_callback else auth_resume
+    auth_retry = pstat.workflow.retry_auth_callback if pstat.workflow.retry_auth_callback is not None else auth_resume
     # Iterate over previous steps to look for policy changes
     remaining_steps = pstat.log
-    past_steps = pstat.workflow.steps[:-len(remaining_steps)]
+    past_steps = pstat.workflow.steps[: -len(remaining_steps)]
     # Review last steps and current step
     for step in past_steps + [pstat.log[0]]:
         if step.resume_auth_callback and step.retry_auth_callback is None:
             # Set both to authorize_callback
             auth_resume = step.resume_auth_callback
             auth_retry = step.resume_auth_callback
             continue
-        elif step.resume_auth_callback and step.retry_auth_callback:
+        if step.resume_auth_callback and step.retry_auth_callback:
             auth_resume = step.resume_auth_callback
             auth_retry = step.retry_auth_callback
             continue
-        elif step.resume_auth_callback is None and step.retry_auth_callback:
+        if step.resume_auth_callback is None and step.retry_auth_callback:
             # Only update retry
             auth_retry = step.retry_auth_callback
             continue
-        else: # Both None
-            continue
+        # Both None
+        continue
     return auth_resume, auth_retry
 
 
@@ -209,13 +209,12 @@ def resume_process_endpoint(
     pstat = load_process(process)
     auth_resume, auth_retry = get_auth_callbacks(pstat)
     if process.last_status == ProcessStatus.SUSPENDED:
-        if not auth_resume(user_model):
+        if auth_resume is not None and not auth_resume(user_model):
             raise_status(HTTPStatus.FORBIDDEN, "User is not authorized to resume step")
     elif process.last_status == ProcessStatus.FAILED:
-        if not auth_retry(user_model):
+        if auth_retry is not None and not auth_retry(user_model):
             raise_status(HTTPStatus.FORBIDDEN, "User is not authorized to retry step")
 
-
     broadcast_invalidate_status_counts()
     broadcast_func = api_broadcast_process_data(request)
 

orchestrator/services/celery.py

@@ -18,8 +18,8 @@
 import structlog
 from celery.result import AsyncResult
 from kombu.exceptions import ConnectionError, OperationalError
-from oauth2_lib.fastapi import OIDCUserModel
 
+from oauth2_lib.fastapi import OIDCUserModel
 from orchestrator import app_settings
 from orchestrator.api.error_handling import raise_status
 from orchestrator.db import ProcessTable, db
@@ -47,7 +47,7 @@ def _celery_start_process(
     user_inputs: list[State] | None,
     user: str = SYSTEM_USER,
     user_model: OIDCUserModel | None = None,
-    **kwargs: Any
+    **kwargs: Any,
 ) -> UUID:
     """Client side call of Celery."""
     from orchestrator.services.tasks import NEW_TASK, NEW_WORKFLOW, get_celery_task

orchestrator/services/processes.py

@@ -42,7 +42,6 @@
 from orchestrator.workflow import (
     CALLBACK_TOKEN_KEY,
     DEFAULT_CALLBACK_PROGRESS_KEY,
-    Authorizer,
     Failed,
     ProcessStat,
     ProcessStatus,
@@ -575,7 +574,7 @@ def resume_process(
         raise
 
     resume_func = get_execution_context()["resume"]
-    return resume_func(process, user_inputs=user_inputs, user=user, user_model=user_model, broadcast_func=broadcast_func)
+    return resume_func(process, user_inputs=user_inputs, user=user, broadcast_func=broadcast_func)
 
 
 def ensure_correct_callback_token(pstat: ProcessStat, *, token: str) -> None:

orchestrator/workflow.py

@@ -13,21 +13,19 @@
 
 
 from __future__ import annotations
+
 import contextvars
 import functools
 import inspect
 import secrets
 from collections.abc import Callable
 from dataclasses import asdict, dataclass
-from functools import update_wrapper
-from http import HTTPStatus
 from itertools import dropwhile
 from typing import (
     Any,
     Generic,
     NoReturn,
     Protocol,
-    TypeAlias,
     TypeVar,
     cast,
     overload,
@@ -42,7 +40,6 @@
 
 from nwastdlib import const, identity
 from oauth2_lib.fastapi import OIDCUserModel
-from orchestrator.api.error_handling import raise_status
 from orchestrator.config.assignee import Assignee
 from orchestrator.db import db, transactional
 from orchestrator.services.settings import get_engine_settings
@@ -181,7 +178,7 @@ def __repr__(self) -> str:
 
 
 def _handle_simple_input_form_generator(f: StateInputStepFunc) -> StateInputFormGenerator:
-    """Processes f into a form generator and injects a pre-hook for user authorization"""
+    """Processes f into a form generator and injects a pre-hook for user authorization."""
     if inspect.isgeneratorfunction(f):
         return cast(StateInputFormGenerator, f)
     if inspect.isgenerator(f):
@@ -219,7 +216,9 @@ def wrapping_function() -> NoReturn:
     wrapping_function.description = description
     wrapping_function.authorize_callback = allow if authorize_callback is None else authorize_callback
     # If no retry auth policy is given, defer to policy for process creation.
-    wrapping_function.retry_auth_callback = wrapping_function.authorize_callback if retry_auth_callback is None else retry_auth_callback
+    wrapping_function.retry_auth_callback = (
+        wrapping_function.authorize_callback if retry_auth_callback is None else retry_auth_callback
+    )
 
     if initial_input_form is None:
         # We always need a form to prevent starting a workflow when no input is needed.
@@ -288,9 +287,16 @@ def wrapper(state: State) -> Process:
     return decorator
 
 
-def inputstep(name: str, assignee: Assignee, resume_auth_callback: Authorizer | None = None, retry_auth_callback: Authorizer | None = None) -> Callable[[InputStepFunc], Step]:
+def inputstep(
+    name: str,
+    assignee: Assignee,
+    resume_auth_callback: Authorizer | None = None,
+    retry_auth_callback: Authorizer | None = None,
+) -> Callable[[InputStepFunc], Step]:
     """Add user input step to workflow.
 
+    Any authorization callbacks will be attached to the resulting Step.
+
     IMPORTANT: In contrast to other workflow steps, the `@inputstep` wrapped function will not run in the
     workflow engine! This means that it must be free of side effects!
 
@@ -317,7 +323,14 @@ def wrapper(state: State) -> FormGenerator:
         def suspend(state: State) -> Process:
             return Suspend(state)
 
-        return make_step_function(suspend, name, wrapper, assignee, resume_auth_callback=resume_auth_callback, retry_auth_callback=retry_auth_callback)
+        return make_step_function(
+            suspend,
+            name,
+            wrapper,
+            assignee,
+            resume_auth_callback=resume_auth_callback,
+            retry_auth_callback=retry_auth_callback,
+        )
 
     return decorator
 
@@ -497,8 +510,8 @@ def workflow(
     description: str,
     initial_input_form: InputStepFunc | None = None,
     target: Target = Target.SYSTEM,
-    authorize_callback: Authorizer| None = None,
-    retry_auth_callback: Authorizer| None = None,
+    authorize_callback: Authorizer | None = None,
+    retry_auth_callback: Authorizer | None = None,
 ) -> Callable[[Callable[[], StepList]], Workflow]:
     """Transform an initial_input_form and a step list into a workflow.
 
@@ -519,7 +532,13 @@ def create_service_port():
 
     def _workflow(f: Callable[[], StepList]) -> Workflow:
         return make_workflow(
-            f, description, initial_input_form_in_form_inject_args, target, f(), authorize_callback=authorize_callback, retry_auth_callback=retry_auth_callback
+            f,
+            description,
+            initial_input_form_in_form_inject_args,
+            target,
+            f(),
+            authorize_callback=authorize_callback,
+            retry_auth_callback=retry_auth_callback,
         )
 
     return _workflow

orchestrator/workflows/utils.py

@@ -20,7 +20,6 @@
 from pydantic import field_validator, model_validator
 from sqlalchemy import select
 
-from oauth2_lib.fastapi import OIDCUserModel
 from orchestrator.db import ProductTable, SubscriptionTable, db
 from orchestrator.forms.validators import ProductId
 from orchestrator.services import subscriptions
@@ -242,7 +241,7 @@ def _create_workflow(f: Callable[[], StepList]) -> Workflow:
             Target.CREATE,
             steplist,
             authorize_callback=authorize_callback,
-            retry_auth_callback=retry_auth_callback
+            retry_auth_callback=retry_auth_callback,
         )
 
     return _create_workflow

test/unit_tests/api/test_processes.py

@@ -6,8 +6,6 @@
 from uuid import uuid4
 
 import pytest
-from pydantic_forms.core import FormPage
-from pydantic_forms.types import State
 from sqlalchemy import select
 
 from oauth2_lib.fastapi import OIDCUserModel
@@ -25,6 +23,7 @@
 from orchestrator.settings import app_settings
 from orchestrator.targets import Target
 from orchestrator.workflow import ProcessStatus, done, init, inputstep, step, workflow
+from pydantic_forms.core import FormPage
 from test.unit_tests.helpers import URL_STR_TYPE
 from test.unit_tests.workflows import WorkflowInstanceForTests
 
@@ -600,25 +599,25 @@ def unauthorized_workflow():
 def test_inputstep_authorization(test_client):
     def disallow(_: OIDCUserModel | None = None) -> bool:
         return False
-    
+
     def allow(_: OIDCUserModel | None = None) -> bool:
         return True
 
     class ConfirmForm(FormPage):
         confirm: bool
 
     @inputstep("unauthorized_resume", assignee=Assignee.SYSTEM, resume_auth_callback=disallow)
-    def unauthorized_resume(state: State) -> State:
+    def unauthorized_resume(state):
         user_input = yield ConfirmForm
         return user_input.model_dump()
 
     @inputstep("authorized_resume", assignee=Assignee.SYSTEM, resume_auth_callback=allow)
-    def authorized_resume(state: State) -> State:
+    def authorized_resume(state):
         user_input = yield ConfirmForm
         return user_input.model_dump()
 
     @inputstep("noauth_resume", assignee=Assignee.SYSTEM)
-    def noauth_resume(state: State) -> State:
+    def noauth_resume(state):
         user_input = yield ConfirmForm
         return user_input.model_dump()
 
@@ -640,5 +639,5 @@ def test_auth_workflow():
         response = test_client.put(f"/api/processes/{process_id}/resume", json=[{"confirm": True}])
         assert HTTPStatus.FORBIDDEN == response.status_code
 
-    #TODO test how this interacts with passing a different callback to @authorize_workflow
-    # These should be as functionally independent as possible.
+    # TODO test how this interacts with passing a different callback to @authorize_workflow
+    # These should be as functionally independent as possible.