feat(odin): Execute incoming actions from Odin

cognite/extractorutils/unstable/core/base.py

@@ -85,6 +85,11 @@ def my_task_function(self, task_context: TaskContext) -> None:
     create_state_store,
 )
 from cognite.extractorutils.unstable.core._dto import (
+    Action,
+    ActionStatus,
+    ActionType,
+    ActionUpdate,
+    AvailableActionWrite,
     CogniteModel,
     ExtractorInfo,
     StartupRequest,
@@ -94,6 +99,7 @@ def my_task_function(self, task_context: TaskContext) -> None:
     Task as DtoTask,
 )
 from cognite.extractorutils.unstable.core._messaging import RuntimeMessage
+from cognite.extractorutils.unstable.core.actions import ActionContext, CustomAction
 from cognite.extractorutils.unstable.core.checkin_worker import CheckinWorker
 from cognite.extractorutils.unstable.core.errors import Error, ErrorLevel
 from cognite.extractorutils.unstable.core.logger import CogniteLogger, RobustFileHandler
@@ -106,6 +112,7 @@ def my_task_function(self, task_context: TaskContext) -> None:
     "CogniteModel",
     "ConfigRevision",
     "ConfigType",
+    "CustomAction",
     "Extractor",
 ]
 
@@ -184,6 +191,9 @@ def __init__(self, config: FullConfig[ConfigType], checkin_worker: CheckinWorker
         self._scheduler = TaskScheduler(self.cancellation_token.create_child_token())
 
         self._tasks: list[Task] = []
+        self._custom_actions: list[CustomAction] = []
+        self._running_task_tokens: dict[str, CancellationToken] = {}
+        self._running_task_tokens_lock = RLock()
         self._start_time: datetime
 
         self.metrics: BaseMetrics = self._load_metrics(config.metrics_class)
@@ -195,6 +205,7 @@ def __init__(self, config: FullConfig[ConfigType], checkin_worker: CheckinWorker
         )
 
         self.__init_tasks__()
+        self.__init_actions__()
 
     def _setup_cancellation_watcher(self, cancel_event: MpEvent) -> None:
         """Starts a daemon thread to watch the inter-process event."""
@@ -335,6 +346,36 @@ def __init_tasks__(self) -> None:
         """
         pass
 
+    def __init_actions__(self) -> None:
+        """
+        This method should be overridden by subclasses to register custom actions.
+
+        It is called automatically after ``__init_tasks__`` during initialization.
+
+        Subclasses should call ``self.add_action(...)`` to register custom actions.
+        """
+        pass
+
+    def add_action(self, action: CustomAction) -> None:
+        """
+        Register a custom action that Odin can invoke on this extractor.
+
+        Args:
+            action: The custom action to register.
+        """
+        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 isinstance(t, ScheduledTask) 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 scheduled task action."
+            )
+
+        self._custom_actions.append(action)
+
     def _set_runtime_message_queue(self, queue: Queue) -> None:
         self._runtime_messages = queue
 
@@ -343,6 +384,22 @@ def _attach_runtime_controls(self, *, cancel_event: MpEvent, message_queue: Queu
         self._setup_cancellation_watcher(cancel_event)
 
     def _get_startup_request(self) -> StartupRequest:
+        available_actions: list[AvailableActionWrite] = []
+
+        for t in self._tasks:
+            if isinstance(t, ScheduledTask):
+                available_actions.append(
+                    AvailableActionWrite(name=f"Start {t.name}", type=ActionType.start_task, task=t.name)
+                )
+                available_actions.append(
+                    AvailableActionWrite(name=f"Stop {t.name}", type=ActionType.stop_task, task=t.name)
+                )
+
+        available_actions.extend(
+            AvailableActionWrite(name=a.name, type=ActionType.custom, description=a.description)
+            for a in self._custom_actions
+        )
+
         return StartupRequest(
             external_id=self.connection_config.integration.external_id,
             active_config_revision=self.current_config_revision,
@@ -358,6 +415,7 @@ def _get_startup_request(self) -> StartupRequest:
             ]
             if len(self._tasks) > 0
             else None,
+            available_actions=available_actions or None,
             timestamp=int(self._start_time.timestamp() * 1000),
         )
 
@@ -408,6 +466,9 @@ def add_task(self, task: Task) -> None:
         Args:
             task: The task to add. It should be an instance of ``StartupTask``, ``ContinuousTask``, or ``ScheduledTask``
         """
+        if any(t.name == task.name for t in self._tasks):
+            raise ValueError(f"Task '{task.name}' is already registered.")
+
         # Store this for later, since we'll override it with the wrapped version
         target = task.target
 
@@ -445,14 +506,156 @@ def run_task(task_context: TaskContext) -> None:
                 self._scheduler.schedule_task(
                     name=t.name,
                     schedule=t.schedule,
-                    task=lambda: t.target(
-                        TaskContext(
-                            task=task,
-                            extractor=self,
-                            cancellation_token=self.cancellation_token.create_child_token(),
-                        )
-                    ),
+                    task=lambda: self._run_task_with_token(t),
+                )
+
+    def _run_task_with_token(self, task: ScheduledTask) -> None:
+        """
+        Create a fresh child cancellation token for a scheduled task run.
+
+        Stores the token in ``_running_task_tokens`` for external cancellation (e.g. a
+        stop_task action), then executes the task. Called by both the scheduler and any
+        future start_task action handler.
+        """
+        child_token = self.cancellation_token.create_child_token()
+        with self._running_task_tokens_lock:
+            self._running_task_tokens[task.name] = child_token
+        try:
+            task.target(TaskContext(task=task, extractor=self, cancellation_token=child_token))
+        finally:
+            with self._running_task_tokens_lock:
+                if self._running_task_tokens.get(task.name) is child_token:
+                    self._running_task_tokens.pop(task.name, None)
+
+    def _handle_actions(self, actions: list[Action]) -> None:
+        for action in actions:
+            Thread(
+                target=self._dispatch_single_action,
+                args=(action,),
+                name=f"Action-{action.external_id}",
+                daemon=True,
+            ).start()
+
+    def _dispatch_single_action(self, action: Action) -> None:
+        scheduled_start_names = {f"Start {t.name}" for t in self._tasks if isinstance(t, ScheduledTask)}
+        scheduled_stop_names = {f"Stop {t.name}" for t in self._tasks if isinstance(t, ScheduledTask)}
+        custom_names = {a.name for a in self._custom_actions}
+
+        if action.action_name in scheduled_start_names:
+            self._handle_start_task_action(action)
+        elif action.action_name in scheduled_stop_names:
+            self._handle_stop_task_action(action)
+        elif action.action_name in custom_names:
+            self._handle_custom_action(action)
+        else:
+            self._checkin_worker.queue_action_update(
+                ActionUpdate(
+                    external_id=action.external_id,
+                    status=ActionStatus.failed,
+                    result_message=f"No action named '{action.action_name}' registered",
                 )
+            )
+
+    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, ScheduledTask)),
+            None,
+        )
+        if task is None:
+            self._checkin_worker.queue_action_update(
+                ActionUpdate(
+                    external_id=action.external_id,
+                    status=ActionStatus.failed,
+                    result_message=f"No scheduled task named '{task_name}' found",
+                )
+            )
+            return
+
+        with self._running_task_tokens_lock:
+            if task_name in self._running_task_tokens:
+                self._checkin_worker.queue_action_update(
+                    ActionUpdate(
+                        external_id=action.external_id,
+                        status=ActionStatus.failed,
+                        result_message=f"Task '{task_name}' is already running",
+                    )
+                )
+                return
+
+        self._checkin_worker.queue_action_update(
+            ActionUpdate(external_id=action.external_id, status=ActionStatus.running)
+        )
+
+        try:
+            self._run_task_with_token(task)
+            self._checkin_worker.queue_action_update(
+                ActionUpdate(external_id=action.external_id, status=ActionStatus.succeeded)
+            )
+        except Exception as e:
+            self._checkin_worker.queue_action_update(
+                ActionUpdate(
+                    external_id=action.external_id,
+                    status=ActionStatus.failed,
+                    result_message=str(e),
+                )
+            )
+
+    def _handle_stop_task_action(self, action: Action) -> None:
+        task_name = action.action_name[len("Stop ") :]
+        with self._running_task_tokens_lock:
+            token = self._running_task_tokens.get(task_name)
+        if token is None:
+            self._checkin_worker.queue_action_update(
+                ActionUpdate(
+                    external_id=action.external_id,
+                    status=ActionStatus.failed,
+                    result_message=f"Task '{task_name}' is not currently running",
+                )
+            )
+            return
+
+        token.cancel()
+        self._checkin_worker.queue_action_update(
+            ActionUpdate(external_id=action.external_id, status=ActionStatus.canceled)
+        )
+
+    def _handle_custom_action(self, action: Action) -> None:
+        custom = next((a for a in self._custom_actions if a.name == action.action_name), None)
+        if custom is None:
+            self._checkin_worker.queue_action_update(
+                ActionUpdate(
+                    external_id=action.external_id,
+                    status=ActionStatus.failed,
+                    result_message=f"No custom action named '{action.action_name}' registered",
+                )
+            )
+            return
+
+        self._checkin_worker.queue_action_update(
+            ActionUpdate(external_id=action.external_id, status=ActionStatus.running)
+        )
+
+        ctx = ActionContext(
+            action=custom,
+            extractor=self,
+            external_id=action.external_id,
+            call_metadata=action.call_metadata,
+        )
+
+        try:
+            custom.target(ctx)
+            self._checkin_worker.queue_action_update(
+                ActionUpdate(external_id=action.external_id, status=ActionStatus.succeeded)
+            )
+        except Exception as e:
+            self._checkin_worker.queue_action_update(
+                ActionUpdate(
+                    external_id=action.external_id,
+                    status=ActionStatus.failed,
+                    result_message=str(e),
+                )
+            )
 
     def start(self) -> None:
         """
@@ -467,6 +670,8 @@ def start(self) -> None:
         self._load_state_store()
         self.state_store.start()
 
+        self._checkin_worker.set_action_dispatcher(self._handle_actions)
+
         Thread(target=self._run_checkin, name="ExtractorCheckin", daemon=True).start()
         if self.metrics_push_manager:
             self.metrics_push_manager.start()