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

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

@@ -3,35 +3,52 @@ import logging
 import os
 import sys
 
-from .constants import ENV_TASK_BROKER_URI, ENV_GRANT_TOKEN, DEFAULT_TASK_BROKER_URI
-from .task_runner import TaskRunner
+os.environ["WEBSOCKETS_MAX_LOG_SIZE"] = "256"
 
-logging.basicConfig(
-    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+from .constants import (
+    DEFAULT_MAX_CONCURRENCY,
+    DEFAULT_TASK_TIMEOUT,
+    ENV_MAX_CONCURRENCY,
+    ENV_MAX_PAYLOAD_SIZE,
+    ENV_TASK_BROKER_URI,
+    ENV_GRANT_TOKEN,
+    DEFAULT_TASK_BROKER_URI,
+    DEFAULT_MAX_PAYLOAD_SIZE,
+    ENV_TASK_TIMEOUT,
 )
-
-logger = logging.getLogger(__name__)
+from .logs import setup_logging
+from .task_runner import TaskRunner, TaskRunnerOpts
 
 
 async def main():
-    task_broker_uri = os.getenv(ENV_TASK_BROKER_URI, DEFAULT_TASK_BROKER_URI)
+    setup_logging()
+    logger = logging.getLogger(__name__)
+
+    logger.info("Starting runner...")
+
     grant_token = os.getenv(ENV_GRANT_TOKEN, "")
 
     if not grant_token:
         logger.error(f"{ENV_GRANT_TOKEN} environment variable is required")
         sys.exit(1)
 
-    runner = TaskRunner(
-        task_broker_uri=task_broker_uri,
-        grant_token=grant_token,
+    opts = TaskRunnerOpts(
+        grant_token,
+        os.getenv(ENV_TASK_BROKER_URI, DEFAULT_TASK_BROKER_URI),
+        int(os.getenv(ENV_MAX_CONCURRENCY, DEFAULT_MAX_CONCURRENCY)),
+        int(os.getenv(ENV_MAX_PAYLOAD_SIZE, DEFAULT_MAX_PAYLOAD_SIZE)),
+        int(os.getenv(ENV_TASK_TIMEOUT, DEFAULT_TASK_TIMEOUT)),
     )
 
+    task_runner = TaskRunner(opts)
+
     try:
-        await runner.start()
-    except KeyboardInterrupt:
-        logger.info("Shutting down...")
+        await task_runner.start()
+    except (KeyboardInterrupt, asyncio.CancelledError):
+        logger.info("Shutting down runner...")
     finally:
-        await runner.stop()
+        await task_runner.stop()
+        logger.info("Runner stopped")
 
 
 if __name__ == "__main__":