diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 07b5967c91a09..7f6ad1479d4e2 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -264,6 +264,18 @@ repos: additional_dependencies: ['pnpm@10.28.1'] pass_filenames: false require_serial: true + - id: check-ts-sdk-supervisor-schema + name: Check TypeScript SDK supervisor schema is up to date + entry: ./ts-sdk/scripts/ci/prek/check_supervisor_schema.py + language: node + files: > + (?x) + ^task-sdk/src/airflow/sdk/execution_time/schema/schema\.json$| + ^ts-sdk/src/generated/supervisor\.ts$| + ^ts-sdk/scripts/generate-supervisor\.mjs$ + additional_dependencies: ['pnpm@10.28.1'] + pass_filenames: false + require_serial: true - id: check-ci-workflows-in-sync name: Check ci-arm.yml and ci-amd.yml stay in sync entry: ./scripts/ci/prek/check_ci_workflows_in_sync.py diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index ae736e5d6621e..e6ce8aa3d066e 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4515,6 +4515,18 @@ "format": "date-time", "title": "Timestamp", "type": "string" + }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key" } }, "required": [ diff --git a/ts-sdk/README.md b/ts-sdk/README.md index f76ad3ba9597a..3ac7360aca609 100644 --- a/ts-sdk/README.md +++ b/ts-sdk/README.md @@ -23,10 +23,8 @@ Public TypeScript interfaces for writing Apache Airflow task handlers. **Status:** alpha · API will change · Node 22+ · ESM-only -This package defines the user-facing task handler contract: task registration, -runtime context types, and the `TaskClient` interface used for Airflow -Variables, Connections, and XCom. Runtime transports implement this interface -separately. +This package defines the user-facing task handler contract and the coordinator +runtime used to execute registered TypeScript handlers from Airflow. ## Install @@ -50,7 +48,7 @@ registerTask({ dagId: "example_dag", taskId: "say_hello" }, sayHello); Non-`undefined` return values are pushed to XCom under the `"return_value"` key by the active runtime, matching Python `@task` behavior. -## Intended Coordinator Usage +## Coordinator Usage Airflow runs TypeScript task bundles through the Python-side `airflow.sdk.coordinators.node.NodeCoordinator`. Declaring Airflow Dags in @@ -95,10 +93,10 @@ queue_to_coordinator = {"typescript": "ts"} Each configured bundle directory must contain `bundle.mjs` and `airflow-metadata.yaml`. -TypeScript handlers: +TypeScript entrypoint: ```ts -import { registerTask, type TaskHandlerArgs } from "@apache-airflow/ts-sdk"; +import { registerTask, startCoordinator, type TaskHandlerArgs } from "@apache-airflow/ts-sdk"; export async function extract({ client }: TaskHandlerArgs) { const connection = await client.getConnection("sales_db"); @@ -123,15 +121,34 @@ export async function transform({ client }: TaskHandlerArgs) { registerTask({ dagId: "sales_pipeline", taskId: "extract" }, extract); registerTask({ dagId: "sales_pipeline", taskId: "transform" }, transform); + +await startCoordinator(); ``` The Python stub defines the Dag dependency graph. The TypeScript handler does the work and uses `TaskClient` for task-time Airflow data access. Register each handler with the Python Dag's `dag_id` and the stub task's `task_id`. The handler function is the reusable task implementation; `registerTask` binds that -handler to a Python stub Dag/task identity for coordinator mode. A future -TypeScript Dag authoring API can attach the same handlers without changing the -handler code. +handler to a Python stub Dag/task identity for coordinator mode. + +For larger projects, keep one Airflow entrypoint that imports every module that +registers tasks, then starts the coordinator: + +```ts +import "./sales/tasks"; +import "./billing/tasks"; +import { startCoordinator } from "@apache-airflow/ts-sdk"; + +await startCoordinator(); +``` + +Airflow launches the bundled entrypoint with `--comm=host:port` and +`--logs=host:port`. `startCoordinator()` connects to those sockets, receives +the task startup message, finds the registered handler for the Dag/task pair, +and reports the terminal task state back to Airflow. + +See [`example/`](example/) for a coordinator-runtime example that builds a +`bundle.mjs` with `esbuild` and uses a Python stub Dag. ## TaskClient diff --git a/ts-sdk/eslint.config.js b/ts-sdk/eslint.config.js index 9fed412ab7e19..4d1f851d186b5 100644 --- a/ts-sdk/eslint.config.js +++ b/ts-sdk/eslint.config.js @@ -22,7 +22,7 @@ import tseslint from "typescript-eslint"; export default tseslint.config( { - ignores: ["dist/**", "node_modules/**", "coverage/**"], + ignores: ["dist/**", "node_modules/**", "coverage/**", "src/generated/**"], }, js.configs.recommended, ...tseslint.configs.recommended, diff --git a/ts-sdk/example/.gitignore b/ts-sdk/example/.gitignore new file mode 100644 index 0000000000000..8ee10058ff835 --- /dev/null +++ b/ts-sdk/example/.gitignore @@ -0,0 +1,3 @@ +.pnpm-store/ +dist/ +node_modules/ diff --git a/ts-sdk/example/README.md b/ts-sdk/example/README.md new file mode 100644 index 0000000000000..181135a52312c --- /dev/null +++ b/ts-sdk/example/README.md @@ -0,0 +1,101 @@ + + +# TypeScript Coordinator Runtime Example + +This example shows the coordinator-mode shape for TypeScript task handlers: + +- `dags/typescript_example.py` declares the Airflow Dag and stub tasks. +- `src/main.ts` registers TypeScript handlers for the same Dag/task IDs and + starts the coordinator runtime. +- `dist/bundle.mjs` is the generated Node.js bundle that Airflow launches. + +The TypeScript SDK does not include a packer yet, so this example builds the +bundle with `esbuild` and writes the Airflow metadata file manually. + +## Build + +Build the SDK first so the example can import the local package: + +```bash +cd ts-sdk +pnpm install +pnpm run build +``` + +Build the example bundle: + +```bash +cd ts-sdk/example +pnpm install +pnpm run build +``` + +Create the metadata file next to the generated bundle: + +```bash +node --input-type=module > dist/airflow-metadata.yaml <<'EOF' +import { SUPERVISOR_API_VERSION } from "@apache-airflow/ts-sdk"; + +console.log(`sdk: + supervisor_schema_version: "${SUPERVISOR_API_VERSION}"`); +EOF +``` + +The coordinator expects this layout: + +```text +ts-sdk/example/dist/ + bundle.mjs + airflow-metadata.yaml +``` + +## Airflow Configuration + +Configure Airflow to route the `typescript` queue to the Node coordinator and +point it at the example bundle directory: + +```bash +export AIRFLOW__SDK__COORDINATORS='{ + "node": { + "classpath": "airflow.sdk.coordinators.node.NodeCoordinator", + "kwargs": {"bundles_root": ["/absolute/path/to/airflow/ts-sdk/example/dist"]} + } +}' +export AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{"typescript": "node"}' +``` + +Copy `dags/typescript_example.py` into your Airflow Dags folder. + +The example also uses one Variable and one Connection: + +```bash +airflow variables set typescript_example_greeting "hello from Airflow" +airflow connections add typescript_example_http \ + --conn-type http \ + --conn-host example.com \ + --conn-login user \ + --conn-password pass +``` + +Then start Airflow and trigger the Dag: + +```bash +airflow dags trigger typescript_example +``` diff --git a/ts-sdk/example/dags/typescript_example.py b/ts-sdk/example/dags/typescript_example.py new file mode 100644 index 0000000000000..6f6ae0922d06c --- /dev/null +++ b/ts-sdk/example/dags/typescript_example.py @@ -0,0 +1,45 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from airflow.sdk import dag, task + + +@task +def python_start(): + return "hello from Python" + + +@task.stub(queue="typescript") +def build_message(): ... + + +@task.stub(queue="typescript") +def read_connection(): ... + + +@dag(dag_id="typescript_example", schedule=None, catchup=False, tags=["typescript", "example"]) +def typescript_example(): + start = python_start() + message = build_message() + read_connection() + + start >> message + + +typescript_example() diff --git a/ts-sdk/example/package.json b/ts-sdk/example/package.json new file mode 100644 index 0000000000000..5594727113c13 --- /dev/null +++ b/ts-sdk/example/package.json @@ -0,0 +1,19 @@ +{ + "name": "@apache-airflow/ts-sdk-example", + "private": true, + "version": "0.0.0", + "type": "module", + "license": "Apache-2.0", + "scripts": { + "build": "esbuild src/main.ts --bundle --platform=node --format=esm --target=node22 --outfile=dist/bundle.mjs", + "typecheck": "tsc --noEmit" + }, + "dependencies": { + "@apache-airflow/ts-sdk": "file:.." + }, + "devDependencies": { + "@types/node": "^22.19.17", + "esbuild": "^0.28.1", + "typescript": "^6.0.2" + } +} diff --git a/ts-sdk/example/src/main.ts b/ts-sdk/example/src/main.ts new file mode 100644 index 0000000000000..1c480b8056b85 --- /dev/null +++ b/ts-sdk/example/src/main.ts @@ -0,0 +1,55 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { registerTask, startCoordinator, type TaskHandlerArgs } from "@apache-airflow/ts-sdk"; + +const DAG_ID = "typescript_example"; + +export async function buildMessage({ client }: TaskHandlerArgs) { + const upstream = await client.getXCom({ + key: "return_value", + taskId: "python_start", + }); + const greeting = await client.getVariable("typescript_example_greeting"); + const message = `${greeting ?? "hello from TypeScript"}; upstream=${upstream ?? "missing"}`; + + await client.setXCom({ key: "typescript_message", value: message }); + + return { + message, + upstream, + }; +} + +export async function readConnection({ client }: TaskHandlerArgs) { + const connection = await client.getConnection("typescript_example_http"); + + return { + id: connection?.id ?? null, + type: connection?.type ?? null, + host: connection?.host ?? null, + login: connection?.login ?? null, + hasPassword: connection?.password != null, + }; +} + +registerTask({ dagId: DAG_ID, taskId: "build_message" }, buildMessage); +registerTask({ dagId: DAG_ID, taskId: "read_connection" }, readConnection); + +await startCoordinator(); diff --git a/ts-sdk/example/tsconfig.json b/ts-sdk/example/tsconfig.json new file mode 100644 index 0000000000000..01f4c7b15121b --- /dev/null +++ b/ts-sdk/example/tsconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../tsconfig.json", + "compilerOptions": { + "noEmit": true, + "rootDir": "." + }, + "include": ["src/**/*.ts"] +} diff --git a/ts-sdk/package.json b/ts-sdk/package.json index 2a36a73058deb..190df17101456 100644 --- a/ts-sdk/package.json +++ b/ts-sdk/package.json @@ -2,7 +2,7 @@ "name": "@apache-airflow/ts-sdk", "version": "0.1.0-alpha.0", "packageManager": "pnpm@10.28.1", - "description": "Public TypeScript Task SDK interfaces for Apache Airflow task handlers", + "description": "TypeScript Task SDK for Apache Airflow task handlers", "license": "Apache-2.0", "type": "module", "main": "./dist/index.js", @@ -20,6 +20,10 @@ ".": { "types": "./dist/index.d.ts", "import": "./dist/index.js" + }, + "./coordinator": { + "types": "./dist/coordinator/index.d.ts", + "import": "./dist/coordinator/index.js" } }, "files": [ @@ -38,7 +42,8 @@ "test": "vitest run", "test:watch": "vitest", "build": "pnpm run clean && tsc -p tsconfig.build.json", - "prepack": "pnpm run build" + "prepack": "pnpm run build", + "generate:supervisor": "node scripts/generate-supervisor.mjs" }, "keywords": [ "airflow", @@ -51,10 +56,14 @@ "engines": { "node": ">=22" }, + "dependencies": { + "@msgpack/msgpack": "^3.1.2" + }, "devDependencies": { "@eslint/js": "^10.0.1", "@types/node": "^22.19.17", "eslint": "^10.4.0", + "json-schema-to-typescript": "^15.0.4", "prettier": "^3.8.3", "tsx": "^4.21.0", "typescript": "^6.0.2", diff --git a/ts-sdk/pnpm-lock.yaml b/ts-sdk/pnpm-lock.yaml index 808dc51c4516d..887acdcc7e418 100644 --- a/ts-sdk/pnpm-lock.yaml +++ b/ts-sdk/pnpm-lock.yaml @@ -7,6 +7,10 @@ settings: importers: .: + dependencies: + '@msgpack/msgpack': + specifier: ^3.1.2 + version: 3.1.3 devDependencies: '@eslint/js': specifier: ^10.0.1 @@ -17,6 +21,9 @@ importers: eslint: specifier: ^10.4.0 version: 10.4.0 + json-schema-to-typescript: + specifier: ^15.0.4 + version: 15.0.4 prettier: specifier: ^3.8.3 version: 3.8.3 @@ -35,6 +42,10 @@ importers: packages: + '@apidevtools/json-schema-ref-parser@11.9.3': + resolution: {integrity: sha512-60vepv88RwcJtSHrD6MjIL6Ta3SOYbgfnkHb+ppAVK+o9mXprRtulx7VlRl3lN3bbvysAfCS7WMVfhUYemB0IQ==} + engines: {node: '>= 16'} + '@babel/helper-string-parser@7.29.7': resolution: {integrity: sha512-Pb5ijPrZ89GDH8223L4UP8i6QApWxs04RbPQJTeWDV0/keR2E36MeKnyr6LYmUUvqRRI+Iv87SuF1W6ErINzYw==} engines: {node: '>=6.9.0'} @@ -290,6 +301,13 @@ packages: '@jridgewell/trace-mapping@0.3.31': resolution: {integrity: sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw==} + '@jsdevtools/ono@7.1.3': + resolution: {integrity: sha512-4JQNk+3mVzK3xh2rqd6RB4J46qUR19azEHBneZyTZM+c456qOrbbM/5xcR8huNCCcbVt7+UmizG6GuUvPvKUYg==} + + '@msgpack/msgpack@3.1.3': + resolution: {integrity: sha512-47XIizs9XZXvuJgoaJUIE2lFoID8ugvc0jzSHP+Ptfk8nTbnR8g788wv48N03Kx0UkAv559HWRQ3yzOgzlRNUA==} + engines: {node: '>= 18'} + '@napi-rs/wasm-runtime@1.1.4': resolution: {integrity: sha512-3NQNNgA1YSlJb/kMH1ildASP9HW7/7kYnRI2szWJaofaS1hWmbGI4H+d3+22aGzXXN9IJ+n+GiFVcGipJP18ow==} peerDependencies: @@ -412,6 +430,9 @@ packages: '@types/json-schema@7.0.15': resolution: {integrity: sha512-5+fP8P8MFNC+AyZCDxrB2pkZFPGzqQWUzpSeuuVLvm8VMcorNYavBqoFcxK8bQz4Qsbn4oUEEem4wDLfcysGHA==} + '@types/lodash@4.17.24': + resolution: {integrity: sha512-gIW7lQLZbue7lRSWEFql49QJJWThrTFFeIMJdp3eH4tKoxm1OvEPg02rm4wCCSHS0cL3/Fizimb35b7k8atwsQ==} + '@types/node@22.19.17': resolution: {integrity: sha512-wGdMcf+vPYM6jikpS/qhg6WiqSV/OhG+jeeHT/KlVqxYfD40iYJf9/AE1uQxVWFvU7MipKRkRv8NSHiCGgPr8Q==} @@ -525,6 +546,9 @@ packages: ajv@6.15.0: resolution: {integrity: sha512-fgFx7Hfoq60ytK2c7DhnF8jIvzYgOMxfugjLOSMHjLIPgenqa7S7oaagATUq99mV6IYvN2tRmC0wnTYX6iPbMw==} + argparse@2.0.1: + resolution: {integrity: sha512-8+9WqebbFzpX9OR+Wa6O29asIogeRMzcGtAINdpMHHyAg10f05aSFVBbcEqGf/PXw1EjAZ+q2/bEBg3DvurK3Q==} + assertion-error@2.0.1: resolution: {integrity: sha512-Izi8RQcffqCeNVgFigKli1ssklIbpHnCYc6AknXGYoB6grJqyeby7jv12JUQgmTAnIDnbck1uxksT4dzN3PWBA==} engines: {node: '>=12'} @@ -718,9 +742,18 @@ packages: js-tokens@10.0.0: resolution: {integrity: sha512-lM/UBzQmfJRo9ABXbPWemivdCW8V2G8FHaHdypQaIy523snUjog0W71ayWXTjiR+ixeMyVHN2XcpnTd/liPg/Q==} + js-yaml@4.3.0: + resolution: {integrity: sha512-1td788aAnnZ5qs7V2QIRl1owjtYpbKt749Y3xauqQgwIIGF/xXWz1wMTEBx5O3LK3lXLVuqXPdPxj2BoFHaW9Q==} + hasBin: true + json-buffer@3.0.1: resolution: {integrity: sha512-4bV5BfR2mqfQTJm+V5tPPdf+ZpuhiIvTuAB5g8kcrXOZpTT/QwwVRWBywX1ozr6lEuPdbHxwaJlm9G6mI2sfSQ==} + json-schema-to-typescript@15.0.4: + resolution: {integrity: sha512-Su9oK8DR4xCmDsLlyvadkXzX6+GGXJpbhwoLtOGArAG61dvbW4YQmSEno2y66ahpIdmLMg6YUf/QHLgiwvkrHQ==} + engines: {node: '>=16.0.0'} + hasBin: true + json-schema-traverse@0.4.1: resolution: {integrity: sha512-xbbCH5dCYU5T8LcEhhuh7HJ88HXuW3qsI3Y0zOZFKfZEHcpWiHU/Jxzk629Brsab/mMiHQti9wMP+845RPe3Vg==} @@ -808,6 +841,9 @@ packages: resolution: {integrity: sha512-iPZK6eYjbxRu3uB4/WZ3EsEIMJFMqAoopl3R+zuq0UjcAm/MO6KCweDgPfP3elTztoKP3KtnVHxTn2NHBSDVUw==} engines: {node: '>=10'} + lodash@4.18.1: + resolution: {integrity: sha512-dMInicTPVE8d1e5otfwmmjlxkZoUpiVLwyeTdUsi/Caj/gfzzblBcCE5sRHV/AsjuCmxWrte2TNGSYuCeCq+0Q==} + magic-string@0.30.21: resolution: {integrity: sha512-vd2F4YUyEXKGcLHoq+TEyCjxueSeHnFxyyjNp80yg0XV4vUhnDer/lvvlqM/arB5bXQN5K2/3oinyCRyx8T2CQ==} @@ -822,6 +858,9 @@ packages: resolution: {integrity: sha512-MULkVLfKGYDFYejP07QOurDLLQpcjk7Fw+7jXS2R2czRQzR56yHRveU5NDJEOviH+hETZKSkIk5c+T23GjFUMg==} engines: {node: 18 || 20 || >=22} + minimist@1.2.8: + resolution: {integrity: sha512-2yyAR8qBkN3YuheJanUpWC5U3bb5osDywNB8RzDVlDwDHbocAJveqqj1u8+SVD7jkWT4yvsHCpWqqWqAxb0zCA==} + ms@2.1.3: resolution: {integrity: sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA==} @@ -1076,6 +1115,12 @@ packages: snapshots: + '@apidevtools/json-schema-ref-parser@11.9.3': + dependencies: + '@jsdevtools/ono': 7.1.3 + '@types/json-schema': 7.0.15 + js-yaml: 4.3.0 + '@babel/helper-string-parser@7.29.7': optional: true @@ -1251,6 +1296,10 @@ snapshots: '@jridgewell/sourcemap-codec': 1.5.5 optional: true + '@jsdevtools/ono@7.1.3': {} + + '@msgpack/msgpack@3.1.3': {} + '@napi-rs/wasm-runtime@1.1.4(@emnapi/core@1.10.0)(@emnapi/runtime@1.10.0)': dependencies: '@emnapi/core': 1.10.0 @@ -1331,6 +1380,8 @@ snapshots: '@types/json-schema@7.0.15': {} + '@types/lodash@4.17.24': {} + '@types/node@22.19.17': dependencies: undici-types: 6.21.0 @@ -1495,6 +1546,8 @@ snapshots: json-schema-traverse: 0.4.1 uri-js: 4.4.1 + argparse@2.0.1: {} + assertion-error@2.0.1: {} ast-v8-to-istanbul@1.0.2: @@ -1707,8 +1760,24 @@ snapshots: js-tokens@10.0.0: optional: true + js-yaml@4.3.0: + dependencies: + argparse: 2.0.1 + json-buffer@3.0.1: {} + json-schema-to-typescript@15.0.4: + dependencies: + '@apidevtools/json-schema-ref-parser': 11.9.3 + '@types/json-schema': 7.0.15 + '@types/lodash': 4.17.24 + is-glob: 4.0.3 + js-yaml: 4.3.0 + lodash: 4.18.1 + minimist: 1.2.8 + prettier: 3.8.3 + tinyglobby: 0.2.16 + json-schema-traverse@0.4.1: {} json-stable-stringify-without-jsonify@1.0.1: {} @@ -1775,6 +1844,8 @@ snapshots: dependencies: p-locate: 5.0.0 + lodash@4.18.1: {} + magic-string@0.30.21: dependencies: '@jridgewell/sourcemap-codec': 1.5.5 @@ -1795,6 +1866,8 @@ snapshots: dependencies: brace-expansion: 5.0.6 + minimist@1.2.8: {} + ms@2.1.3: {} nanoid@3.3.11: {} diff --git a/ts-sdk/scripts/ci/prek/check_supervisor_schema.py b/ts-sdk/scripts/ci/prek/check_supervisor_schema.py new file mode 100755 index 0000000000000..81e6fe1ca5e32 --- /dev/null +++ b/ts-sdk/scripts/ci/prek/check_supervisor_schema.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python3 +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[4] / "scripts" / "ci" / "prek")) + +from common_prek_utils import AIRFLOW_ROOT_PATH, console, run_command + +if __name__ not in ("__main__", "__mp_main__"): + raise SystemExit( + "This file is intended to be executed as an executable program. You cannot use it as a module." + f"To run this script, run the ./{__file__} command" + ) + +# Path of the generated file, relative to the repo root (for git diff / messages). +GENERATED = "ts-sdk/src/generated/supervisor.ts" + +if __name__ == "__main__": + directory = AIRFLOW_ROOT_PATH / "ts-sdk" + run_command(["pnpm", "config", "set", "store-dir", ".pnpm-store"], cwd=directory) + run_command(["pnpm", "install", "--frozen-lockfile", "--config.confirmModulesPurge=false"], cwd=directory) + # Regenerate, then format exactly as `pnpm run format` would, so the diff + # reflects a stale schema and not raw-vs-prettier formatting noise. + run_command(["pnpm", "run", "generate:supervisor"], cwd=directory) + run_command(["pnpm", "exec", "prettier", "--write", "src/generated/supervisor.ts"], cwd=directory) + + diff = subprocess.run( + ["git", "diff", "--", GENERATED], + cwd=AIRFLOW_ROOT_PATH, + capture_output=True, + text=True, + check=False, + ) + if diff.stdout.strip(): + message = ( + f"{GENERATED} is out of date with the supervisor wire schema " + "(task-sdk/src/airflow/sdk/execution_time/schema/schema.json).\n" + "Regenerate it with `pnpm run generate:supervisor` in ts-sdk/ and commit the result." + ) + if console: + console.print(f"[red]{message}[/]") + console.print(diff.stdout) + else: + print(message, file=sys.stderr) + print(diff.stdout, file=sys.stderr) + raise SystemExit(1) diff --git a/ts-sdk/scripts/generate-supervisor.mjs b/ts-sdk/scripts/generate-supervisor.mjs new file mode 100644 index 0000000000000..d13f93b5e107f --- /dev/null +++ b/ts-sdk/scripts/generate-supervisor.mjs @@ -0,0 +1,128 @@ +#!/usr/bin/env node +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// Codegen for the Airflow supervisor wire schema. +// +// Reads the canonical supervisor schema from Airflow's Task SDK and emits +// `src/generated/supervisor.ts`. +// +// The input file is Airflow's canonical supervisor JSON Schema. +// We wrap its top-level `$defs` into a synthetic schema so +// json-schema-to-typescript treats each entry as an exported interface, +// preserving cross-references between schema definitions. + +import { compile } from "json-schema-to-typescript"; +import { readFileSync, writeFileSync, mkdirSync } from "node:fs"; +import { dirname, join } from "node:path"; +import { fileURLToPath } from "node:url"; + +const HERE = dirname(fileURLToPath(import.meta.url)); +const ROOT = join(HERE, ".."); +const SCHEMA_PATH = join(ROOT, "../task-sdk/src/airflow/sdk/execution_time/schema/schema.json"); +const OUT_PATH = join(ROOT, "src/generated/supervisor.ts"); + +const HEADER = `/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// AUTO-GENERATED by scripts/generate-supervisor.mjs — do not edit by hand. +// Source: ../task-sdk/src/airflow/sdk/execution_time/schema/schema.json +// +// Re-run with: pnpm run generate:supervisor +`; + +const raw = JSON.parse(readFileSync(SCHEMA_PATH, "utf8")); +const { api_version: apiVersion, $defs: schemas } = raw; + +if (!apiVersion || !schemas) { + throw new Error(`Unexpected schema shape: missing api_version or $defs`); +} + +// Wrap into a single document with $defs so each named schema becomes +// an exported interface, and intra-schema refs ($defs/AssetResponse etc.) +// still resolve. `unreachableDefinitions: true` below emits definitions +// even when the synthetic root does not reference them directly. +const root = { + title: "SupervisorWireSchema", + type: "object", + $defs: Object.fromEntries( + Object.entries(schemas).flatMap(([name, sch]) => { + // Hoist each top-level schema into root.$defs; also hoist its + // own $defs alongside, deduplicating by name. The Airflow + // snapshot embeds the same nested types (AssetResponse, etc.) + // under multiple parents; one merged copy is enough. + const entries = [[name, stripDefs(sch)]]; + for (const [defName, defSchema] of Object.entries(sch.$defs ?? {})) { + entries.push([defName, defSchema]); + } + return entries; + }), + ), + additionalProperties: false, +}; + +function stripDefs(sch) { + // eslint-disable-next-line @typescript-eslint/no-unused-vars + const { $defs: _, ...rest } = sch; + return rest; +} + +// After flattening, every type lives under the root $defs, so "$ref": +// "#/$defs/X" still points at the right schema. No rewrite is needed. + +const ts = await compile(root, "SupervisorWireSchema", { + bannerComment: "", + additionalProperties: false, + style: { tabWidth: 2 }, + declareExternallyReferenced: true, + unreachableDefinitions: true, // emit unreferenced $defs too + enableConstEnums: false, + unknownAny: true, +}); + +const constants = ` +/** Cadwyn schema version this SDK was generated against. + * Not transmitted on the wire — the supervisor learns it out-of-band + * (e.g. bundle metadata) and runs the migrator accordingly. + * Exposed so the SDK author / operator can confirm which schema + * version their build is pinned to. */ +export const SUPERVISOR_API_VERSION = ${JSON.stringify(apiVersion)} as const; +`; + +mkdirSync(dirname(OUT_PATH), { recursive: true }); +writeFileSync(OUT_PATH, HEADER + "\n" + ts + constants, "utf8"); + +console.log(`wrote ${OUT_PATH}`); +console.log(` api_version=${apiVersion}, schemas=${Object.keys(schemas).length}`); diff --git a/ts-sdk/src/coordinator/client.ts b/ts-sdk/src/coordinator/client.ts new file mode 100644 index 0000000000000..2d71976141654 --- /dev/null +++ b/ts-sdk/src/coordinator/client.ts @@ -0,0 +1,193 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import type { CommChannel } from "./comm-channel.js"; +import type { LogChannel } from "./log-channel.js"; +import type { TaskContext } from "../sdk/task.js"; +import type { TaskClient } from "../sdk/client.js"; +import type { ConnectionResult, GetXComOpts, SetXComOpts } from "../sdk/client-types.js"; +import { VariableNotFoundError } from "../sdk/client.js"; +import type { + GetVariable, + GetXCom, + SetXCom, + GetConnection, + ConnectionResult as WireConnectionResult, +} from "./protocol.js"; + +function resolveWireMapIndex( + requestedMapIndex: number | null | undefined, + contextMapIndex: number, +): number | null { + // `mapIndex` is nullable, so only use the context value when the user did + // not provide one. If the user passes null, send null to the supervisor. + const mapIndex = requestedMapIndex === undefined ? contextMapIndex : requestedMapIndex; + return mapIndex == null || mapIndex < 0 ? null : mapIndex; +} + +function fromWireConnection(body: WireConnectionResult): ConnectionResult { + return { + id: body.conn_id, + type: body.conn_type, + host: body.host ?? null, + schema: body.schema ?? null, + login: body.login ?? null, + password: body.password ?? null, + port: body.port ?? null, + extra: body.extra ?? null, + }; +} + +export function createCoordinatorClient( + comm: CommChannel, + ctx: TaskContext, + logs: LogChannel | null = null, +): TaskClient { + async function rpc( + op: string, + expectedType: string | null, + request: unknown, + extract: (body: Record | null) => T, + ): Promise { + logs?.debug(`${op} request`); + const frame = await comm.request(request); + const err = parseFrameError(frame); + if (err) { + if (isNotFound(err)) { + logs?.debug(`${op} not found`, { error: err.code }); + return null; + } + logs?.warning(`${op} failed`, { error: err.code }); + throw new Error(`${op} failed: ${err.code}`); + } + const body = frame.body as Record | null; + if (expectedType !== null && body?.type !== expectedType) { + logs?.error(`${op} unexpected response type`, { + expected: expectedType, + got: body?.type ?? null, + }); + throw new Error(`${op}: unexpected response type ${JSON.stringify(body?.type)}`); + } + logs?.debug(`${op} ok`); + return extract(body); + } + + const client: TaskClient = { + // ---- Variables ---- + + async getVariable(key: string): Promise { + const msg: GetVariable = { type: "GetVariable", key }; + return rpc("GetVariable", "VariableResult", msg, (body) => (body!.value as string) ?? null); + }, + + async getVariableOrThrow(key: string): Promise { + const value = await client.getVariable(key); + if (value == null) throw new VariableNotFoundError(key); + return value; + }, + + // ---- XCom ---- + + async getXCom(opts: GetXComOpts): Promise { + const msg: GetXCom = { + type: "GetXCom", + key: opts.key, + dag_id: opts.dagId ?? ctx.dagId, + task_id: opts.taskId ?? ctx.taskId, + run_id: opts.runId ?? ctx.runId, + map_index: resolveWireMapIndex(opts.mapIndex, ctx.mapIndex), + include_prior_dates: opts.includePriorDates ?? false, + }; + return rpc("GetXCom", "XComResult", msg, (body) => (body!.value as T) ?? null); + }, + + async setXCom(opts: SetXComOpts): Promise { + const msg: SetXCom = { + type: "SetXCom", + key: opts.key, + value: opts.value, + dag_id: opts.dagId ?? ctx.dagId, + task_id: opts.taskId ?? ctx.taskId, + run_id: opts.runId ?? ctx.runId, + map_index: resolveWireMapIndex(opts.mapIndex, ctx.mapIndex), + }; + await rpc("SetXCom", null, msg, () => undefined); + }, + + // ---- Connections ---- + + async getConnection(connId: string): Promise { + const msg: GetConnection = { type: "GetConnection", conn_id: connId }; + return rpc("GetConnection", "ConnectionResult", msg, (body) => + fromWireConnection(body as unknown as WireConnectionResult), + ); + }, + }; + return client; +} + +// -------- Error handling (two functions) -------- +// +// parseFrameError: extract a structured error from the frame (once). +// isNotFound: decide if the error means "absent" (return null to caller) +// or "failed" (throw). + +interface FrameError { + code: string; + statusCode?: number; +} + +/** Extract the error code and optional HTTP status from a response frame. + * Returns `null` for non-error frames. */ +function parseFrameError(frame: { body: unknown; error?: unknown }): FrameError | null { + // Case 1: error field on the frame itself + if (frame.error != null) { + if (typeof frame.error === "string") return { code: frame.error }; + if (typeof frame.error === "object") { + const e = frame.error as Record; + if (typeof e.error === "string") { + const detail = e.detail as { status_code?: number } | undefined; + return { code: e.error, statusCode: detail?.status_code }; + } + } + } + // Case 2: ErrorResponse in body + const body = frame.body as Record | null; + if (body?.type === "ErrorResponse" && typeof body.error === "string") { + const detail = body.detail as { status_code?: number } | undefined; + return { code: body.error, statusCode: detail?.status_code }; + } + return null; +} + +// Exact supervisor ErrorType codes that mean "absent", not "failed". +const NOT_FOUND_CODES = new Set(["VARIABLE_NOT_FOUND", "XCOM_NOT_FOUND", "CONNECTION_NOT_FOUND"]); + +/** Is this error a "not found" (caller should get null, not a throw)? + * Covers both dedicated NOT_FOUND codes and API_SERVER_ERROR with 404. */ +function isNotFound(err: FrameError): boolean { + if (NOT_FOUND_CODES.has(err.code)) return true; + // The supervisor wraps API server 404s as API_SERVER_ERROR with + // detail.status_code=404 (supervisor.py: WatchedSubprocess.handle_requests). + // Dag / Dag run lookups hit this path. + // TODO: If the TS client adds APIs beyond variables, XCom, and connections, + // make not-found handling operation-specific instead of treating every + // API_SERVER_ERROR 404 as null. + return err.code === "API_SERVER_ERROR" && err.statusCode === 404; +} diff --git a/ts-sdk/src/coordinator/comm-channel.ts b/ts-sdk/src/coordinator/comm-channel.ts new file mode 100644 index 0000000000000..858265f7d5720 --- /dev/null +++ b/ts-sdk/src/coordinator/comm-channel.ts @@ -0,0 +1,266 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// Comm socket client — length-prefixed msgpack frames over TCP. +// Mirrors the Airflow supervisor's comm socket protocol. +// +// The channel is the sole reader on the socket. The task sends +// requests and awaits id-correlated replies. The supervisor's only +// unprompted frame is the greeting (StartupDetails / +// DagFileParseRequest), pre-caught into the `greeting` promise that +// `connect()` awaits — the protocol sends nothing else +// supervisor-initiated (comms.py: "No messages are sent to task +// process except in response to a request"). + +import type { Socket } from "node:net"; +import { encodeRequest, encodeResponse, FrameReader, type Frame } from "./frames.js"; +import { connectTcp } from "./tcp-connect.js"; +import { Deferred } from "./deferred.js"; +import type { LogChannel } from "./log-channel.js"; + +/** What `CommChannel.connect` resolves to: the live channel plus the + * supervisor's first frame (StartupDetails / DagFileParseRequest), + * already in hand so the caller never has to manage a "frame arrived + * with no consumer" window. */ +export interface CommConnection { + channel: CommChannel; + firstFrame: Frame; +} + +export interface SendResponseOptions { + timeoutMs?: number; +} + +export interface RequestOptions { + timeoutMs?: number; +} + +export const COORDINATOR_REQUEST_TIMEOUT_MS = 30_000; + +export class CommChannel { + private readonly sock: Socket; + private readonly reader = new FrameReader(); + private readonly logs: LogChannel | null; + private nextId = 0; + private pendingReplies = new Map void>(); + private closed = false; + private closeError: Error | null = null; + + // The greeting (first supervisor-initiated frame). The promise is + // its own buffer: arriving before `connect()` awaits is fine — it + // stays settled with the value, so there is no race to handle. + private readonly greeting = new Deferred(); + + private constructor(sock: Socket, logs: LogChannel | null) { + this.sock = sock; + this.logs = logs; + // A `new Socket()` (from `connectTcp`) starts paused: it buffers + // inbound bytes and emits no `data` until a listener attaches + // and flips it to flowing. Attaching synchronously here — same + // tick as construction, before the event loop can deliver a + // read, and as the only reader — loses nothing, double-reads + // nothing. + sock.on("data", (chunk) => this.handleData(chunk)); + sock.on("close", () => this.handleClose(null)); + sock.on("error", (err) => this.handleClose(err)); + } + + /** Connect and wait for the supervisor's greeting; rejects if the + * socket dies before it arrives. */ + static async connect(addr: string, logs: LogChannel | null = null): Promise { + const sock = await connectTcp(addr); + const channel = new CommChannel(sock, logs); + const firstFrame = await channel.greeting.promise; + return { channel, firstFrame }; + } + + /** Send a request to the supervisor and await its matching response. */ + async request(body: unknown, opts: RequestOptions = {}): Promise { + if (this.closed) { + throw this.closeError ?? new Error("Comm channel closed"); + } + + const id = this.nextId++; + const type = describeFrameType(body); + const timeoutMs = opts.timeoutMs ?? COORDINATOR_REQUEST_TIMEOUT_MS; + + // A reply that times out on its own and removes its pending entry once done. + const reply = new Deferred() + .rejectAfter( + timeoutMs, + () => new Error(`Timed out waiting for ${type} response after ${timeoutMs} ms`), + ) + .onSettle(() => this.pendingReplies.delete(id)); + + // How the reply gets fulfilled: the reader matches this id and calls us. + this.pendingReplies.set(id, (frame) => { + this.logs?.debug("Response received", { + id, + request_type: type, + response_type: describeFrameType(frame.body), + error: frame.error ?? null, + }); + reply.resolve(frame); + }); + + this.logs?.debug("Sending request", { id, type }); + try { + this.sock.write(encodeRequest(id, body), (err) => { + if (err) reply.reject(err); + }); + } catch (err) { + reply.reject(err as Error); + } + + return reply.promise; + } + + /** Send a response for an incoming supervisor request. */ + async sendResponse( + id: number, + body: unknown, + error?: unknown, + opts: SendResponseOptions = {}, + ): Promise { + this.logs?.debug("Sending response", { + id, + type: describeFrameType(body), + error: error ?? null, + }); + if (this.closed) { + throw this.closeError ?? new Error("Comm channel closed"); + } + + const done = new Deferred(); + if (opts.timeoutMs !== undefined) { + // A wedged terminal write must not hang the process: destroy the socket + // and fail so the runtime exits non-zero instead of waiting forever. + done.rejectAfter(opts.timeoutMs, () => { + const err = new Error(`Timed out sending response after ${opts.timeoutMs} ms`); + this.sock.destroy(err); + return err; + }); + } + + try { + this.sock.write(encodeResponse(id, body, error), (err) => + err ? done.reject(err) : done.resolve(), + ); + } catch (err) { + done.reject(err as Error); + } + + return done.promise; + } + + async close(): Promise { + return new Promise((resolve) => { + if (this.closed) { + resolve(); + return; + } + this.sock.end(() => resolve()); + }); + } + + // -- internals -- + + private handleData(chunk: Buffer): void { + let frames: Frame[]; + try { + frames = this.reader.push(chunk); + } catch (err) { + // Frame decode failure — protocol violation or socket + // corruption. Surface it so it's not a silent dropped chunk. + this.logs?.error("Frame decode failed", { + error: (err as Error).message ?? String(err), + pending_bytes: this.reader.pending, + }); + this.handleClose(err as Error); + return; + } + for (const frame of frames) { + this.logs?.debug("Handling frame", { id: frame.id }); + this.route(frame); + } + } + + private route(frame: Frame): void { + // Route by pending-request lookup, not frame arity. If the id + // matches a request we sent, it's the response. Otherwise it's + // supervisor-initiated (the greeting). This works because the + // greeting always arrives before any request is sent, so id=0 + // can never collide with a pending request. + const pending = this.pendingReplies.get(frame.id); + if (pending) { + this.pendingReplies.delete(frame.id); + pending(frame); + return; + } + this.deliverSupervisorFrame(frame); + } + + private deliverSupervisorFrame(frame: Frame): void { + // The supervisor's only unprompted frame is the greeting. + if (!this.greeting.settled) { + this.greeting.resolve(frame); + return; + } + // Anything else supervisor-initiated is a protocol anomaly — + // comms.py guarantees "No messages are sent to task process + // except in response to a request". Surface it; never buffer. + this.logs?.error("Unexpected supervisor-initiated frame after greeting", { + id: frame.id, + type: describeFrameType(frame.body), + }); + } + + private handleClose(err: Error | null): void { + if (this.closed) return; + this.closed = true; + this.closeError = err; + if (err) { + this.logs?.warning("Comm channel closed with error", { + error: err.message, + pending_replies: this.pendingReplies.size, + }); + } + // Before the greeting this rejects so `connect()` throws; + // after it, a no-op — the Deferred settles at most once, so no + // guard is needed here. + this.greeting.reject(err ?? new Error("Comm channel closed before first frame")); + for (const [, resolver] of this.pendingReplies) { + resolver({ + id: -1, + body: null, + error: err?.message ?? "closed", + isResponse: true, + }); + } + this.pendingReplies.clear(); + } +} + +function describeFrameType(body: unknown): string { + if (body && typeof body === "object" && "type" in body) { + const t = (body as { type?: unknown }).type; + if (typeof t === "string") return t; + } + return "unknown"; +} diff --git a/ts-sdk/src/coordinator/deferred.ts b/ts-sdk/src/coordinator/deferred.ts new file mode 100644 index 0000000000000..da524b6db47b8 --- /dev/null +++ b/ts-sdk/src/coordinator/deferred.ts @@ -0,0 +1,92 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// Coordinator-scoped (like `tcp-connect.ts`): its only caller is +// `comm-channel.ts`, so it lives next to it rather than in a shared util. + +/** + * A promise whose settlement is triggered from *outside* its executor + * (the classic "deferred"). Use it when the producer and consumer of a + * one-time value are decoupled in time and code location — e.g. a value + * that arrives on a socket but is awaited elsewhere. + * + * Owns the settle-at-most-once invariant: `resolve`/`reject` after the + * first settlement are no-ops, so callers never hand-maintain that rule. + */ +export class Deferred { + readonly promise: Promise; + private done = false; + private res!: (v: T) => void; + private rej!: (e: Error) => void; + private readonly onSettleFns: Array<() => void> = []; + + constructor() { + this.promise = new Promise((res, rej) => { + this.res = res; + this.rej = rej; + }); + } + + /** Whether `resolve`/`reject` has fired — for callers that need to + * branch on it. Not needed to guard `resolve`/`reject`; those are + * already idempotent. */ + get settled(): boolean { + return this.done; + } + + /** Run `fn` once, when this settles either way — the deferred's + * `finally`. Runs immediately if already settled. Returns `this` so a + * timeout / cleanup can be attached fluently at construction. */ + onSettle(fn: () => void): this { + if (this.done) { + fn(); + } else { + this.onSettleFns.push(fn); + } + return this; + } + + /** Reject with `makeError()` after `ms`, unless already settled. The + * timer auto-clears on settle, so a fulfilled deferred never fires it. */ + rejectAfter(ms: number, makeError: () => Error): this { + const timer = setTimeout(() => this.reject(makeError()), ms); + return this.onSettle(() => clearTimeout(timer)); + } + + resolve(v: T): void { + this.settle(() => this.res(v)); + } + + reject(e: Error): void { + this.settle(() => this.rej(e)); + } + + private settle(apply: () => void): void { + if (this.done) return; + this.done = true; + try { + apply(); + } finally { + for (const fn of this.onSettleFns) { + fn(); + } + this.onSettleFns.length = 0; + } + } +} diff --git a/ts-sdk/src/coordinator/frames.ts b/ts-sdk/src/coordinator/frames.ts new file mode 100644 index 0000000000000..397c3da512cf7 --- /dev/null +++ b/ts-sdk/src/coordinator/frames.ts @@ -0,0 +1,147 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// Task SDK wire frame codec. +// +// Mirrors the Airflow supervisor's length-prefixed msgpack IPC format. +// +// Frame format on the wire: +// +// +---------+----------------------------+ +// | len (4) | msgpack payload (variable) | +// +---------+----------------------------+ +// +// `len` is big-endian uint32 — the byte length of the msgpack payload. +// Payload is a msgpack array: +// - Request: [id: int, body: map] (arity 2) +// - Response: [id: int, body: map, error: map?] (arity 3) +// +// The body is a map with a `type` key naming the message. Both +// supervisor and runtime maintain independent id counters starting +// at 0. Routing uses the pending-request map (not arity) — if the +// id matches a request we sent, it's the response; otherwise it's +// supervisor-initiated. See `comm-channel.ts`. + +import { encode, decode } from "@msgpack/msgpack"; + +/** + * Maximum frame payload size in bytes (2^32 − 1). The length prefix is a + * big-endian uint32, so anything larger would silently truncate. + */ +export const MAX_FRAME_SIZE = 0xffff_ffff; + +export interface Frame { + id: number; + body: unknown; + error?: unknown; + /** Whether the frame's msgpack array was arity 3 (response) or 2 (request). */ + isResponse: boolean; +} + +export function encodeRequest(id: number, body: unknown): Buffer { + return encodeFrame(id, body, undefined, /* isResponse */ false); +} + +export function encodeResponse(id: number, body?: unknown, error?: unknown): Buffer { + return encodeFrame(id, body, error, /* isResponse */ true); +} + +function encodeFrame(id: number, body: unknown, error: unknown, isResponse: boolean): Buffer { + const array = isResponse ? [id, body ?? null, error ?? null] : [id, body ?? null]; + const payload = Buffer.from(encode(array)); + if (payload.length > MAX_FRAME_SIZE) { + throw new RangeError( + `Frame payload ${payload.length} bytes exceeds MAX_FRAME_SIZE (${MAX_FRAME_SIZE})`, + ); + } + const framed = Buffer.alloc(4 + payload.length); + framed.writeUInt32BE(payload.length, 0); + payload.copy(framed, 4); + return framed; +} + +/** Decode a single framed payload (length prefix already stripped). */ +export function decodePayload(payload: Buffer): Frame { + const value = decode(payload); + if (!Array.isArray(value)) { + throw new Error(`Expected msgpack array frame, got ${typeof value}`); + } + const arity = value.length; + if (arity < 2) { + throw new Error(`Unexpected Task SDK frame arity ${arity}`); + } + const [id, body, error] = value as [number, unknown, unknown?]; + if (typeof id !== "number") { + throw new Error(`Frame id must be number, got ${typeof id}`); + } + // Specific per-failure messages above are deliberate — this is a + // cross-language wire boundary and a vague decode error there costs + // hours. Construction itself is a single expression: error is + // omitted (not set to null) when absent so `"error" in frame` + // stays a faithful arity signal. + return { + id, + body, + isResponse: arity >= 3, + ...(error != null ? { error } : {}), + }; +} + +/** + * Try to consume one full frame (length prefix + payload) from `buf`. + * Returns the frame plus any trailing bytes belonging to the next frame, + * or `null` if `buf` doesn't yet contain a complete frame. + */ +export function tryTakeFrame(buf: Buffer): { frame: Frame; rest: Buffer } | null { + if (buf.length < 4) return null; + const len = buf.readUInt32BE(0); + if (len > MAX_FRAME_SIZE) { + throw new RangeError( + `Incoming frame length ${len} bytes exceeds MAX_FRAME_SIZE (${MAX_FRAME_SIZE})`, + ); + } + if (buf.length < 4 + len) return null; + const payload = buf.subarray(4, 4 + len); + const rest = buf.subarray(4 + len); + return { + frame: decodePayload(Buffer.from(payload)), + rest: Buffer.from(rest), + }; +} + +/** Accumulate bytes across socket reads and emit complete frames. */ +export class FrameReader { + private buf: Buffer = Buffer.alloc(0); + + push(chunk: Buffer): Frame[] { + this.buf = this.buf.length === 0 ? chunk : Buffer.concat([this.buf, chunk]); + const frames: Frame[] = []; + while (true) { + const taken = tryTakeFrame(this.buf); + if (!taken) break; + frames.push(taken.frame); + this.buf = taken.rest; + } + return frames; + } + + get pending(): number { + return this.buf.length; + } +} diff --git a/ts-sdk/src/coordinator/index.ts b/ts-sdk/src/coordinator/index.ts new file mode 100644 index 0000000000000..815e0c849b972 --- /dev/null +++ b/ts-sdk/src/coordinator/index.ts @@ -0,0 +1,29 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// Coordinator-mode public API. Re-exported through the package root. +// +// TaskClient and related types are exported from the package root. This +// barrel only exports coordinator-specific entry points. + +export { startCoordinator, type StartCoordinatorOptions } from "./runtime.js"; +/** Cadwyn schema version this SDK was generated against. Not sent on + * the wire — exposed so callers can read it for bundle metadata, + * health checks, or to confirm which schema their build is pinned to. */ +export { SUPERVISOR_API_VERSION } from "./protocol.js"; diff --git a/ts-sdk/src/coordinator/log-channel.ts b/ts-sdk/src/coordinator/log-channel.ts new file mode 100644 index 0000000000000..e64737e3b671b --- /dev/null +++ b/ts-sdk/src/coordinator/log-channel.ts @@ -0,0 +1,127 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// Log channel — newline-delimited JSON log records over TCP. +// +// The Airflow coordinator's `_bridge` reads lines from this socket, +// parses each as JSON, and re-emits through structlog using the same +// handler used for ordinary Python task logs +// (`process_log_messages_from_subprocess`). Required fields are +// `event`, `level`, `logger`, `timestamp`. Extra fields pass through +// as structured log keys. +// +// The `logger` field becomes the bracketed name column in structlog's +// ConsoleRenderer output (e.g. `[ts-sdk.runtime] Coordinator runtime +// started`). Use hierarchical names — `ts-sdk.runtime`, `ts-sdk.comm`, +// `ts-sdk.client` — so SDK-emitted lines are visibly distinct from +// user task logs (which typically use the task's module name). + +import type { Socket } from "node:net"; +import { connectTcp } from "./tcp-connect.js"; + +export type LogLevel = "debug" | "info" | "warning" | "error"; + +export interface LogRecord { + event: string; + level: LogLevel; + logger: string; + timestamp: string; + [key: string]: unknown; +} + +const DEFAULT_LOGGER_NAME = "ts-sdk"; + +export class LogChannel { + private readonly sock: Socket; + private readonly name: string; + private readonly isRoot: boolean; + + private constructor(sock: Socket, name: string, isRoot: boolean) { + this.sock = sock; + this.name = name; + this.isRoot = isRoot; + if (isRoot) { + sock.on("error", (err) => { + process.stderr.write(`[${this.name}] log socket error: ${err.message}\n`); + }); + } + } + + static async connect(addr: string, name: string = DEFAULT_LOGGER_NAME): Promise { + return new LogChannel(await connectTcp(addr), name, true); + } + + /** Create a sibling logger that shares the underlying socket but + * carries a hierarchical name (`parent.suffix`). Only the root + * owns the socket — children's `close()` is a no-op. */ + child(suffix: string): LogChannel { + return new LogChannel(this.sock, `${this.name}.${suffix}`, false); + } + + /** Name reported in the `logger` field of every record this + * instance emits. Useful for tests. */ + get loggerName(): string { + return this.name; + } + + send( + record: Omit & { + timestamp?: string; + logger?: string; + }, + ): void { + // Drop late records after the log socket has closed. + if (this.sock.writableEnded) return; + // Prepend the logger name to the event message so it surfaces in + // the Airflow UI task log view, which renders the message text but + // hides the `logger` JSON field. The field is still emitted for + // JSON consumers (grep/jq). Remove the prefix here if the + // supervisor-side renderer ever starts showing the logger column. + const line = JSON.stringify({ + logger: this.name, + ...record, + event: `[${this.name}] ${record.event}`, + timestamp: record.timestamp ?? new Date().toISOString(), + }); + this.sock.write(Buffer.from(line + "\n", "utf8")); + } + + debug(event: string, args: Record = {}): void { + this.send({ event, level: "debug", ...args }); + } + + info(event: string, args: Record = {}): void { + this.send({ event, level: "info", ...args }); + } + + warning(event: string, args: Record = {}): void { + this.send({ event, level: "warning", ...args }); + } + + error(event: string, args: Record = {}): void { + this.send({ event, level: "error", ...args }); + } + + async close(): Promise { + if (!this.isRoot) return; + return new Promise((resolve) => { + this.sock.end(() => resolve()); + }); + } +} diff --git a/ts-sdk/src/coordinator/protocol.ts b/ts-sdk/src/coordinator/protocol.ts new file mode 100644 index 0000000000000..44d0aee172356 --- /dev/null +++ b/ts-sdk/src/coordinator/protocol.ts @@ -0,0 +1,114 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// Coordinator message types over the generated supervisor wire schema. +// +// The generated schema is authoritative for field names and payload shapes. +// This module keeps the runtime decoder next to the generated type exports +// used by the coordinator client. + +import type { + StartupDetails, + DagFileParseRequest, + ErrorResponse, + DagFileParsingResult, + RetryTask, + SucceedTask, + TaskState, +} from "../generated/supervisor.js"; + +export { SUPERVISOR_API_VERSION } from "../generated/supervisor.js"; + +// -------- Re-exports — supervisor frames -------- + +export type { + StartupDetails, + DagFileParseRequest, + ErrorResponse, +} from "../generated/supervisor.js"; + +// -------- Re-exports — runtime responses -------- + +export type { + DagFileParsingResult, + RetryTask, + SucceedTask, + TaskState, +} from "../generated/supervisor.js"; + +// -------- Re-exports — client RPC payloads -------- + +export type { + VariableResult, + XComResult, + ConnectionResult, + GetVariable, + GetXCom, + SetXCom, + GetConnection, +} from "../generated/supervisor.js"; + +// -------- Frames from supervisor -------- + +type WithRequiredType = Omit & { + type: NonNullable; +}; + +export type MsgFromSupervisor = + | WithRequiredType + | WithRequiredType + | WithRequiredType; + +// -------- Frames to supervisor -------- + +export type RuntimeTaskState = WithRequiredType; +export type RuntimeRetryTask = WithRequiredType; +export type RuntimeSucceedTask = WithRequiredType & { + task_outlets: NonNullable; + outlet_events: NonNullable; +}; +export type RuntimeDagFileParsingResult = WithRequiredType; + +// -------- Decoder: raw map → typed message -------- + +export function asMsgFromSupervisor(raw: unknown): MsgFromSupervisor { + const body = normalizeBody(raw); + switch (body.type) { + case "StartupDetails": + case "DagFileParseRequest": + case "ErrorResponse": + return body as unknown as MsgFromSupervisor; + default: + throw new Error(`Unsupported supervisor message type: ${JSON.stringify(body.type)}`); + } +} + +function normalizeBody(raw: unknown): { type: string; [k: string]: unknown } { + if (raw === null || typeof raw !== "object") { + throw new Error(`Frame body must be a map, got ${typeof raw}`); + } + const mapLike = raw as Record; + const type = mapLike["type"]; + if (typeof type !== "string") { + throw new Error( + `Frame body missing string 'type'; got keys: ${Object.keys(mapLike).join(",")}`, + ); + } + return { ...mapLike, type }; +} diff --git a/ts-sdk/src/coordinator/runtime.ts b/ts-sdk/src/coordinator/runtime.ts new file mode 100644 index 0000000000000..d083bf4acce85 --- /dev/null +++ b/ts-sdk/src/coordinator/runtime.ts @@ -0,0 +1,369 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// Coordinator runtime entrypoint. +// +// Invoked by Airflow's coordinator subprocess path: +// +// node my-bundle.mjs --comm=host:port --logs=host:port +// +// where `my-bundle.mjs` is a user-bundled Node script that imports +// the SDK, calls `registerTask(...)` for each handler, then calls +// `startCoordinator()`. +// +// Lifecycle: +// 1. Parse --comm / --logs from argv +// 2. Connect both TCP sockets +// 3. Read the first frame from comm: +// - DagFileParseRequest → respond with DagFileParsingResult, exit +// - StartupDetails → run task, respond Succeed or Fail, exit +// +import { createCoordinatorClient } from "./client.js"; +import { CommChannel } from "./comm-channel.js"; +import { LogChannel } from "./log-channel.js"; +import { + asMsgFromSupervisor, + SUPERVISOR_API_VERSION, + type RuntimeDagFileParsingResult, + type RuntimeRetryTask, + type RuntimeSucceedTask, + type RuntimeTaskState, + type StartupDetails, +} from "./protocol.js"; +import { getRegisteredTask, listRegisteredTasks } from "../sdk/registry.js"; +import type { TaskContext, TaskHandlerArgs } from "../sdk/task.js"; +import type { JsonValue } from "../sdk/client-types.js"; + +export const ABORT_GRACE_PERIOD_MS = 30_000; +export const COORDINATOR_RESPONSE_TIMEOUT_MS = 30_000; + +/** Options for `startCoordinator()`. */ +export interface StartCoordinatorOptions { + /** Comm socket address (host:port). Must be supplied together with `logsAddr`; otherwise parsed from argv. */ + commAddr?: string; + /** Logs socket address (host:port). Must be supplied together with `commAddr`; otherwise parsed from argv. */ + logsAddr?: string; + /** Source argv. Defaults to `process.argv`. */ + argv?: readonly string[]; +} + +interface ParsedArgs { + commAddr: string; + logsAddr: string; +} + +type SignalListener = NodeJS.SignalsListener; + +interface ProcessSignalSource { + on(signal: NodeJS.Signals, listener: SignalListener): ProcessSignalSource; + off(signal: NodeJS.Signals, listener: SignalListener): ProcessSignalSource; +} + +interface RuntimeAbortOptions { + signalSource?: ProcessSignalSource; + exitProcess?: (code: number) => never; +} + +export interface RuntimeAbort { + readonly signal: AbortSignal; + dispose(): void; +} + +const ABORT_SIGNALS: readonly NodeJS.Signals[] = ["SIGTERM", "SIGINT"]; +const ABORT_FORCE_EXIT_CODE = 1; + +export function parseArgs(argv: readonly string[]): ParsedArgs { + let commAddr: string | null = null; + let logsAddr: string | null = null; + for (const arg of argv) { + if (arg.startsWith("--comm=")) { + commAddr = arg.slice("--comm=".length); + } else if (arg.startsWith("--logs=")) { + logsAddr = arg.slice("--logs=".length); + } + } + if (!commAddr) throw new Error("Missing --comm=host:port"); + if (!logsAddr) throw new Error("Missing --logs=host:port"); + return { commAddr, logsAddr }; +} + +/** Start the coordinator runtime. Resolves when the subprocess has + * delivered its terminal frame and closed both sockets. */ +export async function startCoordinator(opts: StartCoordinatorOptions = {}): Promise { + const argv = opts.argv ?? process.argv; + const parsed = + opts.commAddr && opts.logsAddr + ? { commAddr: opts.commAddr, logsAddr: opts.logsAddr } + : parseArgs(argv); + + let logs: LogChannel | null = null; + let comm: CommChannel | null = null; + let runtimeAbort: RuntimeAbort | null = null; + + try { + // Connect log channel first so early failures are captured. + // Root logger is `ts-sdk`; subsystems use child names (`ts-sdk.runtime`, + // `ts-sdk.comm`, `ts-sdk.client`) so structlog's ConsoleRenderer prints + // them as a distinct `[name]` column on the supervisor side. + logs = await LogChannel.connect(parsed.logsAddr); + const runtimeLogs = logs.child("runtime"); + const tasks = listRegisteredTasks(); + runtimeLogs.info("Coordinator runtime started", { + registered_tasks: tasks, + count: tasks.length, + // Cadwyn schema version this SDK was generated against. Logged + // for operator visibility; not sent on the wire. + supervisor_api_version: SUPERVISOR_API_VERSION, + }); + + const connection = await CommChannel.connect(parsed.commAddr, logs.child("comm")); + comm = connection.channel; + const firstFrame = connection.firstFrame; + runtimeLogs.debug("Connected comm socket", { commAddr: parsed.commAddr }); + runtimeAbort = createRuntimeAbort(runtimeLogs); + + const body = asMsgFromSupervisor(firstFrame.body); + + if (body.type === "DagFileParseRequest") { + runtimeLogs.info("Received Dag parse request", { + file: body.file, + bundle_path: body.bundle_path, + }); + const response = handleParse(body, runtimeLogs); + await sendSupervisorResponse(firstFrame.id, response, comm, runtimeLogs); + } else if (body.type === "StartupDetails") { + runtimeLogs.info("Received task startup details", { + dag_id: body.ti.dag_id, + task_id: body.ti.task_id, + run_id: body.ti.run_id, + try_number: body.ti.try_number, + map_index: body.ti.map_index ?? -1, + bundle: body.bundle_info.name, + }); + const response = await handleTask( + body, + comm, + runtimeLogs, + logs.child("client"), + runtimeAbort.signal, + ); + await sendSupervisorResponse(firstFrame.id, response, comm, runtimeLogs); + if (response.type === "SucceedTask") { + runtimeLogs.info("Task succeeded", { task_id: body.ti.task_id }); + } + } else { + const errMsg = `First frame must be DagFileParseRequest or StartupDetails, got ${body.type}`; + runtimeLogs.error("Unexpected first frame", { type: body.type }); + await sendSupervisorResponse(firstFrame.id, null, comm, runtimeLogs, { + error: "protocol_error", + detail: errMsg, + }); + } + } finally { + runtimeAbort?.dispose(); + await comm?.close(); + await logs?.close(); + } +} + +export function createRuntimeAbort( + logs: LogChannel | null, + opts: RuntimeAbortOptions = {}, +): RuntimeAbort { + const controller = new AbortController(); + const signalSource = opts.signalSource ?? process; + const exitProcess = opts.exitProcess ?? ((code: number): never => process.exit(code)); + let disposed = false; + let forceExitTimer: ReturnType | null = null; + const signalListeners = new Map(); + + const handleSignal = (signal: NodeJS.Signals): void => { + if (disposed) return; + if (controller.signal.aborted) { + logs?.warning("Additional abort signal received", { signal }); + return; + } + + logs?.warning("Abort signal received", { + signal, + grace_period_ms: ABORT_GRACE_PERIOD_MS, + }); + controller.abort(new Error(`Task aborted by ${signal}`)); + forceExitTimer = setTimeout(() => { + logs?.error("Abort grace period expired; forcing process exit", { + signal, + grace_period_ms: ABORT_GRACE_PERIOD_MS, + }); + exitProcess(ABORT_FORCE_EXIT_CODE); + }, ABORT_GRACE_PERIOD_MS); + }; + + for (const signal of ABORT_SIGNALS) { + const listener: SignalListener = () => handleSignal(signal); + signalListeners.set(signal, listener); + signalSource.on(signal, listener); + } + + return { + signal: controller.signal, + dispose: () => { + disposed = true; + for (const [signal, listener] of signalListeners) { + signalSource.off(signal, listener); + } + signalListeners.clear(); + if (forceExitTimer !== null) { + clearTimeout(forceExitTimer); + forceExitTimer = null; + } + }, + }; +} + +function handleParse( + request: { file: string; bundle_path: string }, + logs: LogChannel, +): RuntimeDagFileParsingResult { + // TypeScript-native Dag parsing is not yet supported. + // Respond with an empty result so the Python-stub-Dag workflow works. + logs.info("Parse-mode response (TS Dag parsing not yet supported)", { + registered_tasks: listRegisteredTasks(), + }); + const response: RuntimeDagFileParsingResult = { + type: "DagFileParsingResult", + fileloc: request.file, + serialized_dags: [], + }; + return response; +} + +async function handleTask( + details: StartupDetails, + comm: CommChannel, + logs: LogChannel, + clientLogs: LogChannel, + signal: AbortSignal, +): Promise { + const ti = details.ti; + const handler = getRegisteredTask(ti.dag_id, ti.task_id); + + if (!handler) { + logs.warning("No handler registered for task", { + dag_id: ti.dag_id, + task_id: ti.task_id, + available: listRegisteredTasks(), + }); + // A missing handler means this bundle cannot run the task, so retrying the + // same bundle/configuration mismatch would not help. + const response: RuntimeTaskState = { + type: "TaskState", + state: "removed", + end_date: new Date().toISOString(), + }; + return response; + } + + const ctx = buildContext(details, signal); + const client = createCoordinatorClient(comm, ctx, clientLogs); + const args: TaskHandlerArgs = { ctx, client }; + // Startup-details fields already logged above (`Received task + // startup details`); this line just marks the handler-call boundary. + logs.debug("Dispatching to handler", { task_id: ctx.taskId }); + + try { + const result = await handler(args); + if (result !== undefined) { + await client.setXCom({ key: "return_value", value: result as JsonValue }); + } + // SucceedTask MUST include task_outlets and outlet_events as + // empty lists — the Execution API's TISuccessStatePayload + // tagged-union validator rejects null for these fields. + const response: RuntimeSucceedTask = { + type: "SucceedTask", + end_date: new Date().toISOString(), + task_outlets: [], + outlet_events: [], + }; + return response; + } catch (err) { + const message = (err as Error).message ?? String(err); + logs.error("Task failed", { + task_id: ctx.taskId, + error: message, + stack: (err as Error).stack ?? null, + }); + return buildFailureResponse(details, message); + } +} + +async function sendSupervisorResponse( + id: number, + body: unknown, + comm: CommChannel, + logs: LogChannel, + error?: unknown, +): Promise { + try { + await comm.sendResponse(id, body, error, { timeoutMs: COORDINATOR_RESPONSE_TIMEOUT_MS }); + } catch (err) { + logs.error("Failed to send response to supervisor", { + response_type: responseType(body), + error: (err as Error).message ?? String(err), + }); + throw err; + } +} + +function buildFailureResponse( + details: StartupDetails, + message: string, +): RuntimeRetryTask | RuntimeTaskState { + const endDate = new Date().toISOString(); + if (details.ti_context.should_retry) { + return { + type: "RetryTask", + end_date: endDate, + retry_reason: message.slice(0, 500), + }; + } + return { + type: "TaskState", + state: "failed", + end_date: endDate, + }; +} + +function responseType(body: unknown): string { + if (body && typeof body === "object" && "type" in body) { + const type = (body as { type?: unknown }).type; + if (typeof type === "string") return type; + } + return "unknown"; +} + +function buildContext(details: StartupDetails, signal: AbortSignal): TaskContext { + return { + dagId: details.ti.dag_id, + taskId: details.ti.task_id, + runId: details.ti.run_id, + tryNumber: details.ti.try_number, + mapIndex: details.ti.map_index ?? -1, + signal, + }; +} diff --git a/ts-sdk/src/coordinator/tcp-connect.ts b/ts-sdk/src/coordinator/tcp-connect.ts new file mode 100644 index 0000000000000..1ca368fe95b98 --- /dev/null +++ b/ts-sdk/src/coordinator/tcp-connect.ts @@ -0,0 +1,53 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// Opening a raw TCP connection to a `host:port` the coordinator handed +// us. Both coordinator channels — the comm socket (`comm-channel.ts`) +// and the log socket (`log-channel.ts`) — connect the same way, so the +// connect dance lives here once instead of being duplicated in each. +// +// Coordinator-scoped on purpose: this is only how coordinator mode +// reaches Airflow, so keep it next to its only callers. + +import { Socket } from "node:net"; + +/** Connect a `setNoDelay` TCP socket to `addr` (`host:port`). Resolves + * with the connected socket; rejects if the address is malformed or + * the connection attempt errors. The returned socket is raw — the + * caller owns all framing and lifecycle from here. */ +export async function connectTcp(addr: string): Promise { + const [host, portStr] = splitHostPort(addr); + return new Promise((resolve, reject) => { + const sock = new Socket(); + sock.once("connect", () => { + sock.setNoDelay(true); + resolve(sock); + }); + sock.once("error", reject); + sock.connect(Number.parseInt(portStr, 10), host); + }); +} + +/** Split a `host:port` address on its last colon (so IPv6 hosts like + * `::1` survive). Exported for unit testing. */ +export function splitHostPort(addr: string): [string, string] { + const idx = addr.lastIndexOf(":"); + if (idx < 0) throw new Error(`Address must be host:port, got ${addr}`); + return [addr.slice(0, idx), addr.slice(idx + 1)]; +} diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts new file mode 100644 index 0000000000000..fe61155f39e4c --- /dev/null +++ b/ts-sdk/src/generated/supervisor.ts @@ -0,0 +1,1871 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// AUTO-GENERATED by scripts/generate-supervisor.mjs — do not edit by hand. +// Source: ../task-sdk/src/airflow/sdk/execution_time/schema/schema.json +// +// Re-run with: pnpm run generate:supervisor + +export type Name = string; +export type Id = number; +export type Timestamp = string; +export type Extra = { + [k: string]: JsonValue; +} | null; +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "JsonValue". + */ +export type JsonValue = unknown; +export type Name1 = string; +export type Uri = string; +export type Group = string; +export type Extra1 = { + [k: string]: JsonValue; +} | null; +export type RunId = string; +export type DagId = string; +export type LogicalDate = string | null; +export type StartDate = string; +export type EndDate = string | null; +export type State = string; +export type DataIntervalStart = string | null; +export type DataIntervalEnd = string | null; +export type PartitionKey = string | null; +export type CreatedDagruns = DagRunAssetReference[]; +export type SourceTaskId = string | null; +export type SourceDagId = string | null; +export type SourceRunId = string | null; +export type SourceMapIndex = number | null; +export type PartitionKey1 = string | null; +export type AssetEvents = AssetEventResponse[]; +export type Type = "AssetEventsResult"; +export type Name2 = string | null; +export type Uri1 = string | null; +export type Type1 = string; +export type Name3 = string; +export type Uri2 = string; +export type Name4 = string; +export type Uri3 = string; +export type Group1 = string; +export type Extra3 = { + [k: string]: JsonValue; +} | null; +export type Type2 = "AssetResult"; +export type Type3 = "AssetStateStoreResult"; +export type Assets = AssetResult[]; +export type Type4 = "AssetsByAliasResult"; +export type State1 = "awaiting_input" | null; +export type Timeout = string | null; +export type NextMethod = string; +export type NextKwargs = { + [k: string]: JsonValue; +} | null; +export type RenderedMapIndex = string | null; +export type Type5 = "AwaitInputTask"; +export type Name5 = string; +export type Version = string | null; +export type VersionData = { + [k: string]: unknown; +} | null; +export type Name6 = string; +export type Type6 = "ClearAssetStateStoreByName"; +export type Uri4 = string; +export type Type7 = "ClearAssetStateStoreByUri"; +export type TiId = string; +export type Type8 = "ClearTaskStateStore"; +export type ConnId = string; +export type ConnType = string; +export type Host = string | null; +export type Schema = string | null; +export type Login = string | null; +export type Password = string | null; +export type Port = number | null; +export type Extra4 = string | null; +export type Type9 = "ConnectionResult"; +export type TiId1 = string; +/** + * @minItems 1 + */ +export type Options = [string, ...string[]]; +export type Subject = string; +export type Body = string | null; +export type Defaults = string[] | null; +export type Multiple = boolean | null; +export type Params = { + [k: string]: unknown; +} | null; +export type AssignedUsers = HITLUser[] | null; +export type Id1 = string; +export type Name7 = string; +export type Type10 = "CreateHITLDetailPayload"; +export type Count = number; +export type Type11 = "DRCount"; +export type Filepath = string; +export type BundleName = string; +export type BundleVersion = string | null; +export type Msg = string | null; +export type DagId1 = string; +export type RunId1 = string; +export type DagId2 = string; +export type RunId2 = string; +export type LogicalDate1 = string | null; +export type DataIntervalStart1 = string | null; +export type DataIntervalEnd1 = string | null; +export type RunAfter = string; +export type StartDate1 = string | null; +export type EndDate1 = string | null; +export type ClearNumber = number; +/** + * Class with DagRun types. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DagRunType". + */ +export type DagRunType = + | "backfill" + | "scheduled" + | "manual" + | "operator_triggered" + | "asset_triggered" + | "asset_materialization"; +/** + * All possible states that a DagRun can be in. + * + * These are "shared" with TaskInstanceState in some parts of the code, + * so please ensure that their values always match the ones with the + * same name in TaskInstanceState. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DagRunState". + */ +export type DagRunState = "queued" | "running" | "success" | "failed"; +export type Conf = { + [k: string]: unknown; +} | null; +export type TriggeringUserName = string | null; +export type SourceTaskId1 = string | null; +export type SourceDagId1 = string | null; +export type SourceRunId1 = string | null; +export type SourceMapIndex1 = number | null; +export type SourceAliases = AssetAliasReferenceAssetEventDagRun[]; +export type Timestamp1 = string; +export type PartitionKey2 = string | null; +export type ConsumedAssetEvents = AssetEventDagRunReference[]; +export type PartitionKey3 = string | null; +export type PartitionDate = string | null; +export type Note = string | null; +export type TeamName = string | null; +export type Id2 = string; +export type TaskId = string; +export type DagId3 = string; +export type RunId3 = string; +export type TryNumber = number; +export type DagVersionId = string; +export type MapIndex = number; +export type Hostname = string | null; +export type ContextCarrier = { + [k: string]: unknown; +} | null; +export type Queue = string; +export type IsFailureCallback = boolean | null; +export type Type12 = "DagCallbackRequest"; +export type File = string; +export type BundlePath = string; +export type BundleName1 = string; +export type Filepath1 = string; +export type BundleName2 = string; +export type BundleVersion1 = string | null; +export type Msg1 = string | null; +/** + * All possible states that a Task Instance can be in. + * + * Note that None is also allowed, so always use this in a type hint with Optional. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskInstanceState". + */ +export type TaskInstanceState = + | "removed" + | "scheduled" + | "queued" + | "running" + | "success" + | "restarting" + | "failed" + | "up_for_retry" + | "up_for_reschedule" + | "upstream_failed" + | "skipped" + | "deferred" + | "awaiting_input"; +export type TaskRescheduleCount = number; +export type MaxTries = number; +export type Key = string; +export type Value = string | null; +export type Variables = VariableResponse[]; +export type ConnId1 = string; +export type ConnType1 = string; +export type Host1 = string | null; +export type Schema1 = string | null; +export type Login1 = string | null; +export type Password1 = string | null; +export type Port1 = number | null; +export type Extra6 = string | null; +export type Connections = ConnectionResponse[]; +export type NextMethod1 = string | null; +export type NextKwargs1 = + | { + [k: string]: unknown; + } + | string + | null; +export type XcomKeysToClear = string[]; +export type ShouldRetry = boolean; +export type StartDate2 = string | null; +export type Type13 = "TaskCallbackRequest"; +export type Filepath2 = string; +export type BundleName3 = string; +export type BundleVersion2 = string | null; +export type Msg2 = string | null; +export type EmailType = "failure" | "retry"; +export type Type14 = "EmailRequest"; +export type CallbackRequests = (DagCallbackRequest | TaskCallbackRequest | EmailRequest)[]; +export type Type15 = "DagFileParseRequest"; +export type Fileloc = string; +export type LastLoaded = string | null; +export type SerializedDags = LazyDeserializedDAG[]; +export type Warnings = unknown[] | null; +export type ImportErrors = { + [k: string]: string; +} | null; +export type Type16 = "DagFileParsingResult"; +export type DagId4 = string; +export type IsPaused = boolean; +export type BundleName4 = string | null; +export type BundleVersion3 = string | null; +export type RelativeFileloc = string | null; +export type Owners = string | null; +export type Tags = string[]; +export type NextDagrun = string | null; +export type Type17 = "DagResult"; +export type DagId5 = string; +export type RunId4 = string; +export type LogicalDate2 = string | null; +export type DataIntervalStart2 = string | null; +export type DataIntervalEnd2 = string | null; +export type RunAfter1 = string; +export type StartDate3 = string | null; +export type EndDate2 = string | null; +export type ClearNumber1 = number | null; +export type Conf1 = { + [k: string]: unknown; +} | null; +export type TriggeringUserName1 = string | null; +export type ConsumedAssetEvents1 = AssetEventDagRunReference[]; +export type PartitionKey4 = string | null; +export type PartitionDate1 = string | null; +export type Note1 = string | null; +export type TeamName1 = string | null; +export type Type18 = "DagRunResult"; +export type Type19 = "DagRunStateResult"; +export type State2 = "deferred" | null; +export type Classpath = string; +export type TriggerKwargs = + | { + [k: string]: JsonValue; + } + | string + | null; +export type TriggerTimeout = string | null; +export type Queue1 = string | null; +export type NextMethod2 = string; +export type NextKwargs2 = { + [k: string]: JsonValue; +} | null; +export type RenderedMapIndex1 = string | null; +export type Type20 = "DeferTask"; +export type Name8 = string; +export type Key1 = string; +export type Type21 = "DeleteAssetStateStoreByName"; +export type Uri5 = string; +export type Key2 = string; +export type Type22 = "DeleteAssetStateStoreByUri"; +export type TiId2 = string; +export type Key3 = string; +export type Type23 = "DeleteTaskStateStore"; +export type Key4 = string; +export type Type24 = "DeleteVariable"; +export type Key5 = string; +export type DagId6 = string; +export type RunId5 = string; +export type TaskId1 = string; +export type MapIndex1 = number | null; +export type Type25 = "DeleteXCom"; +/** + * Error types used in the API client. + */ +export type ErrorType = + | "CONNECTION_NOT_FOUND" + | "VARIABLE_NOT_FOUND" + | "XCOM_NOT_FOUND" + | "ASSET_NOT_FOUND" + | "TASK_STORE_NOT_FOUND" + | "ASSET_STORE_NOT_FOUND" + | "DAGRUN_ALREADY_EXISTS" + | "PERMISSION_DENIED" + | "GENERIC_ERROR" + | "API_SERVER_ERROR"; +export type Detail = { + [k: string]: unknown; +} | null; +export type Type26 = "ErrorResponse"; +/** + * Error types used in the API client. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ErrorType". + */ +export type ErrorType1 = + | "CONNECTION_NOT_FOUND" + | "VARIABLE_NOT_FOUND" + | "XCOM_NOT_FOUND" + | "ASSET_NOT_FOUND" + | "TASK_STORE_NOT_FOUND" + | "ASSET_STORE_NOT_FOUND" + | "DAGRUN_ALREADY_EXISTS" + | "PERMISSION_DENIED" + | "GENERIC_ERROR" + | "API_SERVER_ERROR"; +export type Name9 = string; +export type Type27 = "GetAssetByName"; +export type Uri6 = string; +export type Type28 = "GetAssetByUri"; +export type Name10 = string | null; +export type Uri7 = string | null; +export type After = string | null; +export type Before = string | null; +export type Limit = number | null; +export type Ascending = boolean; +export type Type29 = "GetAssetEventByAsset"; +export type AliasName = string; +export type After1 = string | null; +export type Before1 = string | null; +export type Limit1 = number | null; +export type Ascending1 = boolean; +export type Type30 = "GetAssetEventByAssetAlias"; +export type Name11 = string; +export type Key6 = string; +export type Type31 = "GetAssetStateStoreByName"; +export type Uri8 = string; +export type Key7 = string; +export type Type32 = "GetAssetStateStoreByUri"; +export type AliasName1 = string; +export type Type33 = "GetAssetsByAlias"; +export type ConnId2 = string; +export type Type34 = "GetConnection"; +export type DagId7 = string; +export type LogicalDates = string[] | null; +export type RunIds = string[] | null; +export type States = string[] | null; +export type Type35 = "GetDRCount"; +export type DagId8 = string; +export type Type36 = "GetDag"; +export type DagId9 = string; +export type RunId6 = string; +export type Type37 = "GetDagRun"; +export type DagId10 = string; +export type RunId7 = string; +export type Type38 = "GetDagRunState"; +export type TiId3 = string; +export type Type39 = "GetHITLDetailResponse"; +export type TiId4 = string; +export type Type40 = "GetPrevSuccessfulDagRun"; +export type DagId11 = string; +export type LogicalDate3 = string; +export type State3 = string | null; +export type Type41 = "GetPreviousDagRun"; +export type DagId12 = string; +export type TaskId2 = string; +export type LogicalDate4 = string | null; +export type MapIndex2 = number; +export type Type42 = "GetPreviousTI"; +export type DagId13 = string; +export type MapIndex3 = number | null; +export type TaskIds = string[] | null; +export type TaskGroupId = string | null; +export type LogicalDates1 = string[] | null; +export type RunIds1 = string[] | null; +export type States1 = string[] | null; +export type Type43 = "GetTICount"; +export type DagId14 = string; +export type RunId8 = string; +export type Type44 = "GetTaskBreadcrumbs"; +export type TiId5 = string; +export type TryNumber1 = number; +export type Type45 = "GetTaskRescheduleStartDate"; +export type TiId6 = string; +export type Key8 = string; +export type Type46 = "GetTaskStateStore"; +export type DagId15 = string; +export type MapIndex4 = number | null; +export type TaskIds1 = string[] | null; +export type TaskGroupId1 = string | null; +export type LogicalDates2 = string[] | null; +export type RunIds2 = string[] | null; +export type Type47 = "GetTaskStates"; +export type Key9 = string; +export type Type48 = "GetVariable"; +export type Prefix = string | null; +export type Limit2 = number; +export type Offset = number; +export type Type49 = "GetVariableKeys"; +export type Key10 = string; +export type DagId16 = string; +export type RunId9 = string; +export type TaskId3 = string; +export type MapIndex5 = number | null; +export type IncludePriorDates = boolean; +export type Type50 = "GetXCom"; +export type Key11 = string; +export type DagId17 = string; +export type RunId10 = string; +export type TaskId4 = string; +export type Type51 = "GetXComCount"; +export type Key12 = string; +export type DagId18 = string; +export type RunId11 = string; +export type TaskId5 = string; +export type Offset1 = number; +export type Type52 = "GetXComSequenceItem"; +export type Key13 = string; +export type DagId19 = string; +export type RunId12 = string; +export type TaskId6 = string; +export type Start = number | null; +export type Stop = number | null; +export type Step = number | null; +export type IncludePriorDates1 = boolean; +export type Type53 = "GetXComSequenceSlice"; +export type TiId7 = string; +/** + * @minItems 1 + */ +export type Options1 = [string, ...string[]]; +export type Subject1 = string; +export type Body1 = string | null; +export type Defaults1 = string[] | null; +export type Multiple1 = boolean | null; +export type Params1 = { + [k: string]: unknown; +} | null; +export type AssignedUsers1 = HITLUser[] | null; +export type Type54 = "HITLDetailRequestResult"; +export type InactiveAssets = AssetProfile[] | null; +export type Type55 = "InactiveAssetsResult"; +export type Name12 = string | null; +export type Type56 = "MaskSecret"; +export type Ok = boolean; +export type Type57 = "OKResponse"; +export type DataIntervalStart3 = string | null; +export type DataIntervalEnd3 = string | null; +export type StartDate4 = string | null; +export type EndDate3 = string | null; +export type Type58 = "PrevSuccessfulDagRunResult"; +export type Type59 = "PreviousDagRunResult"; +export type TaskId7 = string; +export type DagId20 = string; +export type RunId13 = string; +export type LogicalDate5 = string | null; +export type StartDate5 = string | null; +export type EndDate4 = string | null; +export type State4 = string | null; +export type TryNumber2 = number; +export type MapIndex6 = number | null; +export type Duration = number | null; +export type Type60 = "PreviousTIResult"; +export type Key14 = string; +export type Value1 = string | null; +export type Description = string | null; +export type Type61 = "PutVariable"; +export type State5 = "up_for_reschedule" | null; +export type RescheduleDate = string; +export type EndDate5 = string; +export type Type62 = "RescheduleTask"; +export type Type63 = "ResendLoggingFD"; +export type State6 = "up_for_retry" | null; +export type EndDate6 = string; +export type RenderedMapIndex2 = string | null; +export type RetryDelaySeconds = number | null; +export type RetryReason = string | null; +export type Type64 = "RetryTask"; +export type Type65 = "SentFDs"; +export type Fds = number[]; +export type Name13 = string; +export type Key15 = string; +export type Type66 = "SetAssetStateStoreByName"; +export type Uri9 = string; +export type Key16 = string; +export type Type67 = "SetAssetStateStoreByUri"; +export type Type68 = "SetRenderedFields"; +export type RenderedMapIndex3 = string; +export type Type69 = "SetRenderedMapIndex"; +export type TiId8 = string; +export type Key17 = string; +export type ExpiresAt = string | null; +export type Type70 = "SetTaskStateStore"; +export type Key18 = string; +export type DagId21 = string; +export type RunId14 = string; +export type TaskId8 = string; +export type MapIndex7 = number | null; +export type DagResult1 = boolean; +export type MappedLength = number | null; +export type Type71 = "SetXCom"; +export type Tasks = (string | [unknown, unknown])[]; +export type Type72 = "SkipDownstreamTasks"; +export type DagRelPath = string; +export type StartDate6 = string; +export type SentryIntegration = string; +export type Type73 = "StartupDetails"; +export type State7 = "success" | null; +export type EndDate7 = string; +export type TaskOutlets = AssetProfile[] | null; +export type OutletEvents = + | { + [k: string]: unknown; + }[] + | null; +export type RenderedMapIndex4 = string | null; +export type Type74 = "SucceedTask"; +export type Count1 = number; +export type Type75 = "TICount"; +export type Breadcrumbs = { + [k: string]: unknown; +}[]; +export type Type76 = "TaskBreadcrumbsResult"; +export type StartDate7 = string | null; +export type Type77 = "TaskRescheduleStartDate"; +export type State8 = "failed" | "skipped" | "removed"; +export type EndDate8 = string | null; +export type Type78 = "TaskState"; +export type RenderedMapIndex5 = string | null; +export type Type79 = "TaskStateStoreResult"; +export type Type80 = "TaskStatesResult"; +export type LogicalDate6 = string | null; +export type RunAfter2 = string | null; +export type Conf2 = { + [k: string]: unknown; +} | null; +export type ResetDagRun = boolean | null; +export type PartitionKey5 = string | null; +export type Note2 = string | null; +export type DagId22 = string; +export type DagRunId = string; +export type Type81 = "TriggerDagRun"; +export type TiId9 = string; +/** + * @minItems 1 + */ +export type ChosenOptions = [string, ...string[]]; +export type ParamsInput = { + [k: string]: unknown; +} | null; +export type Type82 = "UpdateHITLDetail"; +export type TiId10 = string; +export type Type83 = "ValidateInletsAndOutlets"; +export type Keys = string[]; +export type TotalEntries = number; +export type Type84 = "VariableKeysResult"; +export type Key19 = string; +export type Value2 = string | null; +export type Type85 = "VariableResult"; +export type Len = number; +export type Type86 = "XComCountResponse"; +export type Key20 = string; +export type Type87 = "XComResult"; +export type Type88 = "XComSequenceIndexResult"; +export type Root = JsonValue[]; +export type Type89 = "XComSequenceSliceResult"; + +export interface SupervisorWireSchema {} +/** + * Schema for AssetAliasModel used in AssetEventDagRunReference. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetAliasReferenceAssetEventDagRun". + */ +export interface AssetAliasReferenceAssetEventDagRun { + name: Name; +} +/** + * Asset event schema with fields that are needed for Runtime. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetEventResponse". + */ +export interface AssetEventResponse { + id: Id; + timestamp: Timestamp; + extra?: Extra; + asset: AssetResponse; + created_dagruns: CreatedDagruns; + source_task_id?: SourceTaskId; + source_dag_id?: SourceDagId; + source_run_id?: SourceRunId; + source_map_index?: SourceMapIndex; + partition_key?: PartitionKey1; +} +/** + * Asset schema for responses with fields that are needed for Runtime. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetResponse". + */ +export interface AssetResponse { + name: Name1; + uri: Uri; + group: Group; + extra?: Extra1; +} +/** + * DagRun serializer for asset responses. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DagRunAssetReference". + */ +export interface DagRunAssetReference { + run_id: RunId; + dag_id: DagId; + logical_date?: LogicalDate; + start_date: StartDate; + end_date?: EndDate; + state: State; + data_interval_start?: DataIntervalStart; + data_interval_end?: DataIntervalEnd; + partition_key?: PartitionKey; +} +/** + * Response to GetAssetEvent request. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetEventsResult". + */ +export interface AssetEventsResult { + asset_events: AssetEvents; + type?: Type; +} +/** + * Profile of an asset-like object. + * + * Asset will have name, uri defined, with type set to 'Asset'. + * AssetNameRef will have name defined, type set to 'AssetNameRef'. + * AssetUriRef will have uri defined, type set to 'AssetUriRef'. + * AssetAlias will have name defined, type set to 'AssetAlias'. + * + * Note that 'type' here is distinct from 'asset_type' the user declares on an + * Asset (or subclass). This field is for distinguishing between different + * asset-related types (Asset, AssetRef, or AssetAlias). + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetProfile". + */ +export interface AssetProfile { + name?: Name2; + uri?: Uri1; + type: Type1; +} +/** + * Schema for AssetModel used in AssetEventDagRunReference. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetReferenceAssetEventDagRun". + */ +export interface AssetReferenceAssetEventDagRun { + name: Name3; + uri: Uri2; + extra: Extra2; +} +export interface Extra2 { + [k: string]: JsonValue; +} +/** + * Response to ReadXCom request. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetResult". + */ +export interface AssetResult { + name: Name4; + uri: Uri3; + group: Group1; + extra?: Extra3; + type?: Type2; +} +/** + * Response to GetAssetStateStore; wraps the generated API response for supervisor to worker comms. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetStateStoreResult". + */ +export interface AssetStateStoreResult { + value: JsonValue; + type?: Type3; +} +/** + * Response to GetAssetsByAlias; list of concrete assets resolved from an alias. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetsByAliasResult". + */ +export interface AssetsByAliasResult { + assets: Assets; + type?: Type4; +} +/** + * Park a task instance awaiting human input (Human-in-the-loop), without a trigger. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AwaitInputTask". + */ +export interface AwaitInputTask { + state?: State1; + timeout?: Timeout; + next_method: NextMethod; + next_kwargs?: NextKwargs; + rendered_map_index?: RenderedMapIndex; + type?: Type5; +} +/** + * Schema for telling task which bundle to run with. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "BundleInfo". + */ +export interface BundleInfo { + name: Name5; + version?: Version; + version_data?: VersionData; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ClearAssetStateStoreByName". + */ +export interface ClearAssetStateStoreByName { + name: Name6; + type?: Type6; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ClearAssetStateStoreByUri". + */ +export interface ClearAssetStateStoreByUri { + uri: Uri4; + type?: Type7; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ClearTaskStateStore". + */ +export interface ClearTaskStateStore { + ti_id: TiId; + type?: Type8; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ConnectionResult". + */ +export interface ConnectionResult { + conn_id: ConnId; + conn_type: ConnType; + host?: Host; + schema?: Schema; + login?: Login; + password?: Password; + port?: Port; + extra?: Extra4; + type?: Type9; +} +/** + * Add the input request part of a Human-in-the-loop response. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "CreateHITLDetailPayload". + */ +export interface CreateHITLDetailPayload { + ti_id: TiId1; + options: Options; + subject: Subject; + body?: Body; + defaults?: Defaults; + multiple?: Multiple; + params?: Params; + assigned_users?: AssignedUsers; + type?: Type10; +} +/** + * Schema for a Human-in-the-loop users. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "HITLUser". + */ +export interface HITLUser { + id: Id1; + name: Name7; +} +/** + * Response containing count of Dag Runs matching certain filters. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DRCount". + */ +export interface DRCount { + count: Count; + type?: Type11; +} +/** + * A Class with information about the success/failure DAG callback to be executed. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DagCallbackRequest". + */ +export interface DagCallbackRequest { + filepath: Filepath; + bundle_name: BundleName; + bundle_version: BundleVersion; + msg?: Msg; + dag_id: DagId1; + run_id: RunId1; + context_from_server?: DagRunContext | null; + is_failure_callback?: IsFailureCallback; + type?: Type12; +} +/** + * Class to pass context info from the server to build a Execution context object. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DagRunContext". + */ +export interface DagRunContext { + dag_run?: DagRun | null; + last_ti?: TaskInstance | null; +} +/** + * Schema for DagRun model with minimal required fields needed for Runtime. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DagRun". + */ +export interface DagRun { + dag_id: DagId2; + run_id: RunId2; + logical_date: LogicalDate1; + data_interval_start: DataIntervalStart1; + data_interval_end: DataIntervalEnd1; + run_after: RunAfter; + start_date: StartDate1; + end_date: EndDate1; + clear_number?: ClearNumber; + run_type: DagRunType; + state: DagRunState; + conf?: Conf; + triggering_user_name?: TriggeringUserName; + consumed_asset_events: ConsumedAssetEvents; + partition_key: PartitionKey3; + partition_date?: PartitionDate; + note?: Note; + team_name?: TeamName; +} +/** + * Schema for AssetEvent model used in DagRun. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "AssetEventDagRunReference". + */ +export interface AssetEventDagRunReference { + asset: AssetReferenceAssetEventDagRun; + extra: Extra5; + source_task_id: SourceTaskId1; + source_dag_id: SourceDagId1; + source_run_id: SourceRunId1; + source_map_index: SourceMapIndex1; + source_aliases: SourceAliases; + timestamp: Timestamp1; + partition_key?: PartitionKey2; +} +export interface Extra5 { + [k: string]: JsonValue; +} +/** + * Schema for TaskInstance model with minimal required fields needed for Runtime. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskInstance". + */ +export interface TaskInstance { + id: Id2; + task_id: TaskId; + dag_id: DagId3; + run_id: RunId3; + try_number: TryNumber; + dag_version_id: DagVersionId; + map_index?: MapIndex; + hostname?: Hostname; + context_carrier?: ContextCarrier; + queue?: Queue; +} +/** + * Request for DAG File Parsing. + * + * This is the request that the manager will send to the DAG parser with the dag file and + * any other necessary metadata. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DagFileParseRequest". + */ +export interface DagFileParseRequest { + file: File; + bundle_path: BundlePath; + bundle_name: BundleName1; + callback_requests?: CallbackRequests; + type?: Type15; +} +/** + * Task callback status information. + * + * A Class with information about the success/failure TI callback to be executed. Currently, only failure + * callbacks when tasks are externally killed or experience heartbeat timeouts are run via DagFileProcessorProcess. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskCallbackRequest". + */ +export interface TaskCallbackRequest { + filepath: Filepath1; + bundle_name: BundleName2; + bundle_version: BundleVersion1; + msg?: Msg1; + ti: TaskInstance; + task_callback_type?: TaskInstanceState | null; + context_from_server?: TIRunContext | null; + type?: Type13; +} +/** + * Response schema for TaskInstance run context. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TIRunContext". + */ +export interface TIRunContext { + dag_run: DagRun; + task_reschedule_count?: TaskRescheduleCount; + max_tries: MaxTries; + variables?: Variables; + connections?: Connections; + next_method?: NextMethod1; + next_kwargs?: NextKwargs1; + xcom_keys_to_clear?: XcomKeysToClear; + should_retry?: ShouldRetry; + start_date?: StartDate2; +} +/** + * Variable schema for responses with fields that are needed for Runtime. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "VariableResponse". + */ +export interface VariableResponse { + key: Key; + value: Value; +} +/** + * Connection schema for responses with fields that are needed for Runtime. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ConnectionResponse". + */ +export interface ConnectionResponse { + conn_id: ConnId1; + conn_type: ConnType1; + host: Host1; + schema: Schema1; + login: Login1; + password: Password1; + port: Port1; + extra: Extra6; +} +/** + * Email notification request for task failures/retries. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "EmailRequest". + */ +export interface EmailRequest { + filepath: Filepath2; + bundle_name: BundleName3; + bundle_version: BundleVersion2; + msg?: Msg2; + ti: TaskInstance; + email_type?: EmailType; + context_from_server: TIRunContext; + type?: Type14; +} +/** + * Result of DAG File Parsing. + * + * This is the result of a successful DAG parse, in this class, we gather all serialized DAGs, + * import errors and warnings to send back to the scheduler to store in the DB. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DagFileParsingResult". + */ +export interface DagFileParsingResult { + fileloc: Fileloc; + serialized_dags: SerializedDags; + warnings?: Warnings; + import_errors?: ImportErrors; + type?: Type16; +} +/** + * Lazily build information from the serialized DAG structure. + * + * An object that will present "enough" of the DAG like interface to update DAG db models etc, without having + * to deserialize the full DAG and Task hierarchy. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "LazyDeserializedDAG". + */ +export interface LazyDeserializedDAG { + data: Data; + last_loaded?: LastLoaded; +} +export interface Data { + [k: string]: unknown; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DagResult". + */ +export interface DagResult { + dag_id: DagId4; + is_paused: IsPaused; + bundle_name?: BundleName4; + bundle_version?: BundleVersion3; + relative_fileloc?: RelativeFileloc; + owners?: Owners; + tags: Tags; + next_dagrun?: NextDagrun; + type?: Type17; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DagRunResult". + */ +export interface DagRunResult { + dag_id: DagId5; + run_id: RunId4; + logical_date?: LogicalDate2; + data_interval_start?: DataIntervalStart2; + data_interval_end?: DataIntervalEnd2; + run_after: RunAfter1; + start_date?: StartDate3; + end_date?: EndDate2; + clear_number?: ClearNumber1; + run_type: DagRunType; + state: DagRunState; + conf?: Conf1; + triggering_user_name?: TriggeringUserName1; + consumed_asset_events: ConsumedAssetEvents1; + partition_key?: PartitionKey4; + partition_date?: PartitionDate1; + note?: Note1; + team_name?: TeamName1; + type?: Type18; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DagRunStateResult". + */ +export interface DagRunStateResult { + state: DagRunState; + type?: Type19; +} +/** + * Update a task instance state to deferred. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DeferTask". + */ +export interface DeferTask { + state?: State2; + classpath: Classpath; + trigger_kwargs?: TriggerKwargs; + trigger_timeout?: TriggerTimeout; + queue?: Queue1; + next_method: NextMethod2; + next_kwargs?: NextKwargs2; + rendered_map_index?: RenderedMapIndex1; + type?: Type20; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DeleteAssetStateStoreByName". + */ +export interface DeleteAssetStateStoreByName { + name: Name8; + key: Key1; + type?: Type21; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DeleteAssetStateStoreByUri". + */ +export interface DeleteAssetStateStoreByUri { + uri: Uri5; + key: Key2; + type?: Type22; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DeleteTaskStateStore". + */ +export interface DeleteTaskStateStore { + ti_id: TiId2; + key: Key3; + type?: Type23; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DeleteVariable". + */ +export interface DeleteVariable { + key: Key4; + type?: Type24; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DeleteXCom". + */ +export interface DeleteXCom { + key: Key5; + dag_id: DagId6; + run_id: RunId5; + task_id: TaskId1; + map_index?: MapIndex1; + type?: Type25; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ErrorResponse". + */ +export interface ErrorResponse { + error?: ErrorType; + detail?: Detail; + type?: Type26; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetAssetByName". + */ +export interface GetAssetByName { + name: Name9; + type?: Type27; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetAssetByUri". + */ +export interface GetAssetByUri { + uri: Uri6; + type?: Type28; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetAssetEventByAsset". + */ +export interface GetAssetEventByAsset { + name: Name10; + uri: Uri7; + after?: After; + before?: Before; + limit?: Limit; + ascending?: Ascending; + type?: Type29; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetAssetEventByAssetAlias". + */ +export interface GetAssetEventByAssetAlias { + alias_name: AliasName; + after?: After1; + before?: Before1; + limit?: Limit1; + ascending?: Ascending1; + type?: Type30; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetAssetStateStoreByName". + */ +export interface GetAssetStateStoreByName { + name: Name11; + key: Key6; + type?: Type31; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetAssetStateStoreByUri". + */ +export interface GetAssetStateStoreByUri { + uri: Uri8; + key: Key7; + type?: Type32; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetAssetsByAlias". + */ +export interface GetAssetsByAlias { + alias_name: AliasName1; + type?: Type33; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetConnection". + */ +export interface GetConnection { + conn_id: ConnId2; + type?: Type34; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetDRCount". + */ +export interface GetDRCount { + dag_id: DagId7; + logical_dates?: LogicalDates; + run_ids?: RunIds; + states?: States; + type?: Type35; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetDag". + */ +export interface GetDag { + dag_id: DagId8; + type?: Type36; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetDagRun". + */ +export interface GetDagRun { + dag_id: DagId9; + run_id: RunId6; + type?: Type37; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetDagRunState". + */ +export interface GetDagRunState { + dag_id: DagId10; + run_id: RunId7; + type?: Type38; +} +/** + * Get the response content part of a Human-in-the-loop response. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetHITLDetailResponse". + */ +export interface GetHITLDetailResponse { + ti_id: TiId3; + type?: Type39; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetPrevSuccessfulDagRun". + */ +export interface GetPrevSuccessfulDagRun { + ti_id: TiId4; + type?: Type40; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetPreviousDagRun". + */ +export interface GetPreviousDagRun { + dag_id: DagId11; + logical_date: LogicalDate3; + state?: State3; + type?: Type41; +} +/** + * Request to get previous task instance. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetPreviousTI". + */ +export interface GetPreviousTI { + dag_id: DagId12; + task_id: TaskId2; + logical_date?: LogicalDate4; + map_index?: MapIndex2; + state?: TaskInstanceState | null; + type?: Type42; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetTICount". + */ +export interface GetTICount { + dag_id: DagId13; + map_index?: MapIndex3; + task_ids?: TaskIds; + task_group_id?: TaskGroupId; + logical_dates?: LogicalDates1; + run_ids?: RunIds1; + states?: States1; + type?: Type43; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetTaskBreadcrumbs". + */ +export interface GetTaskBreadcrumbs { + dag_id: DagId14; + run_id: RunId8; + type?: Type44; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetTaskRescheduleStartDate". + */ +export interface GetTaskRescheduleStartDate { + ti_id: TiId5; + try_number?: TryNumber1; + type?: Type45; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetTaskStateStore". + */ +export interface GetTaskStateStore { + ti_id: TiId6; + key: Key8; + type?: Type46; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetTaskStates". + */ +export interface GetTaskStates { + dag_id: DagId15; + map_index?: MapIndex4; + task_ids?: TaskIds1; + task_group_id?: TaskGroupId1; + logical_dates?: LogicalDates2; + run_ids?: RunIds2; + type?: Type47; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetVariable". + */ +export interface GetVariable { + key: Key9; + type?: Type48; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetVariableKeys". + */ +export interface GetVariableKeys { + prefix?: Prefix; + limit?: Limit2; + offset?: Offset; + type?: Type49; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetXCom". + */ +export interface GetXCom { + key: Key10; + dag_id: DagId16; + run_id: RunId9; + task_id: TaskId3; + map_index?: MapIndex5; + include_prior_dates?: IncludePriorDates; + type?: Type50; +} +/** + * Get the number of (mapped) XCom values available. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetXComCount". + */ +export interface GetXComCount { + key: Key11; + dag_id: DagId17; + run_id: RunId10; + task_id: TaskId4; + type?: Type51; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetXComSequenceItem". + */ +export interface GetXComSequenceItem { + key: Key12; + dag_id: DagId18; + run_id: RunId11; + task_id: TaskId5; + offset: Offset1; + type?: Type52; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "GetXComSequenceSlice". + */ +export interface GetXComSequenceSlice { + key: Key13; + dag_id: DagId19; + run_id: RunId12; + task_id: TaskId6; + start: Start; + stop: Stop; + step: Step; + include_prior_dates?: IncludePriorDates1; + type?: Type53; +} +/** + * Response to CreateHITLDetailPayload request. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "HITLDetailRequestResult". + */ +export interface HITLDetailRequestResult { + ti_id: TiId7; + options: Options1; + subject: Subject1; + body?: Body1; + defaults?: Defaults1; + multiple?: Multiple1; + params?: Params1; + assigned_users?: AssignedUsers1; + type?: Type54; +} +/** + * Response of InactiveAssets requests. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "InactiveAssetsResult". + */ +export interface InactiveAssetsResult { + inactive_assets?: InactiveAssets; + type?: Type55; +} +/** + * Add a new value to be redacted in task logs. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "MaskSecret". + */ +export interface MaskSecret { + value: JsonValue; + name?: Name12; + type?: Type56; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "OKResponse". + */ +export interface OKResponse { + ok: Ok; + type?: Type57; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "PrevSuccessfulDagRunResult". + */ +export interface PrevSuccessfulDagRunResult { + data_interval_start?: DataIntervalStart3; + data_interval_end?: DataIntervalEnd3; + start_date?: StartDate4; + end_date?: EndDate3; + type?: Type58; +} +/** + * Response containing previous Dag run information. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "PreviousDagRunResult". + */ +export interface PreviousDagRunResult { + dag_run?: DagRun | null; + type?: Type59; +} +/** + * Schema for response with previous TaskInstance information. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "PreviousTIResponse". + */ +export interface PreviousTIResponse { + task_id: TaskId7; + dag_id: DagId20; + run_id: RunId13; + logical_date?: LogicalDate5; + start_date?: StartDate5; + end_date?: EndDate4; + state?: State4; + try_number: TryNumber2; + map_index?: MapIndex6; + duration?: Duration; +} +/** + * Response containing previous task instance data. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "PreviousTIResult". + */ +export interface PreviousTIResult { + task_instance?: PreviousTIResponse | null; + type?: Type60; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "PutVariable". + */ +export interface PutVariable { + key: Key14; + value: Value1; + description: Description; + type?: Type61; +} +/** + * Update a task instance state to reschedule/up_for_reschedule. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "RescheduleTask". + */ +export interface RescheduleTask { + state?: State5; + reschedule_date: RescheduleDate; + end_date: EndDate5; + type?: Type62; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ResendLoggingFD". + */ +export interface ResendLoggingFD { + type?: Type63; +} +/** + * Update a task instance state to up_for_retry. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "RetryTask". + */ +export interface RetryTask { + state?: State6; + end_date: EndDate6; + rendered_map_index?: RenderedMapIndex2; + retry_delay_seconds?: RetryDelaySeconds; + retry_reason?: RetryReason; + type?: Type64; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "SentFDs". + */ +export interface SentFDs { + type?: Type65; + fds: Fds; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "SetAssetStateStoreByName". + */ +export interface SetAssetStateStoreByName { + name: Name13; + key: Key15; + value: JsonValue; + type?: Type66; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "SetAssetStateStoreByUri". + */ +export interface SetAssetStateStoreByUri { + uri: Uri9; + key: Key16; + value: JsonValue; + type?: Type67; +} +/** + * Payload for setting RTIF for a task instance. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "SetRenderedFields". + */ +export interface SetRenderedFields { + rendered_fields: RenderedFields; + type?: Type68; +} +export interface RenderedFields { + [k: string]: JsonValue; +} +/** + * Payload for setting rendered_map_index for a task instance. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "SetRenderedMapIndex". + */ +export interface SetRenderedMapIndex { + rendered_map_index: RenderedMapIndex3; + type?: Type69; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "SetTaskStateStore". + */ +export interface SetTaskStateStore { + ti_id: TiId8; + key: Key17; + value: JsonValue; + expires_at: ExpiresAt; + type?: Type70; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "SetXCom". + */ +export interface SetXCom { + key: Key18; + value: JsonValue; + dag_id: DagId21; + run_id: RunId14; + task_id: TaskId8; + map_index?: MapIndex7; + dag_result?: DagResult1; + mapped_length?: MappedLength; + type?: Type71; +} +/** + * Update state of downstream tasks within a task instance to 'skipped', while updating current task to success state. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "SkipDownstreamTasks". + */ +export interface SkipDownstreamTasks { + tasks: Tasks; + type?: Type72; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "StartupDetails". + */ +export interface StartupDetails { + ti: TaskInstance; + dag_rel_path: DagRelPath; + bundle_info: BundleInfo; + start_date: StartDate6; + ti_context: TIRunContext; + sentry_integration: SentryIntegration; + type?: Type73; +} +/** + * Update a task's state to success. Includes task_outlets and outlet_events for registering asset events. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "SucceedTask". + */ +export interface SucceedTask { + state?: State7; + end_date: EndDate7; + task_outlets?: TaskOutlets; + outlet_events?: OutletEvents; + rendered_map_index?: RenderedMapIndex4; + type?: Type74; +} +/** + * Response containing count of Task Instances matching certain filters. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TICount". + */ +export interface TICount { + count: Count1; + type?: Type75; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskBreadcrumbsResult". + */ +export interface TaskBreadcrumbsResult { + breadcrumbs: Breadcrumbs; + type?: Type76; +} +/** + * Response containing the first reschedule date for a task instance. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskRescheduleStartDate". + */ +export interface TaskRescheduleStartDate { + start_date: StartDate7; + type?: Type77; +} +/** + * Update a task's state. + * + * If a process exits without sending one of these the state will be derived from the exit code: + * - 0 = SUCCESS + * - anything else = FAILED + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskState". + */ +export interface TaskState { + state: State8; + end_date?: EndDate8; + type?: Type78; + rendered_map_index?: RenderedMapIndex5; +} +/** + * Response to GetTaskStateStore; wraps the generated API response for supervisor to worker comms. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskStateStoreResult". + */ +export interface TaskStateStoreResult { + value: JsonValue; + type?: Type79; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskStatesResult". + */ +export interface TaskStatesResult { + task_states: TaskStates; + type?: Type80; +} +export interface TaskStates { + [k: string]: unknown; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TriggerDagRun". + */ +export interface TriggerDagRun { + logical_date?: LogicalDate6; + run_after?: RunAfter2; + conf?: Conf2; + reset_dag_run?: ResetDagRun; + partition_key?: PartitionKey5; + note?: Note2; + dag_id: DagId22; + run_id: DagRunId; + type?: Type81; +} +/** + * Update the response content part of an existing Human-in-the-loop response. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "UpdateHITLDetail". + */ +export interface UpdateHITLDetail { + ti_id: TiId9; + chosen_options: ChosenOptions; + params_input?: ParamsInput; + type?: Type82; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ValidateInletsAndOutlets". + */ +export interface ValidateInletsAndOutlets { + ti_id: TiId10; + type?: Type83; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "VariableKeysResult". + */ +export interface VariableKeysResult { + keys: Keys; + total_entries: TotalEntries; + type?: Type84; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "VariableResult". + */ +export interface VariableResult { + key: Key19; + value?: Value2; + type?: Type85; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "XComCountResponse". + */ +export interface XComCountResponse { + len: Len; + type?: Type86; +} +/** + * Response to ReadXCom request. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "XComResult". + */ +export interface XComResult { + key: Key20; + value: JsonValue; + type?: Type87; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "XComSequenceIndexResult". + */ +export interface XComSequenceIndexResult { + root: JsonValue; + type?: Type88; +} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "XComSequenceSliceResult". + */ +export interface XComSequenceSliceResult { + root: Root; + type?: Type89; +} + +/** Cadwyn schema version this SDK was generated against. + * Not transmitted on the wire — the supervisor learns it out-of-band + * (e.g. bundle metadata) and runs the migrator accordingly. + * Exposed so the SDK author / operator can confirm which schema + * version their build is pinned to. */ +export const SUPERVISOR_API_VERSION = "2026-06-16" as const; diff --git a/ts-sdk/src/index.ts b/ts-sdk/src/index.ts index 30d79bfef853f..7abeb1af02cda 100644 --- a/ts-sdk/src/index.ts +++ b/ts-sdk/src/index.ts @@ -19,7 +19,9 @@ export { registerTask, listRegisteredTasks } from "./sdk/registry.js"; export { VariableNotFoundError } from "./sdk/client.js"; +export { startCoordinator, SUPERVISOR_API_VERSION } from "./coordinator/index.js"; export type { TaskClient } from "./sdk/client.js"; export type { ConnectionResult, GetXComOpts, JsonValue, SetXComOpts } from "./sdk/client-types.js"; +export type { StartCoordinatorOptions } from "./coordinator/index.js"; export type { TaskRegistration } from "./sdk/registry.js"; export type { TaskContext, TaskHandler, TaskHandlerArgs } from "./sdk/task.js"; diff --git a/ts-sdk/tests/coordinator/client.test.ts b/ts-sdk/tests/coordinator/client.test.ts new file mode 100644 index 0000000000000..0511ffd4179f0 --- /dev/null +++ b/ts-sdk/tests/coordinator/client.test.ts @@ -0,0 +1,204 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { describe, it, expect } from "vitest"; +import { createCoordinatorClient } from "../../src/coordinator/client.js"; +import type { CommChannel } from "../../src/coordinator/comm-channel.js"; +import type { TaskContext } from "../../src/sdk/task.js"; + +function fakeComm(frames: { body: unknown; error?: unknown }[]): CommChannel { + let i = 0; + return { + request: async () => frames[i++], + } as unknown as CommChannel; +} + +const FAKE_CTX: TaskContext = { + dagId: "d", + taskId: "t", + runId: "r", + tryNumber: 1, + mapIndex: -1, + signal: new AbortController().signal, +}; + +function client(frames: { body: unknown; error?: unknown }[]) { + return createCoordinatorClient(fakeComm(frames), FAKE_CTX); +} + +describe("getVariable not-found contract", () => { + it("returns null for the exact VARIABLE_NOT_FOUND code", async () => { + const c = client([{ body: { type: "ErrorResponse", error: "VARIABLE_NOT_FOUND" } }]); + expect(await c.getVariable("x")).toBeNull(); + }); + + it("throws for a non-not-found ErrorResponse", async () => { + const c = client([{ body: { type: "ErrorResponse", error: "API_SERVER_ERROR" } }]); + await expect(c.getVariable("x")).rejects.toThrow(/API_SERVER_ERROR/); + }); + + it("does NOT treat a value that merely contains 'NOT_FOUND' as absence", async () => { + const c = client([{ body: { type: "VariableResult", key: "x", value: "NOT_FOUND_LOL" } }]); + expect(await c.getVariable("x")).toBe("NOT_FOUND_LOL"); + }); + + it("does NOT treat an error code merely containing the substring as not-found", async () => { + // "SOMETHING_NOT_FOUND_ISH" is not in the exact set → must throw. + const c = client([{ body: { type: "ErrorResponse", error: "SOMETHING_NOT_FOUND_ISH" } }]); + await expect(c.getVariable("x")).rejects.toThrow(); + }); +}); + +describe("getVariableOrThrow", () => { + it("returns the value when present", async () => { + const c = client([{ body: { type: "VariableResult", key: "x", value: "v" } }]); + expect(await c.getVariableOrThrow("x")).toBe("v"); + }); + + it("throws VariableNotFoundError on missing key", async () => { + const c = client([{ body: { type: "ErrorResponse", error: "VARIABLE_NOT_FOUND" } }]); + await expect(c.getVariableOrThrow("x")).rejects.toThrow(/Variable not found: x/); + }); + + it("throws VariableNotFoundError on a null-valued result", async () => { + const c = client([{ body: { type: "VariableResult", key: "x", value: null } }]); + await expect(c.getVariableOrThrow("x")).rejects.toThrow(/Variable not found: x/); + }); +}); + +describe("getXCom not-found contract", () => { + it("returns null for the exact XCOM_NOT_FOUND code", async () => { + const c = client([{ body: { type: "ErrorResponse", error: "XCOM_NOT_FOUND" } }]); + expect(await c.getXCom({ key: "k" })).toBeNull(); + }); +}); + +describe("client is bound to TaskContext", () => { + it("defaults dag/task/run + map_index from ctx; allows override", async () => { + const sent: Record[] = []; + const recordingComm = { + request: async (b: Record) => { + sent.push(b); + return { body: null }; + }, + } as unknown as CommChannel; + const c = createCoordinatorClient(recordingComm, FAKE_CTX); + + await c.setXCom({ key: "echo", value: 1 }); + await c.setXCom({ + key: "echo", + value: 2, + dagId: "other", + taskId: "up", + runId: "rX", + }); + + expect(sent[0]).toMatchObject({ + type: "SetXCom", + key: "echo", + value: 1, + dag_id: "d", + task_id: "t", + run_id: "r", + map_index: null, + }); + expect(sent[1]).toMatchObject({ + dag_id: "other", + task_id: "up", + run_id: "rX", + }); + }); + + it("maps camelCase public XCom options to snake_case supervisor fields", async () => { + const sent: Record[] = []; + const recordingComm = { + request: async (b: Record) => { + sent.push(b); + // setXCom expects body=null; getXCom expects an XComResult. + return b.type === "GetXCom" + ? { body: { type: "XComResult", key: b.key, value: null } } + : { body: null }; + }, + } as unknown as CommChannel; + // ctx with a real map index — to prove -1 from opts wins over a + // mapped ctx value (caller is explicitly asking "the non-mapped row"). + const mappedCtx: TaskContext = { ...FAKE_CTX, mapIndex: 3 }; + const c = createCoordinatorClient(recordingComm, mappedCtx); + + await c.setXCom({ key: "k", value: 1, mapIndex: -1 }); + await c.setXCom({ key: "k", value: 2, mapIndex: null }); + await c.getXCom({ + key: "k", + dagId: "other_dag", + taskId: "upstream", + runId: "manual__1", + mapIndex: -1, + includePriorDates: true, + }); + await c.setXCom({ key: "k", value: 3, mapIndex: 5 }); + + expect(sent[0]).toMatchObject({ type: "SetXCom", map_index: null }); + expect(sent[1]).toMatchObject({ type: "SetXCom", map_index: null }); + expect(sent[2]).toMatchObject({ + type: "GetXCom", + dag_id: "other_dag", + task_id: "upstream", + run_id: "manual__1", + map_index: null, + include_prior_dates: true, + }); + expect(sent[3]).toMatchObject({ type: "SetXCom", map_index: 5 }); + }); +}); + +describe("getConnection", () => { + it("maps wire snake_case connection fields to public camelCase fields", async () => { + const c = client([ + { + body: { + type: "ConnectionResult", + conn_id: "warehouse", + conn_type: "postgres", + host: "db.local", + schema: "analytics", + login: "airflow", + password: "secret", + port: 5432, + extra: "{}", + }, + }, + ]); + + await expect(c.getConnection("warehouse")).resolves.toEqual({ + id: "warehouse", + type: "postgres", + host: "db.local", + schema: "analytics", + login: "airflow", + password: "secret", + port: 5432, + extra: "{}", + }); + }); + + it("returns null for missing connections", async () => { + const c = client([{ body: { type: "ErrorResponse", error: "CONNECTION_NOT_FOUND" } }]); + expect(await c.getConnection("missing")).toBeNull(); + }); +}); diff --git a/ts-sdk/tests/coordinator/comm-channel.test.ts b/ts-sdk/tests/coordinator/comm-channel.test.ts new file mode 100644 index 0000000000000..ae71e659b7ff1 --- /dev/null +++ b/ts-sdk/tests/coordinator/comm-channel.test.ts @@ -0,0 +1,166 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { EventEmitter } from "node:events"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { CommChannel, COORDINATOR_REQUEST_TIMEOUT_MS } from "../../src/coordinator/comm-channel.js"; +import { encodeResponse } from "../../src/coordinator/frames.js"; + +class FakeSocket extends EventEmitter { + writeCallback: ((err?: Error) => void) | undefined; + writeError: Error | undefined; + write = vi.fn((_buf: Buffer, cb?: (err?: Error) => void) => { + if (this.writeError) throw this.writeError; + this.writeCallback = cb; + return true; + }); + end = vi.fn((cb?: () => void) => cb?.()); + destroy = vi.fn(); +} + +function createChannel(sock: FakeSocket): CommChannel { + const ctor = CommChannel as unknown as new (sock: FakeSocket, logs: null) => CommChannel; + return new ctor(sock, null); +} + +describe("CommChannel", () => { + afterEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); + }); + + it("destroys the socket when a response write times out", async () => { + vi.useFakeTimers(); + const sock = new FakeSocket(); + const channel = createChannel(sock); + + const send = channel.sendResponse(7, { type: "TaskState", state: "failed" }, undefined, { + timeoutMs: 10, + }); + const assertion = expect(send).rejects.toThrow("Timed out sending response after 10 ms"); + await vi.advanceTimersByTimeAsync(10); + + await assertion; + expect(sock.destroy).toHaveBeenCalledWith(expect.any(Error)); + }); + + it("clears the response timeout when the write completes", async () => { + vi.useFakeTimers(); + const sock = new FakeSocket(); + const channel = createChannel(sock); + + const send = channel.sendResponse(7, { type: "TaskState", state: "failed" }, undefined, { + timeoutMs: 10, + }); + sock.writeCallback?.(); + + await expect(send).resolves.toBeUndefined(); + await vi.advanceTimersByTimeAsync(10); + expect(sock.destroy).not.toHaveBeenCalled(); + }); + + it("clears the response timeout when the write throws", async () => { + vi.useFakeTimers(); + const sock = new FakeSocket(); + const channel = createChannel(sock); + sock.writeError = new Error("write failed"); + + const send = channel.sendResponse(7, { type: "TaskState", state: "failed" }, undefined, { + timeoutMs: 10, + }); + + await expect(send).rejects.toThrow("write failed"); + await vi.advanceTimersByTimeAsync(10); + expect(sock.destroy).not.toHaveBeenCalled(); + }); + + it("rejects a request when its response times out", async () => { + vi.useFakeTimers(); + const sock = new FakeSocket(); + const channel = createChannel(sock); + + const request = channel.request({ type: "GetVariable", key: "greeting" }, { timeoutMs: 10 }); + const assertion = expect(request).rejects.toThrow( + "Timed out waiting for GetVariable response after 10 ms", + ); + await vi.advanceTimersByTimeAsync(10); + + await assertion; + }); + + it("uses the default request timeout when no override is provided", async () => { + vi.useFakeTimers(); + const sock = new FakeSocket(); + const channel = createChannel(sock); + + const request = channel.request({ type: "GetVariable", key: "greeting" }); + const rejection = vi.fn(); + request.catch(rejection); + const assertion = expect(request).rejects.toThrow( + `Timed out waiting for GetVariable response after ${COORDINATOR_REQUEST_TIMEOUT_MS} ms`, + ); + + await vi.advanceTimersByTimeAsync(COORDINATOR_REQUEST_TIMEOUT_MS - 1); + expect(rejection).not.toHaveBeenCalled(); + await vi.advanceTimersByTimeAsync(1); + + await assertion; + expect(rejection).toHaveBeenCalledOnce(); + }); + + it("clears the request timeout when the response arrives", async () => { + vi.useFakeTimers(); + const sock = new FakeSocket(); + const channel = createChannel(sock); + + const request = channel.request({ type: "GetVariable", key: "greeting" }, { timeoutMs: 10 }); + sock.emit("data", encodeResponse(0, { type: "VariableResult", value: "hello" })); + + await expect(request).resolves.toMatchObject({ + body: { type: "VariableResult", value: "hello" }, + }); + await vi.advanceTimersByTimeAsync(10); + }); + + it("clears the request timeout when the write fails", async () => { + vi.useFakeTimers(); + const sock = new FakeSocket(); + const channel = createChannel(sock); + + const request = channel.request({ type: "GetVariable", key: "greeting" }, { timeoutMs: 10 }); + sock.writeCallback?.(new Error("write failed")); + + await expect(request).rejects.toThrow("write failed"); + await vi.advanceTimersByTimeAsync(10); + }); + + it("does not log clean close events", async () => { + const logs = { debug: vi.fn(), warning: vi.fn() }; + const sock = new FakeSocket(); + const ctor = CommChannel as unknown as new (sock: FakeSocket, logs: unknown) => CommChannel; + const channel = new ctor(sock, logs); + + sock.emit("data", encodeResponse(0, { type: "StartupDetails" })); + await channel.close(); + sock.emit("close"); + + expect(logs.debug).not.toHaveBeenCalledWith("Comm channel closed", expect.anything()); + expect(logs.warning).not.toHaveBeenCalled(); + }); +}); diff --git a/ts-sdk/tests/coordinator/deferred.test.ts b/ts-sdk/tests/coordinator/deferred.test.ts new file mode 100644 index 0000000000000..a4f4454aa8d26 --- /dev/null +++ b/ts-sdk/tests/coordinator/deferred.test.ts @@ -0,0 +1,109 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { afterEach, describe, expect, it, vi } from "vitest"; +import { Deferred } from "../../src/coordinator/deferred.js"; + +describe("Deferred", () => { + afterEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); + }); + + it("settles only once", async () => { + const deferred = new Deferred(); + + deferred.resolve(1); + deferred.resolve(2); + deferred.reject(new Error("late")); + + await expect(deferred.promise).resolves.toBe(1); + expect(deferred.settled).toBe(true); + }); + + it("runs onSettle once after resolve", async () => { + const onSettle = vi.fn(); + const deferred = new Deferred().onSettle(onSettle); + + deferred.resolve(1); + deferred.resolve(2); + + await expect(deferred.promise).resolves.toBe(1); + expect(onSettle).toHaveBeenCalledTimes(1); + }); + + it("runs onSettle once after reject", async () => { + const onSettle = vi.fn(); + const deferred = new Deferred().onSettle(onSettle); + const assertion = expect(deferred.promise).rejects.toThrow("boom"); + + deferred.reject(new Error("boom")); + deferred.resolve(1); + + await assertion; + expect(onSettle).toHaveBeenCalledTimes(1); + }); + + it("runs onSettle immediately when already settled", async () => { + const onSettle = vi.fn(); + const deferred = new Deferred(); + + deferred.resolve(1); + deferred.onSettle(onSettle); + + await expect(deferred.promise).resolves.toBe(1); + expect(onSettle).toHaveBeenCalledTimes(1); + }); + + it("rejects after timeout", async () => { + vi.useFakeTimers(); + const deferred = new Deferred().rejectAfter(10, () => new Error("timeout")); + + const assertion = expect(deferred.promise).rejects.toThrow("timeout"); + await vi.advanceTimersByTimeAsync(10); + + await assertion; + expect(deferred.settled).toBe(true); + }); + + it("clears rejectAfter timer when resolved first", async () => { + vi.useFakeTimers(); + const makeError = vi.fn(() => new Error("timeout")); + const deferred = new Deferred().rejectAfter(10, makeError); + + deferred.resolve(1); + await vi.advanceTimersByTimeAsync(10); + + await expect(deferred.promise).resolves.toBe(1); + expect(makeError).not.toHaveBeenCalled(); + }); + + it("clears rejectAfter timer when rejected first", async () => { + vi.useFakeTimers(); + const makeError = vi.fn(() => new Error("timeout")); + const deferred = new Deferred().rejectAfter(10, makeError); + const assertion = expect(deferred.promise).rejects.toThrow("manual"); + + deferred.reject(new Error("manual")); + await vi.advanceTimersByTimeAsync(10); + + await assertion; + expect(makeError).not.toHaveBeenCalled(); + }); +}); diff --git a/ts-sdk/tests/coordinator/frames.test.ts b/ts-sdk/tests/coordinator/frames.test.ts new file mode 100644 index 0000000000000..4ccabfa79ccb2 --- /dev/null +++ b/ts-sdk/tests/coordinator/frames.test.ts @@ -0,0 +1,81 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { describe, expect, it } from "vitest"; +import { + decodePayload, + encodeRequest, + encodeResponse, + FrameReader, + tryTakeFrame, +} from "../../src/coordinator/frames.js"; + +function stripFramePrefix(buf: Buffer): Buffer { + return buf.subarray(4); +} + +describe("frames", () => { + it("round-trips a request frame", () => { + const encoded = encodeRequest(7, { type: "GetVariable", key: "foo" }); + const prefixedLen = encoded.readUInt32BE(0); + expect(encoded.length).toBe(4 + prefixedLen); + const decoded = decodePayload(stripFramePrefix(encoded)); + expect(decoded.id).toBe(7); + expect(decoded.body).toEqual({ type: "GetVariable", key: "foo" }); + expect(decoded.error).toBeUndefined(); + expect(decoded.isResponse).toBe(false); + }); + + it("round-trips a response frame with error", () => { + const encoded = encodeResponse(3, null, { error: "NotFound", detail: "x" }); + const decoded = decodePayload(stripFramePrefix(encoded)); + expect(decoded.id).toBe(3); + expect(decoded.body).toBeNull(); + expect(decoded.error).toEqual({ error: "NotFound", detail: "x" }); + expect(decoded.isResponse).toBe(true); + }); + + it("FrameReader stitches fragmented chunks", () => { + const full = Buffer.concat([ + encodeRequest(0, { type: "A", n: 1 }), + encodeRequest(1, { type: "B", n: 2 }), + encodeRequest(2, { type: "C", n: 3 }), + ]); + const reader = new FrameReader(); + const split1 = full.subarray(0, 3); + const firstLen = full.readUInt32BE(0); + const split2 = full.subarray(3, 4 + firstLen); + const split3 = full.subarray(4 + firstLen); + const frames = [...reader.push(split1), ...reader.push(split2), ...reader.push(split3)]; + expect(frames.map((f) => f.id)).toEqual([0, 1, 2]); + expect(reader.pending).toBe(0); + }); + + it("tryTakeFrame returns null when buffer is short", () => { + const encoded = encodeRequest(0, { type: "X" }); + expect(tryTakeFrame(encoded.subarray(0, 3))).toBeNull(); + expect(tryTakeFrame(encoded.subarray(0, 5))).toBeNull(); + expect(tryTakeFrame(encoded)).not.toBeNull(); + }); + + it("rejects non-array payloads", () => { + const bogus = Buffer.from([0x2a]); // msgpack fixint 42 + expect(() => decodePayload(bogus)).toThrow(/array frame/); + }); +}); diff --git a/ts-sdk/tests/coordinator/integration.test.ts b/ts-sdk/tests/coordinator/integration.test.ts new file mode 100644 index 0000000000000..42dc3650a4e57 --- /dev/null +++ b/ts-sdk/tests/coordinator/integration.test.ts @@ -0,0 +1,580 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// Pure-Node integration test for the coordinator runtime. +// +// Stands up a minimal in-process "supervisor" (TCP server + +// length-prefixed msgpack frame codec) that mirrors what Airflow's +// real `BaseCoordinator._runtime_subprocess_entrypoint` does, then +// drives the runtime through task success, task failure, retry, abort signaling, +// task-time RPCs, missing handlers, and parse-mode responses. +// +// No Python, no Airflow install — but exercises the same wire format +// the real coordinator speaks. + +import { afterEach, describe, expect, it, vi } from "vitest"; +import { createServer, Socket as NetSocket, type Server, type Socket } from "node:net"; +import { encode, decode } from "@msgpack/msgpack"; +import { + COORDINATOR_RESPONSE_TIMEOUT_MS, + startCoordinator, +} from "../../src/coordinator/runtime.js"; +import { registerTask } from "../../src/sdk/registry.js"; + +interface MockResult { + firstResponse: { id: number; body: unknown; isResponse: boolean } | null; + logRecords: Record[]; + runtimeRequests: { type: string; body: Record }[]; +} + +/** Callback used by `driveSupervisor` to answer runtime-initiated + * requests. Return `{ body, error? }` to reply with that arity-3 + * frame, or `null` to ignore the request (the runtime will hang + * waiting for a response — only useful for negative tests). */ +type Responder = ( + msgType: string, + body: Record, +) => { body: unknown; error?: unknown } | null; + +function frameBytes(id: number, body: unknown, isResponse: boolean): Buffer { + const arr = isResponse ? [id, body, null] : [id, body]; + const payload = Buffer.from(encode(arr)); + const header = Buffer.alloc(4); + header.writeUInt32BE(payload.length, 0); + return Buffer.concat([header, payload]); +} + +function readFrames(buf: Buffer): { + frames: { id: number; body: unknown; arity: number }[]; + rest: Buffer; +} { + const out: { id: number; body: unknown; arity: number }[] = []; + let rest: Buffer = buf; + while (rest.length >= 4) { + const len = rest.readUInt32BE(0); + if (rest.length < 4 + len) break; + const payload = rest.subarray(4, 4 + len); + const arr = decode(Buffer.from(payload)) as unknown[]; + out.push({ id: arr[0] as number, body: arr[1], arity: arr.length }); + rest = Buffer.from(rest.subarray(4 + len)); + } + return { frames: out, rest }; +} + +function makeStartupDetails( + taskId: string, + dagId = "test_dag", + runId = "r1", + tiContext: Record = {}, +): unknown { + return { + type: "StartupDetails", + ti: { + id: "ti-1", + dag_version_id: "dag-version-1", + task_id: taskId, + dag_id: dagId, + run_id: runId, + try_number: 1, + map_index: -1, + hostname: "test-host", + queue: "default", + }, + dag_rel_path: "test.py", + bundle_info: { name: "test", version: null }, + start_date: "2026-04-23T00:00:00Z", + ti_context: tiContext, + sentry_integration: "", + }; +} + +async function listen(): Promise<{ server: Server; port: number }> { + return new Promise((resolve) => { + const server = createServer(); + server.listen(0, "127.0.0.1", () => { + const port = (server.address() as { port: number }).port; + resolve({ server, port }); + }); + }); +} + +async function acceptOne(server: Server): Promise { + return new Promise((resolve) => server.once("connection", resolve)); +} + +function xcomKey(body: Record): string { + return [body["dag_id"], body["run_id"], body["task_id"], body["key"]].join("|"); +} + +function isFrameWithBodyType(chunk: unknown, type: string): boolean { + if (!Buffer.isBuffer(chunk)) return false; + try { + return readFrames(chunk).frames.some((f) => { + const body = f.body as Record | null; + return body?.["type"] === type; + }); + } catch { + return false; + } +} + +async function driveSupervisor(initialFrame: unknown, responder?: Responder): Promise { + const comm = await listen(); + const logs = await listen(); + + const commAccept = acceptOne(comm.server); + const logsAccept = acceptOne(logs.server); + + const runtimeDone = startCoordinator({ + commAddr: `127.0.0.1:${comm.port}`, + logsAddr: `127.0.0.1:${logs.port}`, + argv: [], + }); + + const [commSock, logsSock] = await Promise.all([commAccept, logsAccept]); + + // Send the kickoff frame as a _ResponseFrame (arity 3) — matches what + // Airflow's `_send_startup_details` actually emits on the wire. + commSock.write(frameBytes(0, initialFrame, true)); + + let firstResponse: MockResult["firstResponse"] = null; + const runtimeRequests: MockResult["runtimeRequests"] = []; + const logChunks: Buffer[] = []; + let commBuf: Buffer = Buffer.from(new Uint8Array(0)); + + commSock.on("data", (chunk: Buffer) => { + commBuf = Buffer.from(Buffer.concat([commBuf, chunk])); + const taken = readFrames(commBuf); + commBuf = taken.rest; + for (const f of taken.frames) { + if (f.arity === 2) { + // Runtime-initiated request (mid-task RPC). Forward to the + // responder; reply on the same id with an arity-3 frame. + const body = (f.body ?? {}) as Record; + const msgType = String(body["type"] ?? ""); + runtimeRequests.push({ type: msgType, body }); + const reply = responder?.(msgType, body) ?? + // Default: ack any request with an empty body so the + // runtime never hangs on auto-generated RPCs (e.g. the + // return_value XCom push). + { body: null }; + commSock.write(frameBytes(f.id, reply.body, true)); + continue; + } + // Arity-3: a response from the runtime. The terminal frame + // (SucceedTask / TaskState) lands here. + if (firstResponse === null) { + firstResponse = { id: f.id, body: f.body, isResponse: true }; + } + } + }); + logsSock.on("data", (chunk: Buffer) => logChunks.push(chunk)); + + // Wait for the runtime to finish AND for the comm socket to deliver + // its final bytes (FIN signals that all preceding frames are flushed). + const commSocketEnded = new Promise((resolve) => commSock.on("end", () => resolve())); + await Promise.all([runtimeDone, commSocketEnded]); + + comm.server.close(); + logs.server.close(); + + const lines = Buffer.concat(logChunks).toString("utf8").split("\n").filter(Boolean); + const logRecords = lines.map((l) => JSON.parse(l) as Record); + + return { firstResponse, logRecords, runtimeRequests }; +} + +describe("coordinator runtime integration", () => { + afterEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); + }); + + it("closes the log socket when the comm connection fails during startup", async () => { + const logs = await listen(); + const comm = await listen(); + + const logsSockPromise = acceptOne(logs.server); + const commSockPromise = acceptOne(comm.server); + const runtimeDone = startCoordinator({ + commAddr: `127.0.0.1:${comm.port}`, + logsAddr: `127.0.0.1:${logs.port}`, + argv: [], + }); + const [logsSock, commSock] = await Promise.all([logsSockPromise, commSockPromise]); + const logsSockEnded = new Promise((resolve) => logsSock.on("end", () => resolve())); + logsSock.resume(); + + commSock.write(Buffer.from([0, 0, 0, 1, 0x2a])); + await expect(runtimeDone).rejects.toThrow(/array frame/); + await logsSockEnded; + commSock.destroy(); + logs.server.close(); + comm.server.close(); + }); + + it("dispatches StartupDetails to a registered handler and emits SucceedTask", async () => { + let observedCtx: unknown = null; + registerTask({ dagId: "test_dag", taskId: "say_hello" }, async ({ ctx }) => { + observedCtx = ctx; + return "ok"; + }); + + const result = await driveSupervisor(makeStartupDetails("say_hello")); + + expect(result.firstResponse).not.toBeNull(); + expect(result.firstResponse!.body).toMatchObject({ + type: "SucceedTask", + task_outlets: [], + outlet_events: [], + }); + expect(observedCtx).toMatchObject({ + taskId: "say_hello", + dagId: "test_dag", + runId: "r1", + }); + expect(result.logRecords.some((r) => r["event"] === "[ts-sdk.runtime] Task succeeded")).toBe( + true, + ); + + // Logger names should be hierarchical (`ts-sdk.`) so the + // Python supervisor's ConsoleRenderer prints them as a distinct + // `[name]` column — not hardcoded to "task" (which collides with + // user task logs). + const loggers = new Set(result.logRecords.map((r) => r["logger"])); + expect(loggers.has("ts-sdk.runtime")).toBe(true); + for (const l of loggers) { + expect(l).toMatch(/^ts-sdk(\.|$)/); + } + }); + + it("times out when the terminal task response cannot be written", async () => { + vi.useFakeTimers(); + const originalWrite = NetSocket.prototype.write; + let stalledTerminalWrite: (() => void) | null = null; + + vi.spyOn(NetSocket.prototype, "write").mockImplementation(function ( + this: NetSocket, + ...args: Parameters + ): boolean { + if (isFrameWithBodyType(args[0], "SucceedTask")) { + stalledTerminalWrite?.(); + return true; + } + return Reflect.apply(originalWrite, this, args) as boolean; + }); + + const comm = await listen(); + const logs = await listen(); + const commAccept = acceptOne(comm.server); + const logsAccept = acceptOne(logs.server); + + registerTask({ dagId: "test_dag", taskId: "terminal_timeout" }, async () => undefined); + const runtimeDone = startCoordinator({ + commAddr: `127.0.0.1:${comm.port}`, + logsAddr: `127.0.0.1:${logs.port}`, + argv: [], + }); + const assertion = expect(runtimeDone).rejects.toThrow( + `Timed out sending response after ${COORDINATOR_RESPONSE_TIMEOUT_MS} ms`, + ); + + const [commSock, logsSock] = await Promise.all([commAccept, logsAccept]); + logsSock.resume(); + try { + const stalled = new Promise((resolve) => { + stalledTerminalWrite = resolve; + }); + commSock.write(frameBytes(0, makeStartupDetails("terminal_timeout"), true)); + await stalled; + await vi.advanceTimersByTimeAsync(COORDINATOR_RESPONSE_TIMEOUT_MS); + + await assertion; + } finally { + commSock.destroy(); + logsSock.destroy(); + comm.server.close(); + logs.server.close(); + } + }); + + it("returns TaskState=failed when the handler throws", async () => { + registerTask({ dagId: "test_dag", taskId: "boom" }, async () => { + throw new Error("boom"); + }); + + const result = await driveSupervisor(makeStartupDetails("boom")); + + expect(result.firstResponse!.body).toMatchObject({ + type: "TaskState", + state: "failed", + }); + }); + + it("returns RetryTask when the handler throws and Airflow says the failure is retryable", async () => { + registerTask({ dagId: "test_dag", taskId: "boom_retry" }, async () => { + throw new Error("boom"); + }); + + const result = await driveSupervisor( + makeStartupDetails("boom_retry", "test_dag", "r1", { + should_retry: true, + max_tries: 3, + }), + ); + + expect(result.firstResponse!.body).toMatchObject({ + type: "RetryTask", + retry_reason: "boom", + }); + }); + + it("aborts ctx.signal on SIGTERM and reports a thrown task error", async () => { + let sawAbort = false; + registerTask({ dagId: "test_dag", taskId: "aborted_then_failed" }, async ({ ctx }) => { + process.emit("SIGTERM"); + sawAbort = ctx.signal.aborted; + throw new Error("interrupted"); + }); + + const result = await driveSupervisor(makeStartupDetails("aborted_then_failed")); + + expect(sawAbort).toBe(true); + expect(result.firstResponse!.body).toMatchObject({ + type: "TaskState", + state: "failed", + }); + expect(result.runtimeRequests.filter((r) => r.type === "SetXCom")).toHaveLength(0); + }); + + it("returns RetryTask with the thrown error when a task fails after SIGTERM", async () => { + let sawAbort = false; + registerTask({ dagId: "test_dag", taskId: "aborted_then_failed_retry" }, async ({ ctx }) => { + process.emit("SIGTERM"); + sawAbort = ctx.signal.aborted; + throw new Error("interrupted"); + }); + + const result = await driveSupervisor( + makeStartupDetails("aborted_then_failed_retry", "test_dag", "r1", { + should_retry: true, + max_tries: 3, + }), + ); + + expect(sawAbort).toBe(true); + expect(result.firstResponse!.body).toMatchObject({ + type: "RetryTask", + retry_reason: "interrupted", + }); + expect(result.runtimeRequests.filter((r) => r.type === "SetXCom")).toHaveLength(0); + }); + + it("does not discard a completed task result after SIGTERM", async () => { + let sawAbort = false; + registerTask({ dagId: "test_dag", taskId: "completed_after_sigterm" }, async ({ ctx }) => { + process.emit("SIGTERM"); + sawAbort = ctx.signal.aborted; + return "completed"; + }); + + const result = await driveSupervisor(makeStartupDetails("completed_after_sigterm")); + + expect(sawAbort).toBe(true); + expect(result.firstResponse!.body).toMatchObject({ + type: "SucceedTask", + task_outlets: [], + outlet_events: [], + }); + + const setXComReqs = result.runtimeRequests.filter((r) => r.type === "SetXCom"); + expect(setXComReqs).toHaveLength(1); + expect(setXComReqs[0]!.body).toMatchObject({ + key: "return_value", + value: "completed", + }); + }); + + it("returns TaskState=removed when no handler is registered", async () => { + const result = await driveSupervisor(makeStartupDetails("missing_task")); + + expect(result.firstResponse!.body).toMatchObject({ + type: "TaskState", + state: "removed", + }); + }); + + it("exposes a client that round-trips GetVariable / SetXCom / GetXCom", async () => { + const xcomStore = new Map(); + let observedGreeting: string | null = ""; + + registerTask({ dagId: "test_dag", taskId: "say_hello_client" }, async ({ ctx, client }) => { + // The coordinator-mode handler MUST receive a client. + if (!client) throw new Error("client missing in coordinator mode"); + + observedGreeting = await client.getVariable("e6_greeting"); + await client.setXCom({ + key: "echo", + value: `node says: ${observedGreeting}`, + dagId: ctx.dagId, + taskId: ctx.taskId, + runId: ctx.runId, + }); + const back = await client.getXCom({ + key: "echo", + dagId: ctx.dagId, + taskId: ctx.taskId, + runId: ctx.runId, + }); + if (back !== `node says: ${observedGreeting}`) { + throw new Error(`xcom round-trip mismatch: ${back}`); + } + }); + + const responder: Responder = (msgType, body) => { + if (msgType === "GetVariable") { + return { + body: { + type: "VariableResult", + key: body["key"], + value: "hello from airflow", + }, + }; + } + if (msgType === "SetXCom") { + const k = xcomKey(body); + xcomStore.set(k, body["value"]); + // Supervisor sends an empty arity-3 frame on success. + return { body: null }; + } + if (msgType === "GetXCom") { + const k = xcomKey(body); + const value = xcomStore.get(k) ?? null; + return { + body: { type: "XComResult", key: body["key"], value }, + }; + } + return null; + }; + + const result = await driveSupervisor(makeStartupDetails("say_hello_client"), responder); + + expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" }); + expect(observedGreeting).toBe("hello from airflow"); + + const requestTypes = result.runtimeRequests.map((r) => r.type); + expect(requestTypes).toEqual(["GetVariable", "SetXCom", "GetXCom"]); + + const setReq = result.runtimeRequests.find((r) => r.type === "SetXCom")!.body; + expect(setReq).toMatchObject({ + key: "echo", + value: "node says: hello from airflow", + dag_id: "test_dag", + task_id: "say_hello_client", + run_id: "r1", + }); + }); + + it("returns null from getVariable when the supervisor signals NOT_FOUND", async () => { + let observed: string | null = ""; + registerTask({ dagId: "test_dag", taskId: "missing_variable" }, async ({ client }) => { + observed = await client.getVariable("missing_key"); + }); + + const responder: Responder = (msgType) => { + if (msgType === "GetVariable") { + return { + body: { + type: "ErrorResponse", + error: "VARIABLE_NOT_FOUND", + detail: { key: "missing_key" }, + }, + }; + } + return null; + }; + + const result = await driveSupervisor(makeStartupDetails("missing_variable"), responder); + expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" }); + expect(observed).toBeNull(); + }); + + it("looks up handlers by exact Dag and task id", async () => { + let calledFirstDag = false; + let calledSecondDag = false; + registerTask({ dagId: "test_dag", taskId: "shared_task" }, async () => { + calledFirstDag = true; + }); + registerTask({ dagId: "other_dag", taskId: "shared_task" }, async () => { + calledSecondDag = true; + }); + + await driveSupervisor(makeStartupDetails("shared_task")); + + expect(calledFirstDag).toBe(true); + expect(calledSecondDag).toBe(false); + }); + + it("returns empty serialized_dags for DagFileParseRequest", async () => { + const parseRequest = { + type: "DagFileParseRequest", + file: "/dags/test.mjs", + bundle_path: "/dags", + }; + + const result = await driveSupervisor(parseRequest); + + const body = result.firstResponse!.body as Record; + expect(body.type).toBe("DagFileParsingResult"); + expect(body.serialized_dags).toEqual([]); + }); + + it("auto-pushes return_value XCom when handler returns a value", async () => { + registerTask({ dagId: "test_dag", taskId: "pusher" }, async () => "my-result"); + + const responder: Responder = (msgType, _body) => { + if (msgType === "SetXCom") return { body: null }; + return null; + }; + + const result = await driveSupervisor(makeStartupDetails("pusher"), responder); + + expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" }); + + const setXComReqs = result.runtimeRequests.filter((r) => r.type === "SetXCom"); + expect(setXComReqs).toHaveLength(1); + expect(setXComReqs[0]!.body).toMatchObject({ + key: "return_value", + value: "my-result", + }); + }); + + it("does NOT push return_value XCom when handler returns undefined", async () => { + registerTask({ dagId: "test_dag", taskId: "void_task" }, async () => { + // no return value + }); + + const result = await driveSupervisor(makeStartupDetails("void_task")); + + expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" }); + + const setXComReqs = result.runtimeRequests.filter((r) => r.type === "SetXCom"); + expect(setXComReqs).toHaveLength(0); + }); +}); diff --git a/ts-sdk/tests/coordinator/log-channel.test.ts b/ts-sdk/tests/coordinator/log-channel.test.ts new file mode 100644 index 0000000000000..2b8a79579dc4e --- /dev/null +++ b/ts-sdk/tests/coordinator/log-channel.test.ts @@ -0,0 +1,181 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import * as net from "node:net"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { LogChannel } from "../../src/coordinator/log-channel.js"; + +interface Fixture { + server: net.Server; + port: number; + received: Buffer[]; + sockClosed: Promise; +} + +async function makeServer(): Promise { + const received: Buffer[] = []; + let resolveSockClosed: () => void; + const sockClosed = new Promise((r) => { + resolveSockClosed = r; + }); + const server = net.createServer((sock) => { + sock.on("data", (chunk) => received.push(chunk)); + sock.on("close", () => resolveSockClosed()); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const port = (server.address() as net.AddressInfo).port; + return { server, port, received, sockClosed }; +} + +function readRecords(received: Buffer[]): Record[] { + return Buffer.concat(received) + .toString("utf8") + .split("\n") + .filter(Boolean) + .map((line) => JSON.parse(line) as Record); +} + +describe("LogChannel", () => { + let fx: Fixture; + + beforeEach(async () => { + fx = await makeServer(); + }); + + afterEach(async () => { + fx.server.close(); + }); + + it("defaults logger name to 'ts-sdk' and auto-stamps timestamp", async () => { + const ch = await LogChannel.connect(`127.0.0.1:${fx.port}`); + ch.info("hello"); + await ch.close(); + await fx.sockClosed; + const records = readRecords(fx.received); + expect(records).toHaveLength(1); + const record = records[0]!; + // Logger name kept as a JSON field AND prepended to the event so + // it surfaces in the Airflow UI text renderer. + expect(record).toMatchObject({ + event: "[ts-sdk] hello", + level: "info", + logger: "ts-sdk", + }); + expect(typeof record["timestamp"]).toBe("string"); + expect(new Date(record["timestamp"] as string).toString()).not.toBe("Invalid Date"); + }); + + it("accepts a custom root name", async () => { + const ch = await LogChannel.connect(`127.0.0.1:${fx.port}`, "ts-sdk.runtime"); + ch.warning("started"); + await ch.close(); + await fx.sockClosed; + const records = readRecords(fx.received); + expect(records).toHaveLength(1); + const record = records[0]!; + expect(record).toMatchObject({ + event: "[ts-sdk.runtime] started", + level: "warning", + logger: "ts-sdk.runtime", + }); + }); + + it("child() creates a hierarchical sibling sharing the socket", async () => { + const root = await LogChannel.connect(`127.0.0.1:${fx.port}`); + const comm = root.child("comm"); + const client = root.child("client"); + expect(comm.loggerName).toBe("ts-sdk.comm"); + expect(client.loggerName).toBe("ts-sdk.client"); + + root.info("a"); + comm.debug("b"); + client.error("c"); + await root.close(); + await fx.sockClosed; + + const records = readRecords(fx.received); + expect(records.map((r) => r["logger"])).toEqual(["ts-sdk", "ts-sdk.comm", "ts-sdk.client"]); + expect(records.map((r) => r["level"])).toEqual(["info", "debug", "error"]); + }); + + it("children must not close the shared socket", async () => { + const root = await LogChannel.connect(`127.0.0.1:${fx.port}`); + const child = root.child("comm"); + // Child.close() is a no-op. Root.close() ends the socket. + await child.close(); + // The socket should still be open — write through root to prove it. + root.info("still alive"); + await root.close(); + await fx.sockClosed; + const records = readRecords(fx.received); + expect(records).toHaveLength(1); + expect(records[0]).toMatchObject({ + event: "[ts-sdk] still alive", + logger: "ts-sdk", + }); + }); + + it("handles post-connect socket errors on the root channel", async () => { + const write = vi.spyOn(process.stderr, "write").mockImplementation(() => true); + const root = await LogChannel.connect(`127.0.0.1:${fx.port}`); + root.child("child"); + const sock = (root as unknown as { sock: net.Socket }).sock; + + try { + sock.emit("error", new Error("boom")); + expect(write).toHaveBeenCalledTimes(1); + expect(write.mock.calls[0]?.[0]).toBe("[ts-sdk] log socket error: boom\n"); + } finally { + write.mockRestore(); + await root.close(); + await fx.sockClosed; + } + }); + + it("drops records sent after close instead of writing to the ended socket", async () => { + const root = await LogChannel.connect(`127.0.0.1:${fx.port}`); + const child = root.child("comm"); + root.info("before close"); + await root.close(); + + root.info("after close"); + child.debug("after close via child"); + + await fx.sockClosed; + const records = readRecords(fx.received); + expect(records).toHaveLength(1); + expect(records[0]).toMatchObject({ event: "[ts-sdk] before close" }); + }); + + it("reports close-time socket errors", async () => { + const write = vi.spyOn(process.stderr, "write").mockImplementation(() => true); + const root = await LogChannel.connect(`127.0.0.1:${fx.port}`); + const sock = (root as unknown as { sock: net.Socket }).sock; + + try { + sock.emit("error", Object.assign(new Error("write EPIPE"), { code: "EPIPE" })); + expect(write).toHaveBeenCalledTimes(1); + expect(write.mock.calls[0]?.[0]).toBe("[ts-sdk] log socket error: write EPIPE\n"); + } finally { + write.mockRestore(); + await root.close(); + await fx.sockClosed; + } + }); +}); diff --git a/ts-sdk/tests/coordinator/protocol.test.ts b/ts-sdk/tests/coordinator/protocol.test.ts new file mode 100644 index 0000000000000..330f94b62b05d --- /dev/null +++ b/ts-sdk/tests/coordinator/protocol.test.ts @@ -0,0 +1,101 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { describe, expect, it } from "vitest"; +import { asMsgFromSupervisor } from "../../src/coordinator/protocol.js"; +import { parseArgs, startCoordinator } from "../../src/coordinator/runtime.js"; + +describe("protocol decode", () => { + it("accepts StartupDetails", () => { + const raw = { + type: "StartupDetails", + ti: { + id: "u", + dag_version_id: "dag-version-1", + task_id: "t", + dag_id: "d", + run_id: "r", + try_number: 1, + map_index: -1, + hostname: "test-host", + queue: "default", + }, + dag_rel_path: "dags/my.mjs", + bundle_info: { name: "my", version: "1" }, + start_date: "2026-04-23T00:00:00Z", + ti_context: {}, + }; + expect(asMsgFromSupervisor(raw).type).toBe("StartupDetails"); + }); + + it("accepts DagFileParseRequest", () => { + const raw = { + type: "DagFileParseRequest", + file: "/x.mjs", + bundle_path: "/", + }; + expect(asMsgFromSupervisor(raw).type).toBe("DagFileParseRequest"); + }); + + it("rejects missing type", () => { + expect(() => asMsgFromSupervisor({ ti: {} })).toThrow(/missing string 'type'/); + }); + + it("rejects unknown type", () => { + expect(() => asMsgFromSupervisor({ type: "WeirdMsg" })).toThrow( + /Unsupported supervisor message type/, + ); + }); +}); + +describe("runtime arg parser", () => { + it("parses --comm and --logs from argv", () => { + const r = parseArgs(["node", "bundle.mjs", "--comm=127.0.0.1:5001", "--logs=127.0.0.1:5002"]); + expect(r.commAddr).toBe("127.0.0.1:5001"); + expect(r.logsAddr).toBe("127.0.0.1:5002"); + }); + + it("throws when --comm is missing", () => { + expect(() => parseArgs(["node", "bundle.mjs", "--logs=127.0.0.1:5002"])).toThrow( + /Missing --comm/, + ); + }); + + it("throws when --logs is missing", () => { + expect(() => parseArgs(["node", "bundle.mjs", "--comm=127.0.0.1:5001"])).toThrow( + /Missing --logs/, + ); + }); + + it("requires commAddr and logsAddr overrides to be supplied together", async () => { + await expect( + startCoordinator({ + commAddr: "127.0.0.1:5001", + argv: ["node", "bundle.mjs", "--logs=127.0.0.1:5002"], + }), + ).rejects.toThrow(/Missing --comm/); + + await expect( + startCoordinator({ + logsAddr: "127.0.0.1:5002", + argv: ["node", "bundle.mjs", "--comm=127.0.0.1:5001"], + }), + ).rejects.toThrow(/Missing --logs/); + }); +}); diff --git a/ts-sdk/tests/coordinator/public-api.test.ts b/ts-sdk/tests/coordinator/public-api.test.ts new file mode 100644 index 0000000000000..218fa42508ea4 --- /dev/null +++ b/ts-sdk/tests/coordinator/public-api.test.ts @@ -0,0 +1,35 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { describe, expectTypeOf, it } from "vitest"; +import type { StartCoordinatorOptions } from "../../src/coordinator/index.js"; +import { startCoordinator } from "../../src/coordinator/index.js"; + +describe("coordinator public API", () => { + it("exports the coordinator runtime entrypoint from the coordinator subpath", () => { + expectTypeOf().toEqualTypeOf< + (opts?: StartCoordinatorOptions) => Promise + >(); + expectTypeOf().toEqualTypeOf<{ + commAddr?: string; + logsAddr?: string; + argv?: readonly string[]; + }>(); + }); +}); diff --git a/ts-sdk/tests/coordinator/runtime-signal.test.ts b/ts-sdk/tests/coordinator/runtime-signal.test.ts new file mode 100644 index 0000000000000..c906127b07737 --- /dev/null +++ b/ts-sdk/tests/coordinator/runtime-signal.test.ts @@ -0,0 +1,87 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { EventEmitter } from "node:events"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { ABORT_GRACE_PERIOD_MS, createRuntimeAbort } from "../../src/coordinator/runtime.js"; + +describe("coordinator runtime signal handling", () => { + afterEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); + }); + + it("aborts on SIGTERM and force-exits when the grace period expires", () => { + vi.useFakeTimers(); + const events = new EventEmitter(); + const exitProcess = vi.fn((code: number): never => { + throw new Error(`exit ${code}`); + }); + const runtimeAbort = createRuntimeAbort(null, { + signalSource: events, + exitProcess, + }); + + events.emit("SIGTERM", "SIGTERM"); + + expect(runtimeAbort.signal.aborted).toBe(true); + expect(exitProcess).not.toHaveBeenCalled(); + expect(() => vi.advanceTimersByTime(ABORT_GRACE_PERIOD_MS)).toThrow("exit 1"); + expect(exitProcess).toHaveBeenCalledWith(1); + runtimeAbort.dispose(); + }); + + it("uses the registered signal name when the event does not pass a payload", () => { + vi.useFakeTimers(); + const events = new EventEmitter(); + const exitProcess = vi.fn((code: number): never => { + throw new Error(`exit ${code}`); + }); + const runtimeAbort = createRuntimeAbort(null, { + signalSource: events, + exitProcess, + }); + + events.emit("SIGTERM"); + + expect(runtimeAbort.signal.aborted).toBe(true); + expect((runtimeAbort.signal.reason as Error).message).toBe("Task aborted by SIGTERM"); + runtimeAbort.dispose(); + }); + + it("removes signal listeners and clears the force-exit timer on dispose", () => { + vi.useFakeTimers(); + const events = new EventEmitter(); + const exitProcess = vi.fn((code: number): never => { + throw new Error(`exit ${code}`); + }); + const runtimeAbort = createRuntimeAbort(null, { + signalSource: events, + exitProcess, + }); + + events.emit("SIGINT", "SIGINT"); + runtimeAbort.dispose(); + vi.advanceTimersByTime(ABORT_GRACE_PERIOD_MS); + + expect(exitProcess).not.toHaveBeenCalled(); + expect(events.listenerCount("SIGINT")).toBe(0); + expect(events.listenerCount("SIGTERM")).toBe(0); + }); +}); diff --git a/ts-sdk/tests/coordinator/tcp-connect.test.ts b/ts-sdk/tests/coordinator/tcp-connect.test.ts new file mode 100644 index 0000000000000..8b0ecdec811a6 --- /dev/null +++ b/ts-sdk/tests/coordinator/tcp-connect.test.ts @@ -0,0 +1,50 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { describe, expect, it } from "vitest"; +import { connectTcp, splitHostPort } from "../../src/coordinator/tcp-connect.js"; + +describe("splitHostPort", () => { + it.each([ + // [input, host, port] + ["127.0.0.1:8080", "127.0.0.1", "8080"], + ["localhost:5432", "localhost", "5432"], + // Splits on the LAST colon, so a bare IPv6 host survives intact. + // (Result feeds Node's connect() as-is; the bracketed "[::1]" + // form is intentionally not handled — Node doesn't strip + // brackets, so it wouldn't connect anyway.) + ["::1:8080", "::1", "8080"], + ])("splits %s into host=%s port=%s", (addr, host, port) => { + expect(splitHostPort(addr)).toEqual([host, port]); + }); + + it("throws with the exact host:port message when there is no colon", () => { + expect(() => splitHostPort("localhost")).toThrow("Address must be host:port, got localhost"); + }); + + it("throws on the empty string (no colon)", () => { + expect(() => splitHostPort("")).toThrow(/Address must be host:port/); + }); +}); + +describe("connectTcp", () => { + it("rejects (does not throw synchronously) on a malformed address", async () => { + await expect(connectTcp("nocolon")).rejects.toThrow("Address must be host:port, got nocolon"); + }); +}); diff --git a/ts-sdk/tests/public-api.test.ts b/ts-sdk/tests/public-api.test.ts index 1243092a9b4d5..751b749539722 100644 --- a/ts-sdk/tests/public-api.test.ts +++ b/ts-sdk/tests/public-api.test.ts @@ -22,11 +22,18 @@ import type { ConnectionResult, GetXComOpts, SetXComOpts, + StartCoordinatorOptions, TaskClient, TaskContext, TaskRegistration, } from "../src/index.js"; -import { listRegisteredTasks, registerTask, VariableNotFoundError } from "../src/index.js"; +import { + listRegisteredTasks, + registerTask, + startCoordinator, + SUPERVISOR_API_VERSION, + VariableNotFoundError, +} from "../src/index.js"; describe("public API", () => { it("exports task registration helpers", () => { @@ -42,6 +49,18 @@ describe("public API", () => { expect(err.key).toBe("missing"); }); + it("exports the coordinator runtime entrypoint", () => { + expectTypeOf().toEqualTypeOf< + (opts?: StartCoordinatorOptions) => Promise + >(); + expectTypeOf().toEqualTypeOf<{ + commAddr?: string; + logsAddr?: string; + argv?: readonly string[]; + }>(); + expectTypeOf(SUPERVISOR_API_VERSION).toMatchTypeOf(); + }); + it("uses idiomatic TypeScript names for public client types", () => { expectTypeOf().toEqualTypeOf<{ readonly dagId: string; @@ -94,8 +113,6 @@ describe("public API", () => { function acceptsSetXComOpts(_opts: SetXComOpts): void {} function acceptsTaskRegistration(_registration: TaskRegistration): void {} - // TODO: Add coordinator-runtime tests that validate these camelCase - // public options map to the supervisor schema's snake_case fields. acceptsGetXComOpts({ key: "result", dagId: "example",