This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new bd65f7998b1 TS SDK: add withArgNames for explicit argument renames
(#73191)
bd65f7998b1 is described below
commit bd65f7998b1192380fb0e4689f17f5375ca06b7e
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Sat Sep 19 11:26:41 2026 +0800
TS SDK: add withArgNames for explicit argument renames (#73191)
Folding absorbs spelling differences, so a Python `region_code` reaches a
handler's `regionCode` with nothing declared. What it cannot reach is a name
the Python side never used: a handler that wants to call `run_label` simply
`label`, or one whose argument is a TypeScript reserved word like `enum`.
Until
now the only options were to rename the Python parameter, or to read the
argument under Python's name and shadow it locally.
`withArgNames` states that binding. The mapping comes first and the handler
second, so the map reads as a qualifier on the handler rather than an extra
parameter, and an entry takes precedence over folding. Everything the map
does
not mention still folds, which is the property that keeps this rare.
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, since a
mapping for a name the handler does not have would silently do nothing. That
check needs the parameter type to come from the handler alone, so the map's
type parameter is `NoInfer`: inferring it from the map too would make every
key
correct by construction and the check vacuous.
A mapped name the call did not pass misses rather than falling back to
folding.
An author who stated a binding that is wrong should see the miss, not a
value
the SDK guessed at, so the unbound-argument warning now also names the wire
name that was asked for.
The renames travel with the handler on a global symbol and are read back at
the
dispatch site, which is also why the returned value is a wrapper rather than
the author's own function with a property attached: one handler can be
registered for two tasks that rename differently, and mutating it would let
the
second registration silently change the first.
---
.../language-sdks/typescript.rst | 28 ++++++
.../ts_sdk_tests/test_ts_sdk_dag.py | 30 +++++-
ts-sdk/README.md | 24 +++++
ts-sdk/api-docs/dag-authoring-api.ts | 3 +-
ts-sdk/docs/index.md | 3 +
ts-sdk/example/dags/typescript_taskflow_example.py | 15 ++-
ts-sdk/example/src/main.ts | 3 +-
ts-sdk/example/src/taskflow.ts | 30 +++++-
ts-sdk/src/coordinator/arg-binding.ts | 18 +++-
ts-sdk/src/coordinator/runtime.ts | 2 +
ts-sdk/src/index.ts | 2 +
ts-sdk/src/sdk/arg-names.ts | 112 +++++++++++++++++++++
ts-sdk/tests/coordinator/arg-binding.test.ts | 85 +++++++++++++++-
ts-sdk/tests/coordinator/integration.test.ts | 33 ++++++
ts-sdk/tests/public-api.test.ts | 30 ++++++
ts-sdk/tests/sdk/arg-names.test.ts | 112 +++++++++++++++++++++
16 files changed, 515 insertions(+), 15 deletions(-)
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 ad033147bfe..bc21f61f6a1 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 58e5360a804..69c2e662ca1 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 @@ Two Dags mix Python tasks with ``@task.stub`` TypeScript
tasks, both served by t
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 e2aba6595f9..869072f1ff8 100644
--- a/ts-sdk/README.md
+++ b/ts-sdk/README.md
@@ -241,6 +241,30 @@ const rows = await getClient().getXCom<number>({ 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 da90ccb81d4..79b8ddf5e72 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 ab012a397d3..7d60fb0d945 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 8f8ad5f648d..ce096906d9a 100644
--- a/ts-sdk/example/dags/typescript_taskflow_example.py
+++ b/ts-sdk/example/dags/typescript_taskflow_example.py
@@ -21,7 +21,9 @@ TaskFlow argument binding across the language boundary.
``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.
[email protected](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 d1644ade4e2..20e0ac8edff 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 f51b130a7ea..86f7de18039 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 5e426ec58c0..357932fcd11 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<string, string>;
}
/**
@@ -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<string, string>,
values: ReadonlyMap<string, JsonValue>,
- 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 e7142cb0b6c..deb0cb485a0 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 935d58f95fd..e6d162155e3 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 00000000000..697966b2b17
--- /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<TArgs> = {
+ 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<unknown>): ReadonlyMap<string, string> {
+ 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<string, string>();
+ for (const [key, value] of Object.entries(candidate as Record<string,
unknown>)) {
+ 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<TArgs, TReturn>(
+ // NoInfer, so `TArgs` comes from the handler alone:
+ // inferring it from the map too would make every key correct and check
nothing.
+ names: ArgNameMap<NoInfer<TArgs>>,
+ handler: TaskFunction<TArgs, TReturn>,
+): TaskFunction<TArgs, TReturn> {
+ 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<TArgs, TReturn> = (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<never, unknown>):
ReadonlyMap<string, string> {
+ const carrier = handler as unknown as Record<symbol, unknown>;
+ const names = carrier[ARG_NAMES];
+ return names instanceof Map ? (names as ReadonlyMap<string, string>) :
EMPTY_ARG_NAMES;
+}
+
+const EMPTY_ARG_NAMES: ReadonlyMap<string, string> = new Map();
diff --git a/ts-sdk/tests/coordinator/arg-binding.test.ts
b/ts-sdk/tests/coordinator/arg-binding.test.ts
index 50fc8893b5a..1e3f1161d34 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<string, unknown> = {})
return { name, kind: "xcom" as const, task_id: taskId, ...extra };
}
+const NO_RENAMES: ReadonlyMap<string, string> = 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<string, XComEntry>; signal?: AbortSignal } = {},
+ opts: {
+ upstream?: Record<string, XComEntry>;
+ signal?: AbortSignal;
+ argNames?: Record<string, string>;
+ } = {},
): Promise<BindResult> {
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 cc999bae648..c5a6ffc1eaf 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 ccb37542700..4b8d6be8171 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<TaskFunction<ReportArgs, string>>();
+ expect(
+ new Bundle(new TaskHandler("etl", "report",
report)).getTaskHandler("etl", "report"),
+ ).toBe(report);
+ expectTypeOf<ArgNameMap<ReportArgs>>().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<typeof sdk>().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 00000000000..5ce427079f0
--- /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<TaskFunction<ReportArgs, string>>();
+ });
+
+ 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<ReportArgs>,
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<ReportArgs>),
+ ).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<ArgNameMap<ReportArgs>>().toEqualTypeOf<{
+ readonly label?: string;
+ readonly threshold?: string;
+ }>();
+ });
+});
