feat(core): Add execution flow to native Python runner (no-changelog) (#18485)

packages/@n8n/task-runner-python/src/task_executor.py

@@ -0,0 +1,133 @@
+from queue import Empty
+import multiprocessing
+import traceback
+import textwrap
+
+from .errors import (
+    TaskResultMissingError,
+    TaskRuntimeError,
+    TaskTimeoutError,
+    TaskProcessExitError,
+)
+
+from .message_types.broker import NodeMode, Items
+from .constants import EXECUTOR_USER_OUTPUT_KEY
+
+from multiprocessing.context import SpawnProcess
+
+MULTIPROCESSING_CONTEXT = multiprocessing.get_context("spawn")
+
+
+class TaskExecutor:
+    """Responsible for executing Python code tasks in isolated subprocesses."""
+
+    @staticmethod
+    def create_process(code: str, node_mode: NodeMode, items: Items):
+        """Create a subprocess for executing a Python code task and a queue for communication."""
+
+        fn = (
+            TaskExecutor._all_items
+            if node_mode == "all_items"
+            else TaskExecutor._per_item
+        )
+
+        queue = MULTIPROCESSING_CONTEXT.Queue()
+        process = MULTIPROCESSING_CONTEXT.Process(target=fn, args=(code, items, queue))
+
+        return process, queue
+
+    @staticmethod
+    def execute_process(
+        process: SpawnProcess,
+        queue: multiprocessing.Queue,
+        task_timeout: int,
+        continue_on_fail: bool,
+    ):
+        """Execute a subprocess for a Python code task."""
+
+        try:
+            process.start()
+            process.join(timeout=task_timeout)
+
+            if process.is_alive():
+                TaskExecutor.stop_process(process)
+                raise TaskTimeoutError(task_timeout)
+
+            if process.exitcode != 0:
+                assert process.exitcode is not None
+                raise TaskProcessExitError(process.exitcode)
+
+            try:
+                returned = queue.get_nowait()
+            except Empty:
+                raise TaskResultMissingError()
+
+            if "error" in returned:
+                raise TaskRuntimeError(returned["error"])
+
+            return returned["result"] or []
+
+        except Exception as e:
+            if continue_on_fail:
+                return [{"json": {"error": str(e)}}]
+            raise
+
+    @staticmethod
+    def stop_process(process: SpawnProcess | None):
+        """Stop a running subprocess, gracefully else force-killing."""
+
+        if process is None or not process.is_alive():
+            return
+
+        process.terminate()
+        process.join(timeout=1)  # 1s grace period
+
+        if process.is_alive():
+            process.kill()
+
+    @staticmethod
+    def _all_items(raw_code: str, items: Items, queue: multiprocessing.Queue):
+        """Execute a Python code task in all-items mode."""
+
+        try:
+            code = TaskExecutor._wrap_code(raw_code)
+            globals = {"__builtins__": __builtins__, "_items": items}
+            exec(code, globals)
+            queue.put({"result": globals[EXECUTOR_USER_OUTPUT_KEY]})
+
+        except Exception as e:
+            TaskExecutor._put_error(queue, e)
+
+    @staticmethod
+    def _per_item(raw_code: str, items: Items, queue: multiprocessing.Queue):
+        """Execute a Python code task in per-item mode."""
+
+        try:
+            wrapped_code = TaskExecutor._wrap_code(raw_code)
+            compiled_code = compile(wrapped_code, "<per_item_task_execution>", "exec")
+
+            result = []
+            for index, item in enumerate(items):
+                globals = {"__builtins__": __builtins__, "_item": item}
+                exec(compiled_code, globals)
+                user_output = globals[EXECUTOR_USER_OUTPUT_KEY]
+
+                if user_output is None:
+                    continue
+
+                user_output["pairedItem"] = {"item": index}
+                result.append(user_output)
+
+            queue.put({"result": result})
+
+        except Exception as e:
+            TaskExecutor._put_error(queue, e)
+
+    @staticmethod
+    def _wrap_code(raw_code: str) -> str:
+        indented_code = textwrap.indent(raw_code, "    ")
+        return f"def _user_function():\n{indented_code}\n\n{EXECUTOR_USER_OUTPUT_KEY} = _user_function()"
+
+    @staticmethod
+    def _put_error(queue: multiprocessing.Queue, e: Exception):
+        queue.put({"error": {"message": str(e), "stack": traceback.format_exc()}})