diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst index d374225fca492..d0eeac25f85d0 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -175,6 +175,34 @@ one that pushed ``null`` binds ``null``. A Python ``int`` beyond the ±9007199254740991 a JavaScript number holds exactly is refused rather than bound, so carry such a value across the language boundary as a string. +Explicit renames +~~~~~~~~~~~~~~~~ + +``withArgNames`` states a binding when folding cannot reach it, for a name the Python side never used: +a clearer word than the Dag chose, or a TypeScript reserved word like ``enum``. +The mapping comes first, the handler second: + +.. code-block:: typescript + + interface ReportArgs { + summary: Summary; + label: string; // Python calls this `run_label` + } + + const report = withArgNames({ label: "run_label" }, async ({ summary, label }: ReportArgs) => { + // `label` is the call's `run_label`; `summary` folded as usual. + }); + + bundle.register(new TaskHandler("etl", "report", report)); + +An entry beats folding, and everything the map does not mention still folds, +so ``withArgNames`` should be rare in a real Dag. +A mapped name the call did not pass misses rather than falling back to folding. + +The map's keys are checked against the handler's own parameter type, so ``{ labl: "run_label" }`` is a +compile error naming the right key. Its values are Python names, which ``tsc`` cannot see and does not +check. + .. note:: Being upstream is not the same as being passed. As with the other language SDKs, an XCom *dependency* diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/ts_sdk_tests/test_ts_sdk_dag.py b/airflow-e2e-tests/tests/airflow_e2e_tests/ts_sdk_tests/test_ts_sdk_dag.py index 58e5360a80434..69c2e662ca189 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/ts_sdk_tests/test_ts_sdk_dag.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/ts_sdk_tests/test_ts_sdk_dag.py @@ -28,8 +28,9 @@ round-trips, and task logs reaching the log store. ``typescript_taskflow_example`` covers TaskFlow arguments, including an upstream output pulled before -the handler runs, and shares a ``build_message`` task ID with ``typescript_example`` so that dispatch -keying on the task ID alone would run the wrong handler. +the handler runs and a ``withArgNames`` rename on its ``report`` task, and shares a ``build_message`` +task ID with ``typescript_example`` so that dispatch keying on the task ID alone would run the wrong +handler. """ from __future__ import annotations @@ -153,7 +154,12 @@ def test_second_dag_from_the_same_bundle_succeeded(completed_taskflow_run: _Comp f"expected the run to succeed; got {completed_taskflow_run.state!r}. " f"task states: {completed_taskflow_run.ti_states}" ) - expected = {"make_totals": "success", "summarize": "success", "build_message": "success"} + expected = { + "make_totals": "success", + "summarize": "success", + "report": "success", + "build_message": "success", + } for task_id, want in expected.items(): assert completed_taskflow_run.ti_states.get(task_id) == want, ( f"{task_id!r} expected {want!r}. all task states: {completed_taskflow_run.ti_states}" @@ -190,6 +196,24 @@ def test_summarize_binds_its_call_arguments(completed_taskflow_run: _CompletedRu assert completed_taskflow_run.xcom("summarize", key="summary_line") == "uk: 12 orders" +def test_report_binds_an_explicitly_renamed_argument(completed_taskflow_run: _CompletedRun): + """``report(summary, "nightly")`` reaches a handler that renamed one argument. + + Python names it ``run_label``; the handler destructures ``label``, a word + the ``@task.stub`` signature never uses, so folding could not connect the + two and the binding is stated with ``withArgNames``. The handler throws + unless ``label`` is exactly ``"nightly"``, so a rename that did not take + effect fails this task rather than returning a null. + + ``summary`` is not renamed: folding already covers it, which is the point + that keeps ``withArgNames`` rare. + """ + value = completed_taskflow_run.xcom("report") + assert value == {"label": "nightly", "regionCode": "uk", "healthy": True}, ( + f"unexpected 'report' return_value: {value!r}" + ) + + def test_same_task_id_under_two_dags_runs_its_own_handler( completed_run: _CompletedRun, completed_taskflow_run: _CompletedRun ): diff --git a/ts-sdk/README.md b/ts-sdk/README.md index e82dc090a0535..db2ef2f02cde3 100644 --- a/ts-sdk/README.md +++ b/ts-sdk/README.md @@ -241,6 +241,30 @@ const rows = await getClient().getXCom({ key: "return_value", taskId: "e A Python `int` beyond the ±9007199254740991 a JavaScript number holds exactly is refused rather than bound, so carry such a value across the boundary as a string. +### Explicit renames + +`withArgNames` states a binding folding cannot reach, for a name the Python side never used: +a clearer word than the Dag chose, or a TypeScript reserved word like `enum`. Mapping first, handler second: + +```ts +interface ReportArgs { + summary: Summary; + label: string; // Python calls this `run_label` +} + +const report = withArgNames({ label: "run_label" }, async ({ summary, label }: ReportArgs) => { + // `label` is the call's `run_label`; `summary` folded as usual. +}); + +bundle.register(new TaskHandler("etl", "report", report)); +``` + +An entry beats folding, and everything the map does not mention still folds, +so `withArgNames` should be rare in a real Dag. +The map's keys are checked against the handler's own parameter type, +so `{ labl: "run_label" }` is a compile error naming the right key. +Its values are Python names, which `tsc` cannot see and does not check. + `Dag` is another interface, for a Dag declared natively in TypeScript, and is still a work in progress. Airflow launches the bundled entrypoint with `--comm=host:port` and diff --git a/ts-sdk/api-docs/dag-authoring-api.ts b/ts-sdk/api-docs/dag-authoring-api.ts index da90ccb81d47a..79b8ddf5e7263 100644 --- a/ts-sdk/api-docs/dag-authoring-api.ts +++ b/ts-sdk/api-docs/dag-authoring-api.ts @@ -19,8 +19,9 @@ /** @module Authoring */ -export { Bundle, Dag, getClient, getContext, TaskHandler } from "../src/index.js"; +export { Bundle, Dag, getClient, getContext, TaskHandler, withArgNames } from "../src/index.js"; export type { + ArgNameMap, DagSpec, Registerable, TaskClient, diff --git a/ts-sdk/docs/index.md b/ts-sdk/docs/index.md index ab012a397d3fd..7d60fb0d94547 100644 --- a/ts-sdk/docs/index.md +++ b/ts-sdk/docs/index.md @@ -71,6 +71,9 @@ export async function transform({ regionCode, threshold }: TransformArgs) { } ``` +`withArgNames` states a binding folding cannot reach, for a name the Python side never used. +It should be rare, since folding covers ordinary spelling differences. + `Dag` is another interface, for a Dag declared in TypeScript rather than in Python, and is still a work in progress. ## Coordinators diff --git a/ts-sdk/example/dags/typescript_taskflow_example.py b/ts-sdk/example/dags/typescript_taskflow_example.py index 8f8ad5f648d42..ce096906d9a47 100644 --- a/ts-sdk/example/dags/typescript_taskflow_example.py +++ b/ts-sdk/example/dags/typescript_taskflow_example.py @@ -21,7 +21,9 @@ ``summarize`` is called TaskFlow-style, and every argument its call passes reaches the TypeScript handler by name, including ``make_totals``'s output, which the runtime pulls before the handler runs. -Its ``build_message`` stub shares a ``task_id`` with a task in ``typescript_example`` on purpose: +``report`` shows the one case folding cannot cover: its handler wants a name the Dag never used, +so it states that binding explicitly with ``withArgNames``. +The ``build_message`` stub shares a ``task_id`` with a task in ``typescript_example`` on purpose: a handler binds the ``(dag_id, task_id)`` pair, so the two are different tasks. See ``src/taskflow.ts``. """ @@ -46,6 +48,13 @@ def make_totals(): def summarize(totals: dict, region_code: str, currency: str, threshold: float, dry_run: bool = False): ... +# `run_label` is not a spelling difference. The handler wants to call it +# `label`, a word this signature never uses, which is what `withArgNames` is +# for; folding would never connect the two. +@task.stub(queue="typescript") +def report(summary: dict, run_label: str): ... + + # Same task_id as `typescript_example.build_message`, on purpose. @task.stub(queue="typescript") def build_message(): ... @@ -58,7 +67,9 @@ def build_message(): ... tags=["typescript", "example", "taskflow"], ) def typescript_taskflow_example(): - summarize(make_totals(), "uk", "GBP", 280.0) >> build_message() + summary = summarize(make_totals(), "uk", "GBP", 280.0) + report(summary, "nightly") + summary >> build_message() typescript_taskflow_example() diff --git a/ts-sdk/example/src/main.ts b/ts-sdk/example/src/main.ts index d1644ade4e22e..20e0ac8edffed 100644 --- a/ts-sdk/example/src/main.ts +++ b/ts-sdk/example/src/main.ts @@ -24,7 +24,7 @@ import { Bundle, getClient, TaskHandler } from "apache-airflow-ts-sdk"; -import { buildSummaryMessage, summarize } from "./taskflow.js"; +import { buildSummaryMessage, report, summarize } from "./taskflow.js"; export async function buildMessage() { const client = getClient(); @@ -63,6 +63,7 @@ bundle.register( new TaskHandler("typescript_example", "build_message", buildMessage), new TaskHandler("typescript_example", "read_connection", readConnection), new TaskHandler("typescript_taskflow_example", "summarize", summarize), + new TaskHandler("typescript_taskflow_example", "report", report), new TaskHandler("typescript_taskflow_example", "build_message", buildSummaryMessage), ); await bundle.serve(); diff --git a/ts-sdk/example/src/taskflow.ts b/ts-sdk/example/src/taskflow.ts index f51b130a7ea33..86f7de180397c 100644 --- a/ts-sdk/example/src/taskflow.ts +++ b/ts-sdk/example/src/taskflow.ts @@ -28,7 +28,7 @@ // `buildSummaryMessage` implements a task named `build_message`, exactly as the other Dag has, and // the two share nothing else. -import { getClient, getContext } from "apache-airflow-ts-sdk"; +import { getClient, getContext, withArgNames } from "apache-airflow-ts-sdk"; /** What `make_totals` returns on the Python side. */ export interface Totals { @@ -88,6 +88,34 @@ export async function summarize({ }; } +/** Every argument the Dag's `report(...)` call binds, as the handler wants them. */ +export interface ReportArgs { + summary: Summary; + /** The call's `run_label`, which folding cannot reach, so it is mapped below. */ + label: string; +} + +/** + * Renaming an argument the Python side named something else entirely. + * + * The mapping comes first, the handler second. `summary` is absent from the map + * because folding already reaches it. + */ +export const report = withArgNames( + { label: "run_label" }, + async ({ summary, label }: ReportArgs) => { + if (label !== "nightly") { + throw new Error(`expected run label "nightly" but got "${label}"`); + } + + return { + label, + regionCode: summary.regionCode, + healthy: summary.passed, + }; + }, +); + export async function buildSummaryMessage() { // Nothing was passed to this task, so its upstream's output is read explicitly. const ctx = getContext(); diff --git a/ts-sdk/src/coordinator/arg-binding.ts b/ts-sdk/src/coordinator/arg-binding.ts index 5e426ec58c002..357932fcd111b 100644 --- a/ts-sdk/src/coordinator/arg-binding.ts +++ b/ts-sdk/src/coordinator/arg-binding.ts @@ -79,6 +79,8 @@ export interface ArgBindingDeps { /** The task's abort signal, so a terminated task stops mid-pull. */ readonly signal: AbortSignal; readonly logs: LogChannel; + /** Renames the handler declared with `withArgNames`, which beat folding. */ + readonly argNames: ReadonlyMap; } /** @@ -128,7 +130,7 @@ export async function resolveArgs( // Literals need no request, so only a call that pulls races the abort signal. const entries = pullsUpstream ? await abortable(resolveAll, deps.signal) : await resolveAll(); - return { args: makeArgsProxy(names, byFold, new Map(entries), deps.logs), names }; + return { args: makeArgsProxy(names, byFold, new Map(entries), deps), names }; } /** Airflow omits `value` for a literal whose value is null. */ @@ -223,15 +225,21 @@ function abortError(signal: AbortSignal): Error { * has no way to know which spelling a handler will use: it sees Python's names * and nothing else. Folding on read means binding needs nothing declared on * either side, and no guess about the TypeScript name is ever materialized. + * It is also what lets a `withArgNames` entry take precedence, decided per read. */ function makeArgsProxy( names: readonly string[], byFold: ReadonlyMap, values: ReadonlyMap, - logs: LogChannel, + deps: ArgBindingDeps, ): object { - const resolve = (property: string): string | undefined => - values.has(property) ? property : byFold.get(foldArgName(property)); + const { argNames, logs } = deps; + const resolve = (property: string): string | undefined => { + // An explicit rename wins and never falls back to folding, so a wrong entry misses. + const renamed = argNames.get(property); + if (renamed !== undefined) return values.has(renamed) ? renamed : undefined; + return values.has(property) ? property : byFold.get(foldArgName(property)); + }; // A null prototype so a read never reaches Object.prototype: a Python // argument named `constructor` or `toString` must bind like any other, and a @@ -249,6 +257,8 @@ function makeArgsProxy( // one from a typo. logs.warning("Task argument not bound by this task's call", { requested: property, + // Tells a wrong `withArgNames` entry apart from an argument the call never passed. + renamed_to: argNames.get(property) ?? null, bound: [...names], }); return undefined; diff --git a/ts-sdk/src/coordinator/runtime.ts b/ts-sdk/src/coordinator/runtime.ts index e7142cb0b6c5b..deb0cb485a005 100644 --- a/ts-sdk/src/coordinator/runtime.ts +++ b/ts-sdk/src/coordinator/runtime.ts @@ -54,6 +54,7 @@ import { type RuntimeTaskState, type StartupDetails, } from "./protocol.js"; +import { getArgNames } from "../sdk/arg-names.js"; import { listBundleTasks, type Bundle } from "../sdk/bundle.js"; import { runInTaskScope, type TaskContext } from "../sdk/task.js"; import type { JsonValue } from "../sdk/client-types.js"; @@ -315,6 +316,7 @@ async function handleTask( client, signal: ctx.signal, logs, + argNames: getArgNames(handler), }); } catch (err) { // Before the handler ran, so nothing it might have written is at stake. diff --git a/ts-sdk/src/index.ts b/ts-sdk/src/index.ts index 935d58f95fd2a..e6d162155e31e 100644 --- a/ts-sdk/src/index.ts +++ b/ts-sdk/src/index.ts @@ -20,9 +20,11 @@ export { Dag } from "./sdk/dag.js"; export { Bundle } from "./sdk/bundle.js"; export { TaskHandler } from "./sdk/task-handler.js"; +export { withArgNames } from "./sdk/arg-names.js"; export { getClient, getContext } from "./sdk/task.js"; export { ConnectionNotFoundError, VariableNotFoundError } from "./sdk/client.js"; export { SUPERVISOR_API_VERSION } from "./coordinator/index.js"; +export type { ArgNameMap } from "./sdk/arg-names.js"; export type { Registerable } from "./sdk/bundle.js"; export type { DagSpec, TaskInputs, TaskOptions, TaskRef, TaskSpec } from "./sdk/dag.js"; export type { TaskClient } from "./sdk/client.js"; diff --git a/ts-sdk/src/sdk/arg-names.ts b/ts-sdk/src/sdk/arg-names.ts new file mode 100644 index 0000000000000..697966b2b17e7 --- /dev/null +++ b/ts-sdk/src/sdk/arg-names.ts @@ -0,0 +1,112 @@ +/*! + * 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. + */ + +// Stating a binding explicitly. + +import type { TaskFunction } from "./task.js"; + +/** + * A rename per argument, keyed by the handler's own parameter names. + * + * Every key is checked against the handler's parameter type, so a typo is a + * compile error naming the right key. The values are Python names, which `tsc` + * cannot see and does not check. + */ +export type ArgNameMap = { + readonly [K in keyof TArgs]?: string; +}; + +// Read back off the handler at the dispatch site. A global symbol, as the brands are: +// two resolved copies must agree on the key, or a handler wrapped by one loses its renames. +const ARG_NAMES = Symbol.for("airflow.ts-sdk.arg-names"); + +function validate(names: ArgNameMap): ReadonlyMap { + const candidate: unknown = names; + if ( + typeof candidate !== "object" || + candidate === null || + Array.isArray(candidate) || + ![Object.prototype, null].includes(Object.getPrototypeOf(candidate) as object | null) + ) { + throw new Error("withArgNames(...) takes a plain object mapping argument names to wire names"); + } + const entries = new Map(); + for (const [key, value] of Object.entries(candidate as Record)) { + if (typeof value !== "string" || value.length === 0) { + throw new Error( + `withArgNames(...) maps "${key}" to ${JSON.stringify(value)}; a wire name must be a non-empty string`, + ); + } + entries.set(key, value); + } + return entries; +} + +/** + * Bind an argument explicitly, naming it on both sides. + * + * Mapping first, handler second. An entry takes precedence over folding, and + * everything the map does not mention still folds: + * + * ```ts + * interface ReportArgs { + * label: string; // Python calls this `run_label` + * threshold: number; + * } + * + * const report = withArgNames({ label: "run_label" }, async ({ label, threshold }: ReportArgs) => { + * // `label` is the call's `run_label`; `threshold` folded as usual. + * }); + * + * bundle.register(new TaskHandler("etl", "report", report)); + * ``` + * + * Folding absorbs ordinary spelling differences, so this is for a name that genuinely differs: + * a clearer word than the Dag chose, or a TypeScript reserved word like `enum`. + * It should be rare in a real Dag. + * + * The map's keys are checked against the handler's own parameter type, + * so `{ labl: "run_label" }` is a compile error. + * Its values are Python names, which `tsc` cannot check. + */ +export function withArgNames( + // NoInfer, so `TArgs` comes from the handler alone: + // inferring it from the map too would make every key correct and check nothing. + names: ArgNameMap>, + handler: TaskFunction, +): TaskFunction { + const resolved = validate(names); + if (typeof handler !== "function") { + throw new Error("withArgNames(...) takes the handler function as its second argument"); + } + // A wrapper rather than a property on the author's own function: + // one handler can be registered for two tasks that rename differently. + const wrapped: TaskFunction = (args) => handler(args); + Object.defineProperty(wrapped, ARG_NAMES, { value: resolved }); + return wrapped; +} + +/** Internal: the renames a handler was wrapped with, empty when it has none. */ +export function getArgNames(handler: TaskFunction): ReadonlyMap { + const carrier = handler as unknown as Record; + const names = carrier[ARG_NAMES]; + return names instanceof Map ? (names as ReadonlyMap) : EMPTY_ARG_NAMES; +} + +const EMPTY_ARG_NAMES: ReadonlyMap = new Map(); diff --git a/ts-sdk/tests/coordinator/arg-binding.test.ts b/ts-sdk/tests/coordinator/arg-binding.test.ts index 50fc8893b5a00..1e3f1161d34cb 100644 --- a/ts-sdk/tests/coordinator/arg-binding.test.ts +++ b/ts-sdk/tests/coordinator/arg-binding.test.ts @@ -33,6 +33,8 @@ function xcom(name: string, taskId: string, extra: Record = {}) return { name, kind: "xcom" as const, task_id: taskId, ...extra }; } +const NO_RENAMES: ReadonlyMap = new Map(); + function makeLogs() { const warning = vi.fn(); const logs = { warning } as unknown as LogChannel; @@ -56,12 +58,21 @@ interface BindResult extends BoundArgs { async function bind( bindings: ArgBindings, - opts: { upstream?: Record; signal?: AbortSignal } = {}, + opts: { + upstream?: Record; + signal?: AbortSignal; + argNames?: Record; + } = {}, ): Promise { const { logs, warning } = makeLogs(); const { client, pulls } = makeClient(opts.upstream); const signal = opts.signal ?? new AbortController().signal; - const bound = await resolveArgs(bindings, { client, signal, logs }); + const bound = await resolveArgs(bindings, { + client, + signal, + logs, + argNames: new Map(Object.entries(opts.argNames ?? {})), + }); return { ...bound, warning, pulls }; } @@ -165,12 +176,14 @@ describe("resolveArgs", () => { expect(reigonCode).toBeUndefined(); expect(warning).toHaveBeenCalledWith("Task argument not bound by this task's call", { requested: "runId", + renamed_to: null, bound: ["region_code"], }); // Both the requested name and what the call actually delivered, so a typo // is diagnosable from the task log alone. expect(warning).toHaveBeenCalledWith("Task argument not bound by this task's call", { requested: "reigonCode", + renamed_to: null, bound: ["region_code"], }); }); @@ -293,7 +306,7 @@ describe("resolveArgs", () => { const bound = await resolveArgs( [xcom("a", "t_a"), xcom("b", "t_b"), xcom("c", "t_c"), xcom("d", "t_d")], - { client, signal: new AbortController().signal, logs: makeLogs().logs }, + { client, signal: new AbortController().signal, logs: makeLogs().logs, argNames: NO_RENAMES }, ); expect(peak).toBe(4); @@ -326,6 +339,7 @@ describe("resolveArgs", () => { client, signal: new AbortController().signal, logs: makeLogs().logs, + argNames: NO_RENAMES, }, ), ).rejects.toThrowError(/both fold to "regioncode"/); @@ -344,6 +358,7 @@ describe("resolveArgs", () => { client, signal: controller.signal, logs: makeLogs().logs, + argNames: NO_RENAMES, }), ).rejects.toThrowError( /Aborted while resolving this task's arguments.*Task aborted by SIGTERM/, @@ -361,6 +376,7 @@ describe("resolveArgs", () => { client, signal: controller.signal, logs: makeLogs().logs, + argNames: NO_RENAMES, }); controller.abort(new Error("Task aborted by SIGTERM")); @@ -411,6 +427,69 @@ describe("resolveArgs", () => { ); }); + describe("with explicit renames", () => { + it("binds a name the Python side never used", async () => { + const { args } = await bind([literal("run_label", "nightly")], { + argNames: { label: "run_label" }, + }); + expect((args as { label: string }).label).toBe("nightly"); + }); + + it("still folds everything the map does not mention", async () => { + const { args } = await bind([literal("run_label", "nightly"), literal("region_code", "uk")], { + argNames: { label: "run_label" }, + }); + const { label, regionCode } = args as { label: string; regionCode: string }; + expect({ label, regionCode }).toEqual({ label: "nightly", regionCode: "uk" }); + }); + + it("beats a folded match on the same name", async () => { + // An author who stated a binding meant it, so the map wins over the + // name that would otherwise have folded to it. + const { args } = await bind([literal("run_label", "nightly"), literal("label", "folded")], { + argNames: { label: "run_label" }, + }); + expect((args as { label: string }).label).toBe("nightly"); + }); + + it("leaves the renamed wire name reachable under its own name", async () => { + // Enumeration reports Python's names, so rest destructuring must keep + // resolving them whatever the handler renamed. + const { args } = await bind([literal("run_label", "nightly")], { + argNames: { label: "run_label" }, + }); + expect({ ...(args as object) }).toEqual({ run_label: "nightly" }); + expect("run_label" in (args as object)).toBe(true); + expect("label" in (args as object)).toBe(true); + }); + + it("misses rather than falling back when the mapped name was not passed", async () => { + // Falling back to folding would hand the handler a value the SDK + // guessed at, and hide the fact that the stated binding was wrong. + const { args, warning } = await bind([literal("label", "folded")], { + argNames: { label: "run_label" }, + }); + + expect((args as { label?: string }).label).toBeUndefined(); + expect("label" in (args as object)).toBe(false); + // The wire name the handler asked for, so a wrong entry is diagnosable + // from the task log rather than looking like an argument never passed. + expect(warning).toHaveBeenCalledWith("Task argument not bound by this task's call", { + requested: "label", + renamed_to: "run_label", + bound: ["label"], + }); + }); + + it("renames an XCom-backed argument too", async () => { + const { args } = await bind([xcom("run_totals", "make_totals")], { + argNames: { totals: "run_totals" }, + upstream: { make_totals: { found: true, value: { orders: 12 } } }, + }); + expect((args as { totals: { orders: number } }).totals).toEqual({ orders: 12 }); + }); + }); + it("refuses assignment and deletion", async () => { // The bound object mirrors a call site that already happened, so writing // to it would change nothing an author could observe downstream. diff --git a/ts-sdk/tests/coordinator/integration.test.ts b/ts-sdk/tests/coordinator/integration.test.ts index cc999bae648f8..c5a6ffc1eaf7e 100644 --- a/ts-sdk/tests/coordinator/integration.test.ts +++ b/ts-sdk/tests/coordinator/integration.test.ts @@ -37,6 +37,7 @@ import { } from "../../src/coordinator/runtime.js"; import { Dag } from "../../src/sdk/dag.js"; import { Bundle } from "../../src/sdk/bundle.js"; +import { withArgNames } from "../../src/sdk/arg-names.js"; import { TaskHandler } from "../../src/sdk/task-handler.js"; import { getClient, getContext } from "../../src/sdk/task.js"; @@ -519,6 +520,38 @@ describe("coordinator runtime integration", () => { ).toBe(true); }); + it("honours a handler's explicit renames over folding", async () => { + // `withArgNames` travels with the handler, so the dispatch site reads it + // back off the function the bundle holds. + let observed: unknown = null; + bundle.register( + new TaskHandler( + "py_dag", + "renamed", + withArgNames( + { label: "run_label" }, + async ({ label, regionCode }: { label: string; regionCode: string }) => { + observed = { label, regionCode }; + return observed; + }, + ), + ), + ); + + const result = await driveSupervisor( + makeStartupDetails("renamed", "py_dag", "r1", { + arg_bindings: [ + { name: "run_label", kind: "literal", value: "nightly" }, + { name: "region_code", kind: "literal", value: "uk" }, + ], + }), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" }); + // `label` came from the map, `regionCode` folded as usual. + expect(observed).toEqual({ label: "nightly", regionCode: "uk" }); + }); + it("fails the task when two of its bound names fold alike", async () => { // Reported before the handler runs, so nothing it might have written to // XCom is at stake. diff --git a/ts-sdk/tests/public-api.test.ts b/ts-sdk/tests/public-api.test.ts index ccb37542700c2..4b8d6be8171af 100644 --- a/ts-sdk/tests/public-api.test.ts +++ b/ts-sdk/tests/public-api.test.ts @@ -20,6 +20,7 @@ import { afterEach, describe, expect, expectTypeOf, it, vi } from "vitest"; import { AIRFLOW_METADATA_FLAG } from "../src/coordinator/manifest.js"; import type { + ArgNameMap, ConnectionResult, DagSpec, GetXComOpts, @@ -43,6 +44,7 @@ import { SUPERVISOR_API_VERSION, TaskHandler, VariableNotFoundError, + withArgNames, } from "../src/index.js"; describe("public API", () => { @@ -195,6 +197,34 @@ describe("public API", () => { void rejectsFactoryMisuse; }); + it("exports withArgNames for a name the Python side never used", () => { + interface ReportArgs { + label: string; + threshold: number; + } + const report = withArgNames({ label: "run_label" }, async ({ label }: ReportArgs) => label); + + // Wrapping keeps the handler's own type, so the result registers like any + // other handler and nothing at the registration site has to change. + expectTypeOf(report).toEqualTypeOf>(); + expect( + new Bundle(new TaskHandler("etl", "report", report)).getTaskHandler("etl", "report"), + ).toBe(report); + expectTypeOf>().toEqualTypeOf<{ + readonly label?: string; + readonly threshold?: string; + }>(); + + const rejectsUnknownKeys = () => { + // @ts-expect-error "labl" is not a parameter of ReportArgs; "label" is. + withArgNames({ labl: "run_label" }, async ({ label }: ReportArgs) => label); + }; + void rejectsUnknownKeys; + // Reading the renames back is the runtime's business, not an author's. + expectTypeOf().not.toHaveProperty("getArgNames"); + expect("getArgNames" in sdk).toBe(false); + }); + describe("the task-handler getters", () => { it("throw outside a handler, naming the accessor", () => { // The full scope behaviour is covered in tests/sdk/task-scope.test.ts; diff --git a/ts-sdk/tests/sdk/arg-names.test.ts b/ts-sdk/tests/sdk/arg-names.test.ts new file mode 100644 index 0000000000000..5ce427079f0d3 --- /dev/null +++ b/ts-sdk/tests/sdk/arg-names.test.ts @@ -0,0 +1,112 @@ +/*! + * 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, expectTypeOf, it, vi } from "vitest"; + +import { getArgNames, withArgNames, type ArgNameMap } from "../../src/sdk/arg-names.js"; +import type { TaskFunction } from "../../src/sdk/task.js"; + +interface ReportArgs { + label: string; + threshold: number; +} + +describe("withArgNames", () => { + it("carries the renames the runtime reads back", () => { + const handler = async ({ label }: ReportArgs) => label; + const wrapped = withArgNames({ label: "run_label" }, handler); + + expect([...getArgNames(wrapped)]).toEqual([["label", "run_label"]]); + }); + + it("calls the handler it wraps, and returns its value", async () => { + const handler = vi.fn(async ({ label, threshold }: ReportArgs) => `${label}:${threshold}`); + const wrapped = withArgNames({ label: "run_label" }, handler); + + await expect(wrapped({ label: "nightly", threshold: 0.75 })).resolves.toBe("nightly:0.75"); + expect(handler).toHaveBeenCalledWith({ label: "nightly", threshold: 0.75 }); + }); + + it("keeps the handler's own type, so the wrapped value is still registrable", () => { + const wrapped = withArgNames({ label: "run_label" }, async ({ label }: ReportArgs) => label); + expectTypeOf(wrapped).toEqualTypeOf>(); + }); + + it("does not mutate the function it wraps", () => { + // One handler can be registered for two tasks that rename differently, so + // writing the map onto the author's function would let the second + // registration silently change the first. + const handler = async ({ label }: ReportArgs) => label; + const first = withArgNames({ label: "run_label" }, handler); + const second = withArgNames({ label: "report_label" }, handler); + + expect(getArgNames(handler).size).toBe(0); + expect(getArgNames(first).get("label")).toBe("run_label"); + expect(getArgNames(second).get("label")).toBe("report_label"); + }); + + it("reports no renames for a plain handler", () => { + expect(getArgNames(async () => undefined).size).toBe(0); + }); + + it("rejects a map that is not a plain object", () => { + // Typed, so plain JavaScript is what these catch. + const handler = async () => undefined; + for (const bad of [null, undefined, "run_label", ["run_label"], new Map()]) { + expect(() => withArgNames(bad as never, handler)).toThrowError( + /takes a plain object mapping argument names to wire names/, + ); + } + }); + + it("rejects a wire name that is not a non-empty string", () => { + const handler = async ({ label }: ReportArgs) => label; + for (const bad of [1, "", null, {}]) { + expect(() => + withArgNames({ label: bad } as unknown as ArgNameMap, handler), + ).toThrowError(/a wire name must be a non-empty string/); + } + }); + + it("rejects a second argument that is not a function", () => { + expect(() => + withArgNames({ label: "run_label" }, "not a handler" as unknown as TaskFunction), + ).toThrowError(/takes the handler function as its second argument/); + }); + + it("checks the map's keys against the handler's own parameter type", () => { + // The whole point of the type: a mapping for a name the handler does not + // have would silently do nothing, so it must not compile. + const rejectsUnknownKeys = () => { + // @ts-expect-error "labl" is not a parameter of ReportArgs; "label" is. + withArgNames({ labl: "run_label" }, async ({ label }: ReportArgs) => label); + // @ts-expect-error a wire name is a string, not a task reference. + withArgNames({ label: 1 }, async ({ label }: ReportArgs) => label); + // @ts-expect-error the mapping comes first, the handler second. + withArgNames(async ({ label }: ReportArgs) => label, { label: "run_label" }); + }; + void rejectsUnknownKeys; + + // Every key optional, so a handler renames only what it needs to. + expectTypeOf>().toEqualTypeOf<{ + readonly label?: string; + readonly threshold?: string; + }>(); + }); +});