From 2c4c39a0136be864aa558a7a2a350f39d507f156 Mon Sep 17 00:00:00 2001 From: vikramlc Date: Thu, 22 Oct 2026 16:02:29 +0530 Subject: [PATCH 1/7] feat(odin): Make Actions support opt-out per scheduled/continuous task --- cognite/extractorutils/unstable/core/base.py | 27 +++++++------- cognite/extractorutils/unstable/core/tasks.py | 37 +++++++++++++++++-- .../test_unstable/test_action_registration.py | 34 +++++++++++++++++ 3 files changed, 81 insertions(+), 17 deletions(-) diff --git a/cognite/extractorutils/unstable/core/base.py b/cognite/extractorutils/unstable/core/base.py index 61972b05..19550fb5 100644 --- a/cognite/extractorutils/unstable/core/base.py +++ b/cognite/extractorutils/unstable/core/base.py @@ -63,7 +63,7 @@ def my_task_function(self, task_context: TaskContext) -> None: from multiprocessing.synchronize import Event as MpEvent from threading import RLock, Thread from types import TracebackType -from typing import Generic, TypeVar +from typing import Generic, TypeGuard, TypeVar from humps import pascalize from typing_extensions import Self, assert_never @@ -124,9 +124,13 @@ def my_task_function(self, task_context: TaskContext) -> None: _T = TypeVar("_T", bound=ExtractorConfig) ACTIONABLE_TASK_TYPES: tuple[type[ScheduledTask], type[ContinuousTask]] = (ScheduledTask, ContinuousTask) -"""Task types that get auto-generated ``Start``/``Stop`` actions and are tracked in -``_running_task_tokens``. Kept as a single constant so every ``isinstance`` check that gates this -behavior stays in sync - add a type here (and only here) to make a new task type actionable.""" +"""Task types that can get auto-generated ``Start``/``Stop`` actions. Individual tasks can still opt out +with ``actionable=False``; use ``_is_actionable`` to check both.""" + + +def _is_actionable(task: Task) -> TypeGuard[ScheduledTask | ContinuousTask]: + """Whether ``task`` is of an actionable type and has not opted out of Actions support.""" + return isinstance(task, ACTIONABLE_TASK_TYPES) and task.actionable class FullConfig(Generic[_T]): @@ -392,10 +396,7 @@ def add_action(self, action: CustomAction) -> None: raise ValueError(f"Custom action '{action.name}' is already registered.") reserved_names = { - name - for t in self._tasks - if isinstance(t, ACTIONABLE_TASK_TYPES) - for name in (f"Start {t.name}", f"Stop {t.name}") + name for t in self._tasks if _is_actionable(t) for name in (f"Start {t.name}", f"Stop {t.name}") } if action.name in reserved_names: raise ValueError(f"Custom action name '{action.name}' conflicts with an auto-generated task action.") @@ -413,7 +414,7 @@ def _get_startup_request(self) -> StartupRequest: available_actions: list[AvailableActionWrite] = [] for t in self._tasks: - if isinstance(t, ACTIONABLE_TASK_TYPES): + if _is_actionable(t): available_actions.append( AvailableActionWrite(name=f"Start {t.name}", type=ActionType.start_task, task=t.name) ) @@ -433,7 +434,7 @@ def _get_startup_request(self) -> StartupRequest: tasks=[ DtoTask( type=TaskType.continuous if isinstance(t, ContinuousTask) else TaskType.batch, - action=isinstance(t, ACTIONABLE_TASK_TYPES), + action=_is_actionable(t), description=t.description, name=t.name, ) @@ -495,7 +496,7 @@ def add_task(self, task: Task) -> None: if any(t.name == task.name for t in self._tasks): raise ValueError(f"Task '{task.name}' is already registered.") - if isinstance(task, ACTIONABLE_TASK_TYPES): + if _is_actionable(task): conflict_names = {f"Start {task.name}", f"Stop {task.name}"} for action in self._custom_actions: if action.name in conflict_names: @@ -647,7 +648,7 @@ def _dispatch_single_action(self, action: Action) -> None: ) return - actionable_tasks = [t for t in self._tasks if isinstance(t, ACTIONABLE_TASK_TYPES)] + actionable_tasks = [t for t in self._tasks if _is_actionable(t)] scheduled_start_names = {f"Start {t.name}" for t in actionable_tasks} scheduled_stop_names = {f"Stop {t.name}" for t in actionable_tasks} custom_names = {a.name for a in self._custom_actions} @@ -670,7 +671,7 @@ def _dispatch_single_action(self, action: Action) -> None: def _handle_start_task_action(self, action: Action) -> None: task_name = action.action_name[len("Start ") :] task = next( - (t for t in self._tasks if t.name == task_name and isinstance(t, ACTIONABLE_TASK_TYPES)), + (t for t in self._tasks if t.name == task_name and _is_actionable(t)), None, ) if task is None: diff --git a/cognite/extractorutils/unstable/core/tasks.py b/cognite/extractorutils/unstable/core/tasks.py index 9ba249f1..36719194 100644 --- a/cognite/extractorutils/unstable/core/tasks.py +++ b/cognite/extractorutils/unstable/core/tasks.py @@ -68,10 +68,12 @@ def __init__( name: str, target: TaskTarget, description: str | None = None, + actionable: bool = False, ) -> None: self.name = name self.target = target self.description = description + self.actionable = actionable class ScheduledTask(_Task): @@ -85,6 +87,7 @@ class ScheduledTask(_Task): target: A callable that takes a ``TaskContext`` and performs the task. description: An optional description of the task. schedule: A ``ScheduleConfig`` object that defines the scheduling configuration for the task. + actionable: Whether the task supports Start/Stop actions. Enabled by default; set to ``False`` to opt out. """ def __init__( @@ -94,13 +97,20 @@ def __init__( target: TaskTarget, description: str | None = None, schedule: ScheduleConfig, + actionable: bool = True, ) -> None: - super().__init__(name=name, target=target, description=description) + super().__init__(name=name, target=target, description=description, actionable=actionable) self.schedule = schedule @classmethod def from_interval( - cls, *, interval: str, name: str, target: TaskTarget, description: str | None = None + cls, + *, + interval: str, + name: str, + target: TaskTarget, + description: str | None = None, + actionable: bool = True, ) -> "ScheduledTask": """ Create a scheduled task that runs at regular intervals. @@ -110,16 +120,26 @@ def from_interval( name: The name of the task. target: A callable that takes a ``TaskContext`` and performs the task. description: An optional description of the task. + actionable: Whether the task supports Start/Stop actions. Enabled by default. """ return ScheduledTask( name=name, target=target, description=description, + actionable=actionable, schedule=IntervalConfig(type="interval", expression=TimeIntervalConfig(interval)), ) @classmethod - def from_cron(cls, *, cron: str, name: str, target: TaskTarget, description: str | None = None) -> "ScheduledTask": + def from_cron( + cls, + *, + cron: str, + name: str, + target: TaskTarget, + description: str | None = None, + actionable: bool = True, + ) -> "ScheduledTask": """ Create a scheduled task that runs according to a cron expression. @@ -128,11 +148,13 @@ def from_cron(cls, *, cron: str, name: str, target: TaskTarget, description: str name: The name of the task. target: A callable that takes a ``TaskContext`` and performs the task. description: An optional description of the task. + actionable: Whether the task supports Start/Stop actions. Enabled by default. """ return ScheduledTask( name=name, target=target, description=description, + actionable=actionable, schedule=CronConfig(type="cron", expression=cron), ) @@ -142,6 +164,12 @@ class ContinuousTask(_Task): A task that runs continuously. Continuous tasks are started when the extractor starts and are expected to run until the extractor stops. + + Args: + name: The name of the task. + target: A callable that takes a ``TaskContext`` and performs the task. + description: An optional description of the task. + actionable: Whether the task supports Start/Stop actions. Enabled by default; set to ``False`` to opt out. """ def __init__( @@ -150,8 +178,9 @@ def __init__( name: str, target: TaskTarget, description: str | None = None, + actionable: bool = True, ) -> None: - super().__init__(name=name, target=target, description=description) + super().__init__(name=name, target=target, description=description, actionable=actionable) class StartupTask(_Task): diff --git a/tests/test_unstable/test_action_registration.py b/tests/test_unstable/test_action_registration.py index 153aaf8b..83c49164 100644 --- a/tests/test_unstable/test_action_registration.py +++ b/tests/test_unstable/test_action_registration.py @@ -70,6 +70,40 @@ def test_continuous_task_also_produces_start_stop_actions() -> None: assert by_name["Stop listener"].type == ActionType.stop_task +def test_tasks_are_actionable_by_default() -> None: + assert ScheduledTask.from_interval(interval="1h", name="a", target=lambda _: None).actionable + assert ScheduledTask.from_cron(cron="0 0 * * *", name="b", target=lambda _: None).actionable + assert ContinuousTask(name="c", target=lambda _: None).actionable + assert not StartupTask(name="d", target=lambda _: None).actionable + + +def test_opted_out_tasks_produce_no_actions_and_are_flagged_non_actionable() -> None: + extractor = _make_extractor() + extractor.add_task( + ScheduledTask.from_interval(interval="1h", name="sched", target=lambda _: None, actionable=False) + ) + extractor.add_task(ScheduledTask.from_cron(cron="0 0 * * *", name="cron", target=lambda _: None, actionable=False)) + extractor.add_task(ContinuousTask(name="cont", target=lambda _: None, actionable=False)) + extractor.add_task(ContinuousTask(name="kept", target=lambda _: None)) + req = _startup_request(extractor) + assert req.available_actions is not None + assert {a.name for a in req.available_actions if a.type != ActionType.custom} == {"Start kept", "Stop kept"} + assert req.tasks is not None + flags = {t.name: t.action for t in req.tasks} + assert {n: flags[n] for n in ("sched", "cron", "cont", "kept")} == { + "sched": False, + "cron": False, + "cont": False, + "kept": True, + } + + +def test_opted_out_task_name_does_not_reserve_action_names() -> None: + extractor = _make_extractor() + extractor.add_task(ContinuousTask(name="cont", target=lambda _: None, actionable=False)) + extractor.add_action(CustomAction(name="Start cont", target=lambda _: None)) + + @pytest.mark.parametrize( "action_name,expected_type,expected_task", [ From 8e9e9c7ccada74e5ddcc606f0c4474e917b5f361 Mon Sep 17 00:00:00 2001 From: vikramlc Date: Mon, 26 Oct 2026 13:30:36 +0530 Subject: [PATCH 2/7] test(odin): Add more asserts --- tests/test_unstable/test_action_registration.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_unstable/test_action_registration.py b/tests/test_unstable/test_action_registration.py index 83c49164..19f53dd0 100644 --- a/tests/test_unstable/test_action_registration.py +++ b/tests/test_unstable/test_action_registration.py @@ -102,6 +102,11 @@ def test_opted_out_task_name_does_not_reserve_action_names() -> None: extractor = _make_extractor() extractor.add_task(ContinuousTask(name="cont", target=lambda _: None, actionable=False)) extractor.add_action(CustomAction(name="Start cont", target=lambda _: None)) + req = _startup_request(extractor) + assert req.available_actions is not None + by_name = {a.name: a for a in req.available_actions} + assert by_name["Start cont"].type == ActionType.custom + assert "Stop cont" not in by_name @pytest.mark.parametrize( From 46782ffb409c9771f09970fe408f56bf4a886a5d Mon Sep 17 00:00:00 2001 From: vikramlc Date: Tue, 27 Oct 2026 17:27:16 +0530 Subject: [PATCH 3/7] refactor(odin): Removed unnecessary is_actionable method --- cognite/extractorutils/unstable/core/base.py | 29 ++++++-------------- 1 file changed, 8 insertions(+), 21 deletions(-) diff --git a/cognite/extractorutils/unstable/core/base.py b/cognite/extractorutils/unstable/core/base.py index 19550fb5..9a04bab3 100644 --- a/cognite/extractorutils/unstable/core/base.py +++ b/cognite/extractorutils/unstable/core/base.py @@ -63,7 +63,7 @@ def my_task_function(self, task_context: TaskContext) -> None: from multiprocessing.synchronize import Event as MpEvent from threading import RLock, Thread from types import TracebackType -from typing import Generic, TypeGuard, TypeVar +from typing import Generic, TypeVar from humps import pascalize from typing_extensions import Self, assert_never @@ -123,15 +123,6 @@ def my_task_function(self, task_context: TaskContext) -> None: _T = TypeVar("_T", bound=ExtractorConfig) -ACTIONABLE_TASK_TYPES: tuple[type[ScheduledTask], type[ContinuousTask]] = (ScheduledTask, ContinuousTask) -"""Task types that can get auto-generated ``Start``/``Stop`` actions. Individual tasks can still opt out -with ``actionable=False``; use ``_is_actionable`` to check both.""" - - -def _is_actionable(task: Task) -> TypeGuard[ScheduledTask | ContinuousTask]: - """Whether ``task`` is of an actionable type and has not opted out of Actions support.""" - return isinstance(task, ACTIONABLE_TASK_TYPES) and task.actionable - class FullConfig(Generic[_T]): """ @@ -395,9 +386,7 @@ def add_action(self, action: CustomAction) -> None: if any(a.name == action.name for a in self._custom_actions): raise ValueError(f"Custom action '{action.name}' is already registered.") - reserved_names = { - name for t in self._tasks if _is_actionable(t) for name in (f"Start {t.name}", f"Stop {t.name}") - } + reserved_names = {name for t in self._tasks if t.actionable for name in (f"Start {t.name}", f"Stop {t.name}")} if action.name in reserved_names: raise ValueError(f"Custom action name '{action.name}' conflicts with an auto-generated task action.") @@ -414,7 +403,7 @@ def _get_startup_request(self) -> StartupRequest: available_actions: list[AvailableActionWrite] = [] for t in self._tasks: - if _is_actionable(t): + if t.actionable: available_actions.append( AvailableActionWrite(name=f"Start {t.name}", type=ActionType.start_task, task=t.name) ) @@ -434,7 +423,7 @@ def _get_startup_request(self) -> StartupRequest: tasks=[ DtoTask( type=TaskType.continuous if isinstance(t, ContinuousTask) else TaskType.batch, - action=_is_actionable(t), + action=t.actionable, description=t.description, name=t.name, ) @@ -496,7 +485,7 @@ def add_task(self, task: Task) -> None: if any(t.name == task.name for t in self._tasks): raise ValueError(f"Task '{task.name}' is already registered.") - if _is_actionable(task): + if task.actionable: conflict_names = {f"Start {task.name}", f"Stop {task.name}"} for action in self._custom_actions: if action.name in conflict_names: @@ -560,9 +549,7 @@ def _release_if_owned( if tokens.get(key) is token: tokens.pop(key, None) - def _run_task_with_token( - self, task: ScheduledTask | ContinuousTask, child_token: CancellationToken | None = None - ) -> None: + def _run_task_with_token(self, task: Task, child_token: CancellationToken | None = None) -> None: """ Create a fresh child cancellation token for an actionable task run. @@ -648,7 +635,7 @@ def _dispatch_single_action(self, action: Action) -> None: ) return - actionable_tasks = [t for t in self._tasks if _is_actionable(t)] + actionable_tasks = [t for t in self._tasks if t.actionable] scheduled_start_names = {f"Start {t.name}" for t in actionable_tasks} scheduled_stop_names = {f"Stop {t.name}" for t in actionable_tasks} custom_names = {a.name for a in self._custom_actions} @@ -671,7 +658,7 @@ def _dispatch_single_action(self, action: Action) -> None: def _handle_start_task_action(self, action: Action) -> None: task_name = action.action_name[len("Start ") :] task = next( - (t for t in self._tasks if t.name == task_name and _is_actionable(t)), + (t for t in self._tasks if t.name == task_name and t.actionable), None, ) if task is None: From 7ce1b0fbded796d79ca09a49371f345e8fb9f4cb Mon Sep 17 00:00:00 2001 From: vikramlc Date: Tue, 27 Oct 2026 23:04:58 +0530 Subject: [PATCH 4/7] fix(odin): Revert changes to only allow scheduled task and continuous task with actions --- cognite/extractorutils/unstable/core/base.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/cognite/extractorutils/unstable/core/base.py b/cognite/extractorutils/unstable/core/base.py index 9a04bab3..958f0667 100644 --- a/cognite/extractorutils/unstable/core/base.py +++ b/cognite/extractorutils/unstable/core/base.py @@ -123,6 +123,9 @@ def my_task_function(self, task_context: TaskContext) -> None: _T = TypeVar("_T", bound=ExtractorConfig) +ACTIONABLE_TASK_TYPES: tuple[type[ScheduledTask], type[ContinuousTask]] = (ScheduledTask, ContinuousTask) +"""Task types that can support Start/Stop actions. Whether a given task supports actions is set by its ``actionable`` flag.""" + class FullConfig(Generic[_T]): """ @@ -549,14 +552,17 @@ def _release_if_owned( if tokens.get(key) is token: tokens.pop(key, None) - def _run_task_with_token(self, task: Task, child_token: CancellationToken | None = None) -> None: + def _run_task_with_token( + self, task: ScheduledTask | ContinuousTask, child_token: CancellationToken | None = None + ) -> None: """ - Create a fresh child cancellation token for an actionable task run. + Run a scheduled or continuous task under a fresh child cancellation token. Stores the token in ``_running_task_tokens`` for external cancellation (e.g. a stop_task action), then executes the task. Called by the scheduler (for ``ScheduledTask``), the initial boot-time launch (for ``ContinuousTask``), and the - start_task action handler (for both). + start_task action handler (for both). Tasks that opted out of Actions + (``actionable=False``) also run through here; they just never receive Start/Stop actions. If ``child_token`` is supplied (the caller has already registered it in ``_running_task_tokens`` atomically), this method skips the registration step and @@ -657,8 +663,9 @@ def _dispatch_single_action(self, action: Action) -> None: def _handle_start_task_action(self, action: Action) -> None: task_name = action.action_name[len("Start ") :] + # StartupTask is never actionable; the isinstance check only narrows the type for _run_task_with_token. task = next( - (t for t in self._tasks if t.name == task_name and t.actionable), + (t for t in self._tasks if t.name == task_name and t.actionable and isinstance(t, ACTIONABLE_TASK_TYPES)), None, ) if task is None: From 55a393c6db047df6c867e2e862fd5a2bfab44efb Mon Sep 17 00:00:00 2001 From: vikramlc Date: Tue, 27 Oct 2026 23:08:13 +0530 Subject: [PATCH 5/7] fix(odin): Fix lint issue --- cognite/extractorutils/unstable/core/base.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cognite/extractorutils/unstable/core/base.py b/cognite/extractorutils/unstable/core/base.py index 958f0667..33466ed1 100644 --- a/cognite/extractorutils/unstable/core/base.py +++ b/cognite/extractorutils/unstable/core/base.py @@ -124,7 +124,8 @@ def my_task_function(self, task_context: TaskContext) -> None: _T = TypeVar("_T", bound=ExtractorConfig) ACTIONABLE_TASK_TYPES: tuple[type[ScheduledTask], type[ContinuousTask]] = (ScheduledTask, ContinuousTask) -"""Task types that can support Start/Stop actions. Whether a given task supports actions is set by its ``actionable`` flag.""" +"""Task types that can support Start/Stop actions. +Whether a given task supports actions is set by its ``actionable`` flag.""" class FullConfig(Generic[_T]): From b172f7367f1627e885d2f49faf01e1f82c51b288 Mon Sep 17 00:00:00 2001 From: vikramlc Date: Wed, 28 Oct 2026 12:06:19 +0530 Subject: [PATCH 6/7] test(odin): Add test for start/stop on opted out tasks --- tests/test_unstable/test_action_dispatch.py | 26 +++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/test_unstable/test_action_dispatch.py b/tests/test_unstable/test_action_dispatch.py index 94216274..0d2a3757 100644 --- a/tests/test_unstable/test_action_dispatch.py +++ b/tests/test_unstable/test_action_dispatch.py @@ -44,6 +44,32 @@ def test_dispatch_unrecognised_action_name_reports_failed() -> None: assert "DoesNotExist" in (updates[0].result_message or "") +@pytest.mark.parametrize( + "make_task", + [ + pytest.param( + lambda target: ScheduledTask.from_interval(interval="1h", name="quiet", target=target, actionable=False), + id="scheduled", + ), + pytest.param(lambda target: ContinuousTask(name="quiet", target=target, actionable=False), id="continuous"), + ], +) +@pytest.mark.parametrize("action_name", ["Start quiet", "Stop quiet"]) +def test_dispatch_action_on_opted_out_task_reports_failed(make_task: Callable, action_name: str) -> None: + extractor = _make_extractor() + ran = [] + extractor.add_task(make_task(lambda ctx: ran.append(True))) + + extractor._dispatch_single_action(_make_action("act-1", action_name)) + + updates = _queued_updates(extractor) + assert len(updates) == 1 + assert updates[0].status == ActionStatus.failed + assert action_name in (updates[0].result_message or "") + assert not ran + assert "quiet" not in extractor._running_task_tokens + + @pytest.mark.parametrize( "register,action_name,expected_status", [ From 179cf8b62f1cbcc5b67dffdb119f2ca040b458e7 Mon Sep 17 00:00:00 2001 From: vikramlc Date: Wed, 28 Oct 2026 13:59:59 +0530 Subject: [PATCH 7/7] refactor(odin): Add an exact dispatcher rejection message --- tests/test_unstable/test_action_dispatch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_unstable/test_action_dispatch.py b/tests/test_unstable/test_action_dispatch.py index 0d2a3757..e529528d 100644 --- a/tests/test_unstable/test_action_dispatch.py +++ b/tests/test_unstable/test_action_dispatch.py @@ -65,7 +65,7 @@ def test_dispatch_action_on_opted_out_task_reports_failed(make_task: Callable, a updates = _queued_updates(extractor) assert len(updates) == 1 assert updates[0].status == ActionStatus.failed - assert action_name in (updates[0].result_message or "") + assert updates[0].result_message == f"No action named '{action_name}' registered" assert not ran assert "quiet" not in extractor._running_task_tokens