From af5d485ceb46067ff2566c3d032408c63936b7d1 Mon Sep 17 00:00:00 2001 From: "Guan-Ming (Wesley) Chiu" <105915352+guan404ming@users.noreply.github.com> Date: Fri, 2 Oct 2026 14:29:44 +0800 Subject: [PATCH] TS SDK: Add triggerDagRun to native Dags Co-authored-by: Jason(Zhe-You) Liu <68415893+jason810496@users.noreply.github.com> --- .../lang-sdk/0008-control-flow-constructs.md | 10 +- .../language-sdks/typescript.rst | 44 +++ ts-sdk/src/coordinator/client.ts | 57 +++ ts-sdk/src/coordinator/protocol.ts | 8 + ts-sdk/src/coordinator/runtime.ts | 22 +- ts-sdk/src/coordinator/serde.ts | 44 ++- ts-sdk/src/coordinator/trigger-runner.ts | 232 ++++++++++++ ts-sdk/src/index.ts | 2 + ts-sdk/src/sdk/bundle.ts | 12 + ts-sdk/src/sdk/dag.ts | 71 ++-- ts-sdk/src/sdk/trigger-dag-run.ts | 192 ++++++++++ ts-sdk/tests/coordinator/integration.test.ts | 344 ++++++++++++++++++ ts-sdk/tests/sdk/trigger-dag-run.test.ts | 308 ++++++++++++++++ 13 files changed, 1313 insertions(+), 33 deletions(-) create mode 100644 ts-sdk/src/coordinator/trigger-runner.ts create mode 100644 ts-sdk/src/sdk/trigger-dag-run.ts create mode 100644 ts-sdk/tests/sdk/trigger-dag-run.test.ts diff --git a/airflow-core/adr/lang-sdk/0008-control-flow-constructs.md b/airflow-core/adr/lang-sdk/0008-control-flow-constructs.md index 121020491900d..75c4e78b94c7a 100644 --- a/airflow-core/adr/lang-sdk/0008-control-flow-constructs.md +++ b/airflow-core/adr/lang-sdk/0008-control-flow-constructs.md @@ -30,7 +30,8 @@ Proposed. 2. **A branch selects a task, not a string.** A case is the task reference the SDK already handed back, and the value on the wire is that task's task_id, so the compiler checks the candidate exists and no label has to be kept in step with it. 3. **No default case.** `BranchPythonOperator` has none to serialize, and a one-sided `If` whose condition is false follows nothing at all. A decider returning a ref that is not one of the declared cases is a run-time error the SDK raises, a narrower check than `skip_all_except`, which only rejects a task_id missing from the whole Dag. -4. **Triggering a Dag run is an ordinary DSL task**, it is pure DSL purpose instead of a new runtime. +4. **Triggering a Dag run is an ordinary task that the SDK's own runtime runs.** It has no handler. + The runtime sends the `TriggerDagRun` request itself and follows `TriggerDagRunOperator` for every option it offers, including `wait_for_completion` and `deferrable`. 5. **Grouping keeps Python's semantics**: a scope offering the same task and nesting methods as the Dag, prefixing each task_id with the group id (`prefix_group_id`), and can be ordered against a task or another group, as `group1 >> group2` does in Python. @@ -89,7 +90,12 @@ That operator defaults to skipping every task in its downstream closure and igno a single-ref return cannot express that, and no Lang SDK offers it for now. An author needing several paths together puts them behind one task, or gates each with its own condition. This limitation is accepted rather than open. -- **No Lang SDK needs a deferral mechanism for now** to offer `deferrable` or `wait_for_completion`, because the trigger task runs in Python. +- **Each Lang SDK runtime implements the trigger task**, so a native Dag needs no Python worker. + `deferrable` needs no trigger written in that language: the runtime defers to the standard provider's `DagStateTrigger`, which runs in the Python triggerer, and the task keeps its queue, so it resumes in the same runtime with `next_method` set. +- **The runtime reads the config the trigger task needs from its environment**, since it cannot read `airflow.cfg`. + It reads `AIRFLOW__API__BASE_URL`, `AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE` and `AIRFLOW__TRIGGERER__QUEUES_ENABLED`, and an unset one takes Python's fallback. + The coordinator exports them from Airflow's config (#73842). +- **The trigger's arguments are not Jinja-rendered**, because no Python code runs the task. A value is sent as the author wrote it. - **Nothing extra reaches the Dag JSON**, which carries no branch-candidate field at all. A ref is a task_id by the time the decision is sent, so each SDK stores its cases and nothing else. - **A group edge needs one base type per SDK** that both a task and a group satisfy, since either can sit at the end of an edge. Python already has it: `TaskGroup(TaskGroupMixin, DAGNode)` (`task-sdk/src/airflow/sdk/definitions/taskgroup.py:96`) and every operator inherit `DependencyMixin` (`.../definitions/_internal/mixins.py:35`), where `set_upstream` and `set_downstream` live. 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 99f568b679e1d..a3589d6019ac2 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -443,6 +443,50 @@ Exactly one case is selected. Python's branch callable may return a list of task SDK offers that yet: put the paths that run together behind one task, or gate each with its own condition. +Triggering another Dag +~~~~~~~~~~~~~~~~~~~~~~ + +``triggerDagRun`` starts another Dag's run. Pass it to ``dag.task`` in place of a handler: + +.. code-block:: typescript + + import { triggerDagRun } from "apache-airflow-ts-sdk"; + + const trigger = dag.task(triggerDagRun({ dagId: "downstream_etl", waitForCompletion: true }), { + taskId: "trigger_downstream", + })(); + + trigger.after(loaded); + +The task has no handler to take an id from, so ``taskId`` is required, and the factory takes no +inputs. The trailing spec carries the task's other options, as for any task. + +The task runs in the TypeScript runtime, like the Dag's other tasks, and inherits the Dag's queue. It +behaves as ``TriggerDagRunOperator`` does. ``waitForCompletion`` polls the run every ``pokeInterval`` +seconds until it reaches one of ``allowedStates`` or ``failedStates``. With ``deferrable`` as well, the +task defers to ``DagStateTrigger`` instead of holding a worker slot, and resumes in this runtime when +the run finishes. ``DagStateTrigger`` runs in the Python triggerer, so the triggerer needs the standard +provider installed. + +The task pushes the ``trigger_run_id`` XCom, and the task's "Triggered DAG" link opens the run it +started. + +The task follows these config options, with Python's fallback when one is unset: + +- ``[api] base_url`` is the base of the "Triggered DAG" link. It falls back to ``/``. +- ``[operators] default_deferrable`` is the default of ``deferrable``. It falls back to ``false``. +- ``[triggerer] queues_enabled`` decides whether the deferred trigger gets the task's queue. It falls + back to ``false``, so the trigger gets none. + +Some of what ``TriggerDagRunOperator`` does is not offered: + +- ``logical_date`` and ``run_after``. The triggered run's logical date is the time the task triggers + it, as in Python when neither is set. +- Jinja. Values are sent as written, so ``{{ ds }}`` in ``conf`` reaches the triggered run as that + literal string. +- OpenLineage parent injection (``openlineage_inject_parent_info``). The runtime does not add the + parent task's OpenLineage details to ``conf``. + ``new Dag`` and ``dag.task`` both take a trailing spec of Airflow options: ``{ schedule: "@daily", tags: ["etl"] }`` for the Dag, ``{ retries: 2, retryDelay: 30 }`` for a task. diff --git a/ts-sdk/src/coordinator/client.ts b/ts-sdk/src/coordinator/client.ts index 9103d824d58de..8f47955dcdf86 100644 --- a/ts-sdk/src/coordinator/client.ts +++ b/ts-sdk/src/coordinator/client.ts @@ -31,6 +31,9 @@ import type { SetXCom, GetConnection, SkipDownstreamTasks, + TriggerDagRun, + GetDagRunState, + GetDag, ConnectionResult as WireConnectionResult, } from "./protocol.js"; @@ -86,6 +89,16 @@ export interface CoordinatorClient extends TaskClient { /** Mark direct downstream tasks of the running task as skipped; none is a no-op. */ skipDownstreamTasks(taskIds: readonly string[]): Promise; + + /** + * Trigger a Dag run. `"already_exists"` is the one refusal a trigger task + * handles itself; any other error throws. + */ + triggerDagRun(msg: Omit): Promise<"triggered" | "already_exists">; + /** The state of a Dag run, as `GetDagRunState` reports it. */ + getDagRunState(dagId: string, runId: string): Promise; + /** Whether a Dag is paused. */ + isDagPaused(dagId: string): Promise; } export function createCoordinatorClient( @@ -99,11 +112,16 @@ export function createCoordinatorClient( request: unknown, extract: (body: Record | null) => T, absent: AbsentRowPolicy, + allowedError?: string, ): Promise { logs?.debug(`${op} request`); const frame = await comm.request(request); const err = parseFrameError(frame); if (err) { + if (err.code === allowedError) { + logs?.debug(`${op} answered ${err.code}`); + return null; + } if (absent === "null" && isNotFound(err)) { logs?.debug(`${op} not found`, { error: err.code }); return null; @@ -214,6 +232,45 @@ export function createCoordinatorClient( await rpc("SkipDownstreamTasks", null, msg, () => undefined, "throw"); }, + // ---- Dag runs ---- + + async triggerDagRun(msg: Omit) { + const request: TriggerDagRun = { type: "TriggerDagRun", ...msg }; + const triggered = await rpc( + "TriggerDagRun", + null, + request, + () => "triggered" as const, + "throw", + "DAGRUN_ALREADY_EXISTS", + ); + return triggered ?? "already_exists"; + }, + + async getDagRunState(dagId: string, runId: string): Promise { + const msg: GetDagRunState = { type: "GetDagRunState", dag_id: dagId, run_id: runId }; + const state = await rpc( + "GetDagRunState", + "DagRunStateResult", + msg, + (body) => body!.state as string, + "throw", + ); + return state!; + }, + + async isDagPaused(dagId: string): Promise { + const msg: GetDag = { type: "GetDag", dag_id: dagId }; + const paused = await rpc( + "GetDag", + "DagResult", + msg, + (body) => body!.is_paused === true, + "throw", + ); + return paused!; + }, + // ---- Connections ---- async getConnection(connId: string): Promise { diff --git a/ts-sdk/src/coordinator/protocol.ts b/ts-sdk/src/coordinator/protocol.ts index 1c95fee6f907d..9f2f97b949c3d 100644 --- a/ts-sdk/src/coordinator/protocol.ts +++ b/ts-sdk/src/coordinator/protocol.ts @@ -28,6 +28,7 @@ import type { DagFileParseRequest, ErrorResponse, DagFileParsingResult, + DeferTask, RetryTask, SucceedTask, TaskState, @@ -47,6 +48,7 @@ export type { export type { DagFileParsingResult, + DeferTask, RetryTask, SucceedTask, TaskState, @@ -65,6 +67,11 @@ export type { SetXCom, GetConnection, SkipDownstreamTasks, + TriggerDagRun, + GetDagRunState, + DagRunStateResult, + GetDag, + DagResult, } from "../generated/supervisor.js"; // -------- Frames from supervisor -------- @@ -82,6 +89,7 @@ export type MsgFromSupervisor = export type RuntimeTaskState = WithRequiredType; export type RuntimeRetryTask = WithRequiredType; +export type RuntimeDeferTask = WithRequiredType; export type RuntimeSucceedTask = WithRequiredType & { task_outlets: NonNullable; outlet_events: NonNullable; diff --git a/ts-sdk/src/coordinator/runtime.ts b/ts-sdk/src/coordinator/runtime.ts index fe2b027105a94..8c7227f8e70e0 100644 --- a/ts-sdk/src/coordinator/runtime.ts +++ b/ts-sdk/src/coordinator/runtime.ts @@ -49,16 +49,18 @@ import { asMsgFromSupervisor, SUPERVISOR_API_VERSION, type RuntimeDagFileParsingResult, + type RuntimeDeferTask, type RuntimeRetryTask, type RuntimeSucceedTask, type RuntimeTaskState, type StartupDetails, } from "./protocol.js"; import { getArgNames } from "../sdk/arg-names.js"; -import { bundleDags, bundleDagTaskIds, type Bundle } from "../sdk/bundle.js"; +import { bundleDags, bundleDagTaskIds, getBundleTrigger, type Bundle } from "../sdk/bundle.js"; import { finalizeDag } from "../sdk/dag.js"; import { SERIALIZATION_VERSION } from "../generated/dag-schema-fields.js"; import { computeRelativeFileloc, serializeDag } from "./serde.js"; +import { runTriggerDagRun } from "./trigger-runner.js"; import { runInTaskScope, type TaskContext } from "../sdk/task.js"; import type { JsonValue } from "../sdk/client-types.js"; @@ -195,6 +197,8 @@ export async function startCoordinator( await sendSupervisorResponse(firstFrame.id, response, comm, runtimeLogs); if (response.type === "SucceedTask") { runtimeLogs.info("Task succeeded", { task_id: body.ti.task_id }); + } else if (response.type === "DeferTask") { + runtimeLogs.info("Task deferred", { task_id: body.ti.task_id }); } } else { const errMsg = `First frame must be DagFileParseRequest or StartupDetails, got ${body.type}`; @@ -338,8 +342,22 @@ async function handleTask( logs: LogChannel, clientLogs: LogChannel, signal: AbortSignal, -): Promise { +): Promise { const ti = details.ti; + const trigger = getBundleTrigger(bundle, ti.dag_id, ti.task_id); + if (trigger) { + const ctx = buildContext(details, signal); + const client = createCoordinatorClient(comm, ctx, clientLogs); + const fail = (message: string) => { + logs.error("Task failed", { task_id: ctx.taskId, error: message }); + return buildFailureResponse(details, message); + }; + try { + return await runTriggerDagRun(details, trigger, client, logs, ctx.signal, fail); + } catch (err) { + return fail((err as Error).message ?? String(err)); + } + } const handler = bundle.getTaskHandler(ti.dag_id, ti.task_id); if (!handler) { diff --git a/ts-sdk/src/coordinator/serde.ts b/ts-sdk/src/coordinator/serde.ts index a26bf11837cf5..f52320ae59c80 100644 --- a/ts-sdk/src/coordinator/serde.ts +++ b/ts-sdk/src/coordinator/serde.ts @@ -50,6 +50,7 @@ import { type Dag, type RecordedInputs, type TaskGroupRecord, + type TaskRecord, } from "../sdk/dag.js"; /** A serialized Dag: JSON, by the time it reaches the supervisor as msgpack. */ @@ -83,6 +84,29 @@ const TASK_MODULE = "airflow.sdk.coordinators.node"; */ const TASK_LANGUAGE = "typescript"; +/** The Dags this one triggers, for the UI dependency graph. */ +function serializeDagDependencies(dag: Dag): SerializedValue { + const dependencies: SerializedValue[] = []; + for (const [taskId, record] of getDagTaskRecords(dag)) { + if (!record.trigger) continue; + dependencies.push({ + source: dag.dagId, + target: record.trigger.dagId, + label: taskId, + dependency_type: "trigger", + dependency_id: taskId, + }); + } + return dependencies; +} + +/** Makes the UI draw a trigger task as `TriggerDagRunOperator` with its link. */ +const TRIGGER_DAG_RUN_FIELDS: Readonly> = { + _operator_name: "TriggerDagRunOperator", + ui_color: "#ffefeb", + _operator_extra_links: { "Triggered DAG": "_link_TriggerDagRunLink" }, +}; + /** How one set of authoring fields is written into a serialized object. */ interface FieldRules { readonly fields: Readonly>; @@ -144,13 +168,13 @@ export function serializeDag( serializeTask( dag.dagId, taskId, - withDagQueue(record.spec, dag.spec.queue), + record, graph.downstreamTaskIds.get(taskId), inputs.get(taskId), - record.canSkipDownstream === true, + dag.spec.queue, ), ), - dag_dependencies: [], + dag_dependencies: serializeDagDependencies(dag), task_group: serializeTaskGroups(dag, graph), edge_info: {}, params: [], @@ -182,10 +206,10 @@ function withDagQueue(spec: object, dagQueue: string | undefined): object { function serializeTask( dagId: string, taskId: string, - spec: object, + record: TaskRecord, downstream: ReadonlySet | undefined, inputs: RecordedInputs | undefined, - canSkipDownstream: boolean, + dagQueue: string | undefined, ): SerializedValue { const data: Record = { task_id: taskId, @@ -202,11 +226,17 @@ function serializeTask( is_stub: true, }; const label = `task "${taskId}" of Dag "${dagId}"`; + if (record.trigger) { + if (record.trigger.conf !== undefined) { + toPlainJson(record.trigger.conf, `conf of ${label}`); + } + Object.assign(data, structuredClone(TRIGGER_DAG_RUN_FIELDS)); + } const bindings = serializeArgBindings(inputs, label); if (bindings) data["_arg_bindings"] = bindings; // Lets `NotPreviouslySkippedDep` re-skip a cleared downstream, as Python's SkipMixin does. - if (canSkipDownstream) data["_can_skip_downstream"] = true; - applySchemaFields(data, spec, TASK_FIELD_RULES, label); + if (record.canSkipDownstream) data["_can_skip_downstream"] = true; + applySchemaFields(data, withDagQueue(record.spec, dagQueue), TASK_FIELD_RULES, label); if (downstream?.size) { data["downstream_task_ids"] = [...downstream].sort(); } diff --git a/ts-sdk/src/coordinator/trigger-runner.ts b/ts-sdk/src/coordinator/trigger-runner.ts new file mode 100644 index 0000000000000..5ead546914136 --- /dev/null +++ b/ts-sdk/src/coordinator/trigger-runner.ts @@ -0,0 +1,232 @@ +/*! + * 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. + */ + +// Mirrors `TriggerDagRunOperator.execute` and `execute_complete` in the standard provider. + +import { setTimeout as sleep } from "node:timers/promises"; +import type { CoordinatorClient } from "./client.js"; +import type { LogChannel } from "./log-channel.js"; +import type { + RuntimeDeferTask, + RuntimeRetryTask, + RuntimeSucceedTask, + RuntimeTaskState, + StartupDetails, +} from "./protocol.js"; +import { isPlainRecord } from "../sdk/dag.js"; +import { getBooleanEnv, type TriggerDagRunTask } from "../sdk/trigger-dag-run.js"; + +export type TriggerOutcome = + RuntimeSucceedTask | RuntimeRetryTask | RuntimeTaskState | RuntimeDeferTask; + +/** A failure the task's retries apply to, as `_handle_current_task_failed` decides. */ +export type FailTask = (message: string) => RuntimeRetryTask | RuntimeTaskState; + +const DAG_STATE_TRIGGER = "airflow.providers.standard.triggers.external_task.DagStateTrigger"; +/** `TriggerDagRunLink().xcom_key`: the XCom the "Triggered DAG" extra link reads. */ +const LINK_XCOM_KEY = "_link_TriggerDagRunLink"; +const RUN_ID_XCOM_KEY = "trigger_run_id"; +/** `TRIGGER_FAIL_REPR`: the `next_method` a failed or timed-out trigger resumes with. */ +const TRIGGER_FAIL = "__fail__"; +const EXECUTE_COMPLETE = "execute_complete"; + +export async function runTriggerDagRun( + details: StartupDetails, + trigger: TriggerDagRunTask, + client: CoordinatorClient, + logs: LogChannel, + signal: AbortSignal, + fail: FailTask, +): Promise { + const nextMethod = details.ti_context.next_method; + if (nextMethod) return resume(nextMethod, details.ti_context.next_kwargs, trigger, logs, fail); + + const logicalDate = new Date(); + const runId = trigger.runId ?? `manual__${pythonIsoformat(logicalDate)}`; + + if (trigger.failWhenDagIsPaused && (await client.isDagPaused(trigger.dagId))) { + return fail(`Dag ${trigger.dagId} is paused`); + } + + logs.info("Triggering Dag Run.", { trigger_dag_id: trigger.dagId }); + const [, triggered] = await Promise.all([ + client.setXCom({ key: LINK_XCOM_KEY, value: dagRunUrl(trigger.dagId, runId) }), + client.triggerDagRun({ + dag_id: trigger.dagId, + run_id: runId, + logical_date: logicalDate.toISOString(), + run_after: null, + conf: (trigger.conf as Record | undefined) ?? null, + reset_dag_run: trigger.resetDagRun, + note: trigger.note ?? null, + }), + ]); + if (triggered === "already_exists") { + if (trigger.skipWhenAlreadyExists) { + logs.info( + "Dag Run already exists, skipping task as skip_when_already_exists is set to True.", + { dag_id: trigger.dagId }, + ); + return { type: "TaskState", state: "skipped", end_date: new Date().toISOString() }; + } + logs.error("Dag Run already exists, marking task as failed.", { dag_id: trigger.dagId }); + return { type: "TaskState", state: "failed", end_date: new Date().toISOString() }; + } + logs.info("Dag Run triggered successfully.", { trigger_dag_id: trigger.dagId }); + await client.setXCom({ key: RUN_ID_XCOM_KEY, value: runId }); + + if (!trigger.waitForCompletion) { + if (trigger.deferrable) { + logs.info( + "Ignoring deferrable=True because wait_for_completion=False. " + + "Task will complete immediately without waiting for the triggered DAG run.", + { trigger_dag_id: trigger.dagId }, + ); + } + return succeeded(); + } + + if (trigger.deferrable) { + logs.info("Pausing task as DEFERRED.", { trigger_dag_id: trigger.dagId, run_id: runId }); + return { + type: "DeferTask", + state: "deferred", + classpath: DAG_STATE_TRIGGER, + // `DagStateTrigger.serialize()`, key for key. + trigger_kwargs: { + dag_id: trigger.dagId, + states: [...trigger.allowedStates, ...trigger.failedStates], + poll_interval: trigger.pokeInterval, + run_ids: [runId], + execution_dates: null, + }, + trigger_timeout: null, + // `_defer_task` hands the trigger the task's queue only when triggerer queues are enabled. + queue: getBooleanEnv("AIRFLOW__TRIGGERER__QUEUES_ENABLED", false) + ? (details.ti.queue ?? null) + : null, + next_method: EXECUTE_COMPLETE, + next_kwargs: {}, + }; + } + + while (true) { + logs.info("Waiting for dag run to complete execution in allowed state.", { + dag_id: trigger.dagId, + run_id: runId, + allowed_state: trigger.allowedStates, + }); + await sleep(trigger.pokeInterval * 1000, undefined, { signal }); + const state = await client.getDagRunState(trigger.dagId, runId); + if (includes(trigger.failedStates, state)) { + logs.error("DagRun finished with failed state.", { dag_id: trigger.dagId, state }); + return fail(`${trigger.dagId} failed with failed state ${state}`); + } + if (includes(trigger.allowedStates, state)) { + logs.info("DagRun finished with allowed state.", { dag_id: trigger.dagId, state }); + return succeeded(); + } + logs.debug("DagRun not yet in allowed or failed state.", { dag_id: trigger.dagId, state }); + } +} + +/** `BaseOperator.resume_execution` for this task: `__fail__` or `execute_complete`. */ +function resume( + nextMethod: string, + nextKwargs: unknown, + trigger: TriggerDagRunTask, + logs: LogChannel, + fail: FailTask, +): TriggerOutcome { + const kwargs = isPlainRecord(nextKwargs) ? nextKwargs : {}; + if (nextMethod === TRIGGER_FAIL) { + const traceback = kwargs["traceback"]; + if (Array.isArray(traceback)) logs.error(`Trigger failed:\n${traceback.join("\n")}`); + return fail(String(kwargs["error"] ?? "Unknown")); + } + if (nextMethod !== EXECUTE_COMPLETE) { + return fail(`Task cannot resume with next_method "${nextMethod}"`); + } + const eventData = decodeEvent(kwargs["event"]); + const runIds = eventData?.["run_ids"]; + if (eventData === undefined || !Array.isArray(runIds)) { + return fail(`Task resumed with an event it cannot read: ${JSON.stringify(kwargs["event"])}`); + } + const failedRunIds: string[] = []; + for (const runId of runIds) { + const state = eventData[String(runId)]; + if (includes(trigger.failedStates, state)) { + failedRunIds.push(String(runId)); + continue; + } + if (includes(trigger.allowedStates, state)) { + logs.info("Triggered Dag run finished with allowed state.", { + dag_id: trigger.dagId, + state, + run_id: runId, + }); + } + } + if (failedRunIds.length > 0) { + return fail( + `${trigger.dagId} failed with failed states ${JSON.stringify(trigger.failedStates)} ` + + `for run_ids ${JSON.stringify(failedRunIds)}`, + ); + } + return succeeded(); +} + +/** + * The payload `DagStateTrigger` fired with: `(classpath, data)`. The triggerer + * stores it with serde, which encodes a tuple as `{__classname__, __data__}`. + */ +function decodeEvent(event: unknown): Record | undefined { + const pair = + isPlainRecord(event) && event["__classname__"] === "builtins.tuple" ? event["__data__"] : event; + if (!Array.isArray(pair) || pair.length !== 2 || !isPlainRecord(pair[1])) return undefined; + return pair[1]; +} + +function succeeded(): RuntimeSucceedTask { + return { + type: "SucceedTask", + end_date: new Date().toISOString(), + task_outlets: [], + outlet_events: [], + }; +} + +/** `datetime.isoformat()` of a UTC instant, which is how Python spells a run ID's date. */ +export function pythonIsoformat(date: Date): string { + const iso = date.toISOString(); + const seconds = iso.slice(0, 19); + const millis = date.getUTCMilliseconds(); + const fraction = millis === 0 ? "" : `.${String(millis * 1000).padStart(6, "0")}`; + return `${seconds}${fraction}+00:00`; +} + +/** `build_airflow_dagrun_url`, on `[api] base_url` from the environment, or "/" when unset. */ +export function dagRunUrl(dagId: string, runId: string): string { + const base = process.env["AIRFLOW__API__BASE_URL"] || "/"; + return `${base.replace(/\/+$/, "")}/dags/${dagId}/runs/${runId}`; +} + +function includes(states: readonly string[], state: unknown): boolean { + return typeof state === "string" && states.includes(state); +} diff --git a/ts-sdk/src/index.ts b/ts-sdk/src/index.ts index 8c0f90eb81f67..f71bbbf9ac0a8 100644 --- a/ts-sdk/src/index.ts +++ b/ts-sdk/src/index.ts @@ -18,6 +18,7 @@ */ export { Dag } from "./sdk/dag.js"; +export { triggerDagRun } from "./sdk/trigger-dag-run.js"; export { Bundle } from "./sdk/bundle.js"; export { TaskHandler } from "./sdk/task-handler.js"; export { withArgNames } from "./sdk/arg-names.js"; @@ -26,6 +27,7 @@ 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 { DagRunState, TriggerDagRunSpec, TriggerDagRunTask } from "./sdk/trigger-dag-run.js"; export type { Branch, Condition, diff --git a/ts-sdk/src/sdk/bundle.ts b/ts-sdk/src/sdk/bundle.ts index a3e54f1a23172..d4e5e2f616675 100644 --- a/ts-sdk/src/sdk/bundle.ts +++ b/ts-sdk/src/sdk/bundle.ts @@ -23,6 +23,7 @@ import { brand, DUPLICATE_COPY_HINT, hasBrand } from "./brand.js"; import { Dag, finalizeDag, getDagTaskRecords, isDag } from "./dag.js"; import { getTaskHandlerFunction, isTaskHandler, TaskHandler } from "./task-handler.js"; import type { TaskFunction } from "./task.js"; +import type { TriggerDagRunTask } from "./trigger-dag-run.js"; // Assigned inside Bundle's static block, as Dag does for its tasks. let dagsOf: (bundle: Bundle) => ReadonlyMap; @@ -211,6 +212,17 @@ export function validateOwnBundle(value: unknown, accessor: string): asserts val ); } +/** Internal: the `triggerDagRun(...)` task at `taskId` of a Dag this bundle + * declared in TypeScript, which the runtime runs without a handler. */ +export function getBundleTrigger( + bundle: Bundle, + dagId: string, + taskId: string, +): TriggerDagRunTask | undefined { + const dag = dagsOf(bundle).get(dagId); + return dag === undefined ? undefined : getDagTaskRecords(dag).get(taskId)?.trigger; +} + /** Internal: finalize every Dag this bundle declared in TypeScript, so no task * can be added or wired afterwards. */ export function finalizeBundleDags(bundle: Bundle): void { diff --git a/ts-sdk/src/sdk/dag.ts b/ts-sdk/src/sdk/dag.ts index c285a4c9e887c..eb8f32132a1d9 100644 --- a/ts-sdk/src/sdk/dag.ts +++ b/ts-sdk/src/sdk/dag.ts @@ -32,6 +32,7 @@ import { import { brand, DUPLICATE_COPY_HINT, hasBrand } from "./brand.js"; import type { JsonValue } from "./client-types.js"; import { getCurrentModuleSource } from "./module-source.js"; +import { isTriggerDagRunTask, type TriggerDagRunTask } from "./trigger-dag-run.js"; import { getClient, type TaskFunction } from "./task.js"; /** Internal: whether `value` is an object literal, not an array or a class instance. */ @@ -228,6 +229,9 @@ export interface TaskGroupRef extends Node { handler: (args: TArgs) => TReturn | Promise, options?: TaskOptions, ): TaskFactory; + /** Declare a `triggerDagRun(...)` task in this group, prefixed the same way. */ + task(taskId: string, trigger: TriggerDagRunTask, options?: TaskOptions): () => TaskRef; + task(trigger: TriggerDagRunTask, options: TaskOptions): () => TaskRef; /** Nest a group inside this one. */ taskGroup(groupId: string, options?: TaskGroupOptions): TaskGroupRef; before(...downstream: readonly Node[]): TaskGroupRef; @@ -448,14 +452,16 @@ export interface Branch extends Node { after(...upstream: readonly Node[]): Branch; } -/** Per-task record a Dag retains: the reference, the handler, and its spec. */ -export interface TaskRecord { +/** Per-task record a Dag retains: the reference, what runs it, and its spec. */ +export type TaskRecord = { readonly task: TaskRef; - readonly fn: TaskFunction; readonly spec: TaskSpec; /** Whether this task decides which of its downstream tasks to skip. */ readonly canSkipDownstream?: boolean; -} +} & ( + | { readonly fn: TaskFunction; readonly trigger?: never } + | { readonly fn?: never; readonly trigger: TriggerDagRunTask } +); interface ConditionRecord { whenTrue?: TaskRef; @@ -592,9 +598,16 @@ export class Dag { handler: (args: TArgs) => TReturn | Promise, options?: TaskOptions, ): TaskFactory; + /** + * Declare a task that triggers another Dag's run, from `triggerDagRun(spec)`. + * It has no handler to take an id from, so the id is required. + */ + task(taskId: string, trigger: TriggerDagRunTask, options?: TaskOptions): () => TaskRef; + task(trigger: TriggerDagRunTask, options: TaskOptions): () => TaskRef; task( - taskIdOrHandler: string | ((args: TArgs) => TReturn | Promise), - handlerOrOptions?: ((args: TArgs) => TReturn | Promise) | TaskOptions, + taskIdOrHandler: string | ((args: TArgs) => TReturn | Promise) | TriggerDagRunTask, + handlerOrOptions?: + ((args: TArgs) => TReturn | Promise) | TriggerDagRunTask | TaskOptions, maybeOptions?: TaskOptions, ): TaskFactory { return this.#addTask(undefined, taskIdOrHandler, handlerOrOptions, maybeOptions); @@ -691,8 +704,9 @@ export class Dag { taskId: string, decide: (returned: unknown) => Promise<{ skip: string[]; result: unknown }>, ): void { - const record = this.#tasks.get(taskId)!; - const inner = record.fn; + const { task, spec, fn } = this.#tasks.get(taskId)!; + // A decider is always declared from a handler, never from triggerDagRun. + const inner = fn!; const wrapped: TaskFunction = async (args) => { const { skip, result } = await decide(await inner(args as never)); if (skip.length > 0) { @@ -703,7 +717,7 @@ export class Dag { } return result; }; - this.#tasks.set(taskId, { ...record, canSkipDownstream: true, fn: wrapped }); + this.#tasks.set(taskId, { task, spec, canSkipDownstream: true, fn: wrapped }); } /** @@ -794,14 +808,15 @@ export class Dag { #addTask( groupId: string | undefined, - taskIdOrHandler: string | ((args: TArgs) => TReturn | Promise), - handlerOrOptions?: ((args: TArgs) => TReturn | Promise) | TaskOptions, + taskIdOrHandler: string | ((args: TArgs) => TReturn | Promise) | TriggerDagRunTask, + handlerOrOptions?: + ((args: TArgs) => TReturn | Promise) | TriggerDagRunTask | TaskOptions, maybeOptions?: TaskOptions, ): TaskFactory { const idGiven = typeof taskIdOrHandler === "string"; - const handler = (idGiven ? handlerOrOptions : taskIdOrHandler) as ( - args: TArgs, - ) => TReturn | Promise; + const body = idGiven ? handlerOrOptions : taskIdOrHandler; + const trigger = isTriggerDagRunTask(body) ? body : undefined; + const handler = body as (args: TArgs) => TReturn | Promise; const given = idGiven ? maybeOptions : (handlerOrOptions as TaskOptions | undefined); // Defaulted only when absent: an explicit `null` is a bad spec, not an // omitted one, and #taskSpecOf is what reports it. @@ -812,12 +827,18 @@ export class Dag { // wins and the spec's is dropped without a word. if (idGiven && specTaskId !== undefined) { throw new Error( - `Task "${taskIdOrHandler}" of Dag "${this.dagId}" also carries taskId "${specTaskId}" in ` + - "its spec; give the id once, either positionally or in the spec", + `Task "${taskIdOrHandler as string}" of Dag "${this.dagId}" also carries taskId ` + + `"${specTaskId}" in its spec; give the id once, either positionally or in the spec`, ); } const defaulted = idGiven ? undefined : (specTaskId ?? readFunctionName(handler)); - const declared = idGiven ? taskIdOrHandler : defaulted; + const declared = idGiven ? (taskIdOrHandler as string) : defaulted; + if (declared === undefined && trigger !== undefined) { + throw new Error( + `A triggerDagRun task of Dag "${this.dagId}" has no taskId; name it with ` + + 'dag.task(triggerDagRun({ dagId: "..." }), { taskId: "trigger_downstream" })', + ); + } if (declared === undefined) { throw new Error( `A task of Dag "${this.dagId}" has no id: its handler has no name to take one from. ` + @@ -844,12 +865,12 @@ export class Dag { ); } const taskId = this.#childId(groupId, declared); - if (typeof handler !== "function") { + if (trigger === undefined && typeof handler !== "function") { throw new Error(`handler for Dag "${this.dagId}" task "${taskId}" must be a function`); } // TypeScript already says so, but a plain-JavaScript author lands here // with the argument list Python would take. - if (handler.length > 1) { + if (trigger === undefined && handler.length > 1) { throw new Error( `Handler for Dag "${this.dagId}" task "${taskId}" declares ${handler.length} parameters; ` + "a handler takes one object of named arguments — async ({ rows, region }) => ...", @@ -871,7 +892,7 @@ export class Dag { task, // The runtime dispatches every handler through one instantiation, as it // does a registered TaskHandler. - fn: handler as unknown as TaskFunction, + ...(trigger ? { trigger } : { fn: handler as unknown as TaskFunction }), spec: freezeSpec(spec, () => `The spec for Dag "${this.dagId}" task "${taskId}"`), }); return ((...inputs: unknown[]) => { @@ -883,6 +904,11 @@ export class Dag { "it takes one object naming its inputs: myTask({ rows, region })", ); } + if (trigger !== undefined && inputs[0] !== undefined) { + throw new Error( + `Task "${taskId}" of Dag "${this.dagId}" is a triggerDagRun task, which takes no inputs`, + ); + } this.#wire(taskId, inputs[0]); return task; }) as TaskFactory; @@ -1004,8 +1030,9 @@ export class Dag { dagId: this.dagId, groupId, task: ( - taskIdOrHandler: string | ((args: TArgs) => TReturn | Promise), - handlerOrOptions?: ((args: TArgs) => TReturn | Promise) | TaskOptions, + taskIdOrHandler: string | ((args: TArgs) => TReturn | Promise) | TriggerDagRunTask, + handlerOrOptions?: + ((args: TArgs) => TReturn | Promise) | TriggerDagRunTask | TaskOptions, maybeOptions?: TaskOptions, ) => this.#addTask(groupId, taskIdOrHandler, handlerOrOptions, maybeOptions), taskGroup: (childId: string, options?: TaskGroupOptions) => diff --git a/ts-sdk/src/sdk/trigger-dag-run.ts b/ts-sdk/src/sdk/trigger-dag-run.ts new file mode 100644 index 0000000000000..bba79dcb26022 --- /dev/null +++ b/ts-sdk/src/sdk/trigger-dag-run.ts @@ -0,0 +1,192 @@ +/*! + * 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 { brand, hasBrand } from "./brand.js"; +import { isPlainRecord } from "./dag.js"; +import type { JsonValue } from "./client-types.js"; + +/** A Dag run state, as `allowedStates` and `failedStates` name one. */ +export type DagRunState = "queued" | "running" | "success" | "failed"; + +const DAG_RUN_STATES: ReadonlySet = new Set([ + "queued", + "running", + "success", + "failed", +]); + +/** The `TriggerDagRunOperator` options, named as TypeScript spells them. */ +export interface TriggerDagRunSpec { + /** Identifier of the Dag to trigger. */ + readonly dagId: string; + /** Run ID for the triggered run; generated from the trigger time when unset. */ + readonly runId?: string; + /** Configuration the triggered run is started with. */ + readonly conf?: Readonly>; + /** Clear an existing run with the same ID instead of failing. */ + readonly resetDagRun?: boolean; + /** Hold this task open until the triggered run finishes. */ + readonly waitForCompletion?: boolean; + /** Seconds between checks while waiting. Defaults to 60. */ + readonly pokeInterval?: number; + /** Run states that count as success when waiting. Defaults to `["success"]`. */ + readonly allowedStates?: readonly DagRunState[]; + /** Run states that count as failure when waiting. Defaults to `["failed"]`. */ + readonly failedStates?: readonly DagRunState[]; + /** Skip rather than fail when the run already exists. */ + readonly skipWhenAlreadyExists?: boolean; + /** Fail rather than trigger when the target Dag is paused. */ + readonly failWhenDagIsPaused?: boolean; + /** Note recorded against the triggered run. */ + readonly note?: string; + /** + * While waiting, free the worker slot and defer to `DagStateTrigger`, which + * the Python triggerer runs. Defaults to `[operators] default_deferrable`, or + * to false when that is unset. + */ + readonly deferrable?: boolean; +} + +/** Internal: a trigger task's options with `TriggerDagRunOperator`'s defaults applied. */ +export interface TriggerDagRunTask { + readonly dagId: string; + readonly runId: string | undefined; + readonly conf: Readonly> | undefined; + readonly resetDagRun: boolean; + readonly waitForCompletion: boolean; + readonly pokeInterval: number; + readonly allowedStates: readonly DagRunState[]; + readonly failedStates: readonly DagRunState[]; + readonly skipWhenAlreadyExists: boolean; + readonly failWhenDagIsPaused: boolean; + readonly note: string | undefined; + readonly deferrable: boolean; +} + +/** Internal: whether `value` is a trigger task built by any copy of this package. */ +export function isTriggerDagRunTask(value: unknown): value is TriggerDagRunTask { + return hasBrand(value, "TriggerDagRunTask"); +} + +const OPTION_NAMES: ReadonlySet = new Set([ + "dagId", + "runId", + "conf", + "resetDagRun", + "waitForCompletion", + "pokeInterval", + "allowedStates", + "failedStates", + "skipWhenAlreadyExists", + "failWhenDagIsPaused", + "note", + "deferrable", +]); + +function checkStates(name: string, states: unknown): DagRunState[] | undefined { + if (states === undefined) return undefined; + if (!Array.isArray(states)) { + throw new Error(`triggerDagRun(...) option "${name}" must be an array of Dag run states`); + } + for (const state of states) { + if (typeof state !== "string" || !DAG_RUN_STATES.has(state)) { + throw new Error( + `triggerDagRun(...) option "${name}" holds ${JSON.stringify(state)}, which is not a Dag ` + + `run state; use one of ${[...DAG_RUN_STATES].join(", ")}`, + ); + } + } + return [...(states as DagRunState[])]; +} + +function checkType(name: string, value: unknown, type: "string" | "boolean"): void { + if (value !== undefined && typeof value !== type) { + throw new Error(`triggerDagRun(...) option "${name}" must be a ${type}`); + } +} + +/** Internal: a boolean Airflow option from the environment, read as `conf.getboolean` does. */ +export function getBooleanEnv(name: string, fallback: boolean): boolean { + const raw = process.env[name]; + if (raw === undefined) return fallback; + const value = raw.trim().toLowerCase(); + if (value === "t" || value === "true" || value === "1") return true; + if (value === "f" || value === "false" || value === "0") return false; + throw new Error(`${name} is ${JSON.stringify(raw)}, which is not a boolean; use true or false`); +} + +/** + * A task that triggers another Dag's run, passed to `dag.task` in place of a handler. + * + * ```ts + * dag.task(triggerDagRun({ dagId: "downstream_etl" }), { taskId: "trigger_downstream" })(); + * ``` + */ +export function triggerDagRun(spec: TriggerDagRunSpec): TriggerDagRunTask { + const options: unknown = spec; + if (!isPlainRecord(options)) { + throw new Error("triggerDagRun(...) takes an options object"); + } + for (const name of Object.keys(options)) { + if (!OPTION_NAMES.has(name)) throw new Error(`Unknown option "${name}" for triggerDagRun(...)`); + } + if (typeof options["dagId"] !== "string" || options["dagId"].length === 0) { + throw new Error("triggerDagRun(...) needs the dagId of the Dag to trigger"); + } + const conf = options["conf"]; + if (conf !== undefined && (typeof conf !== "object" || conf === null || Array.isArray(conf))) { + throw new Error('triggerDagRun(...) option "conf" must be an object'); + } + checkType("runId", options["runId"], "string"); + checkType("note", options["note"], "string"); + for (const flag of [ + "resetDagRun", + "waitForCompletion", + "skipWhenAlreadyExists", + "failWhenDagIsPaused", + "deferrable", + ]) { + checkType(flag, options[flag], "boolean"); + } + const pokeInterval = options["pokeInterval"] ?? 60; + if (typeof pokeInterval !== "number" || !Number.isFinite(pokeInterval) || pokeInterval < 0) { + throw new Error('triggerDagRun(...) option "pokeInterval" must be a non-negative number'); + } + const allowedStates = checkStates("allowedStates", options["allowedStates"]); + const failedStates = checkStates("failedStates", options["failedStates"]); + const task: TriggerDagRunTask = { + dagId: options["dagId"], + runId: (options["runId"] as string | undefined) || undefined, + conf: conf as TriggerDagRunTask["conf"], + resetDagRun: options["resetDagRun"] === true, + waitForCompletion: options["waitForCompletion"] === true, + pokeInterval, + // An empty `failedStates` is kept, as Python keeps it. + allowedStates: allowedStates?.length ? allowedStates : ["success"], + failedStates: failedStates ?? ["failed"], + skipWhenAlreadyExists: options["skipWhenAlreadyExists"] === true, + failWhenDagIsPaused: options["failWhenDagIsPaused"] === true, + note: (options["note"] as string | undefined) || undefined, + deferrable: + (options["deferrable"] as boolean | undefined) ?? + getBooleanEnv("AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE", false), + }; + brand(task, "TriggerDagRunTask"); + return Object.freeze(task); +} diff --git a/ts-sdk/tests/coordinator/integration.test.ts b/ts-sdk/tests/coordinator/integration.test.ts index 31d48d017194c..ee945c1c8aa64 100644 --- a/ts-sdk/tests/coordinator/integration.test.ts +++ b/ts-sdk/tests/coordinator/integration.test.ts @@ -36,6 +36,7 @@ import { startCoordinator, } from "../../src/coordinator/runtime.js"; import { Dag } from "../../src/sdk/dag.js"; +import { triggerDagRun } from "../../src/sdk/trigger-dag-run.js"; import { Bundle } from "../../src/sdk/bundle.js"; import { withArgNames } from "../../src/sdk/arg-names.js"; import { TaskHandler } from "../../src/sdk/task-handler.js"; @@ -942,4 +943,347 @@ describe("coordinator runtime integration", () => { const setXComReqs = result.runtimeRequests.filter((r) => r.type === "SetXCom"); expect(setXComReqs).toHaveLength(0); }); + + describe("triggerDagRun", () => { + const DAG_STATE_TRIGGER = "airflow.providers.standard.triggers.external_task.DagStateTrigger"; + const OK = { body: { type: "OKResponse", ok: true } }; + + function replies(queued: Record): Responder { + return (msgType) => queued[msgType]?.shift() ?? { body: null }; + } + + function requestsOf(result: MockResult, type: string) { + return result.runtimeRequests.filter((r) => r.type === type).map((r) => r.body); + } + + afterEach(() => { + vi.unstubAllEnvs(); + }); + + it("triggers the run, pushes the link and run id, and succeeds", async () => { + testDag.task( + triggerDagRun({ dagId: "downstream", conf: { source: "ts" }, note: "from ts" }), + { taskId: "trigger" }, + )(); + + const result = await driveSupervisor( + makeStartupDetails("trigger"), + replies({ TriggerDagRun: [OK] }), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" }); + expect(result.runtimeRequests.map((r) => r.type)).toEqual([ + "SetXCom", + "TriggerDagRun", + "SetXCom", + ]); + const [sent] = requestsOf(result, "TriggerDagRun"); + const runId = sent!["run_id"] as string; + expect(runId).toMatch(/^manual__\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d(\.\d{6})?\+00:00$/); + expect(new Date(sent!["logical_date"] as string).getTime()).toBe( + new Date(runId.slice("manual__".length)).getTime(), + ); + expect(sent).toMatchObject({ + dag_id: "downstream", + conf: { source: "ts" }, + reset_dag_run: false, + note: "from ts", + run_after: null, + }); + expect(requestsOf(result, "SetXCom")).toEqual([ + expect.objectContaining({ + key: "_link_TriggerDagRunLink", + value: `/dags/downstream/runs/${runId}`, + task_id: "trigger", + }), + expect.objectContaining({ key: "trigger_run_id", value: runId, task_id: "trigger" }), + ]); + }); + + it("builds the link on AIRFLOW__API__BASE_URL when the environment sets it", async () => { + vi.stubEnv("AIRFLOW__API__BASE_URL", "https://airflow.example.com/sub/"); + testDag.task(triggerDagRun({ dagId: "downstream", runId: "given" }), { taskId: "trigger" })(); + + const result = await driveSupervisor( + makeStartupDetails("trigger"), + replies({ TriggerDagRun: [OK] }), + ); + + expect(requestsOf(result, "TriggerDagRun")[0]).toMatchObject({ run_id: "given" }); + expect(requestsOf(result, "SetXCom")[0]).toMatchObject({ + value: "https://airflow.example.com/sub/dags/downstream/runs/given", + }); + }); + + it.each([ + [true, "skipped"], + [false, "failed"], + ])( + "with skipWhenAlreadyExists=%s, ends %s when the run exists, retries or not", + async (skip, state) => { + testDag.task(triggerDagRun({ dagId: "downstream", skipWhenAlreadyExists: skip }), { + taskId: "trigger", + })(); + + const result = await driveSupervisor( + makeStartupDetails("trigger", "test_dag", "r1", { should_retry: true }), + replies({ + TriggerDagRun: [{ body: { type: "ErrorResponse", error: "DAGRUN_ALREADY_EXISTS" } }], + }), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "TaskState", state }); + expect(requestsOf(result, "SetXCom").map((b) => b["key"])).toEqual([ + "_link_TriggerDagRunLink", + ]); + }, + ); + + it("fails like any task error when the trigger request itself fails", async () => { + testDag.task(triggerDagRun({ dagId: "missing" }), { taskId: "trigger" })(); + + const result = await driveSupervisor( + makeStartupDetails("trigger", "test_dag", "r1", { should_retry: true }), + replies({ + TriggerDagRun: [ + { + body: { + type: "ErrorResponse", + error: "API_SERVER_ERROR", + detail: { status_code: 404 }, + }, + }, + ], + }), + ); + + expect(result.firstResponse!.body).toMatchObject({ + type: "RetryTask", + retry_reason: "TriggerDagRun failed: API_SERVER_ERROR", + }); + }); + + it("refuses to trigger a paused Dag when failWhenDagIsPaused is set", async () => { + testDag.task(triggerDagRun({ dagId: "downstream", failWhenDagIsPaused: true }), { + taskId: "trigger", + })(); + + const result = await driveSupervisor( + makeStartupDetails("trigger"), + replies({ + GetDag: [{ body: { type: "DagResult", dag_id: "downstream", is_paused: true } }], + }), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "TaskState", state: "failed" }); + expect(requestsOf(result, "GetDag")).toEqual([{ type: "GetDag", dag_id: "downstream" }]); + expect(requestsOf(result, "TriggerDagRun")).toEqual([]); + }); + + it("polls the run's state until it reaches an allowed state", async () => { + testDag.task( + triggerDagRun({ dagId: "downstream", waitForCompletion: true, pokeInterval: 0 }), + { taskId: "trigger" }, + )(); + + const result = await driveSupervisor( + makeStartupDetails("trigger"), + replies({ + TriggerDagRun: [OK], + GetDagRunState: [ + { body: { type: "DagRunStateResult", state: "queued" } }, + { body: { type: "DagRunStateResult", state: "running" } }, + { body: { type: "DagRunStateResult", state: "success" } }, + ], + }), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" }); + const runId = requestsOf(result, "TriggerDagRun")[0]!["run_id"]; + expect(requestsOf(result, "GetDagRunState")).toEqual( + Array(3).fill({ type: "GetDagRunState", dag_id: "downstream", run_id: runId }), + ); + }); + + it("fails, with retries honoured, when the polled run reaches a failed state", async () => { + testDag.task( + triggerDagRun({ dagId: "downstream", waitForCompletion: true, pokeInterval: 0 }), + { taskId: "trigger" }, + )(); + + const result = await driveSupervisor( + makeStartupDetails("trigger", "test_dag", "r1", { should_retry: true }), + replies({ + TriggerDagRun: [OK], + GetDagRunState: [{ body: { type: "DagRunStateResult", state: "failed" } }], + }), + ); + + expect(result.firstResponse!.body).toMatchObject({ + type: "RetryTask", + retry_reason: "downstream failed with failed state failed", + }); + }); + + it("defers to DagStateTrigger with the kwargs Python's serializer writes", async () => { + testDag.task( + triggerDagRun({ + dagId: "downstream", + waitForCompletion: true, + deferrable: true, + pokeInterval: 5, + allowedStates: ["success"], + failedStates: ["failed", "queued"], + }), + { taskId: "trigger" }, + )(); + + const result = await driveSupervisor( + makeStartupDetails("trigger"), + replies({ TriggerDagRun: [OK] }), + ); + + const runId = requestsOf(result, "TriggerDagRun")[0]!["run_id"]; + expect(result.firstResponse!.body).toEqual({ + type: "DeferTask", + state: "deferred", + classpath: DAG_STATE_TRIGGER, + trigger_kwargs: { + dag_id: "downstream", + states: ["success", "failed", "queued"], + poll_interval: 5, + run_ids: [runId], + execution_dates: null, + }, + trigger_timeout: null, + queue: null, + next_method: "execute_complete", + next_kwargs: {}, + }); + expect(requestsOf(result, "GetDagRunState")).toEqual([]); + }); + + it.each([ + ["False", null], + ["True", "default"], + ])( + "with AIRFLOW__TRIGGERER__QUEUES_ENABLED=%s, hands the trigger the queue %s", + async (enabled, queue) => { + vi.stubEnv("AIRFLOW__TRIGGERER__QUEUES_ENABLED", enabled); + testDag.task( + triggerDagRun({ dagId: "downstream", waitForCompletion: true, deferrable: true }), + { taskId: "trigger" }, + )(); + + const result = await driveSupervisor( + makeStartupDetails("trigger"), + replies({ TriggerDagRun: [OK] }), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "DeferTask", queue }); + }, + ); + + it("ignores deferrable without waitForCompletion, as Python does", async () => { + testDag.task(triggerDagRun({ dagId: "downstream", deferrable: true }), { + taskId: "trigger", + })(); + + const result = await driveSupervisor( + makeStartupDetails("trigger"), + replies({ TriggerDagRun: [OK] }), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" }); + }); + + function resumedWith(state: string, runId = "manual__2026-09-29T12:00:00+00:00") { + return { + next_method: "execute_complete", + next_kwargs: { + event: { + __classname__: "builtins.tuple", + __version__: 1, + __data__: [ + DAG_STATE_TRIGGER, + { + dag_id: "downstream", + states: ["success", "failed"], + poll_interval: 60, + run_ids: [runId], + execution_dates: null, + [runId]: state, + }, + ], + }, + }, + }; + } + + it("succeeds on resume when the triggered run finished in an allowed state", async () => { + testDag.task( + triggerDagRun({ dagId: "downstream", waitForCompletion: true, deferrable: true }), + { taskId: "trigger" }, + )(); + + const result = await driveSupervisor( + makeStartupDetails("trigger", "test_dag", "r1", resumedWith("success")), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "SucceedTask" }); + expect(result.runtimeRequests).toEqual([]); + }); + + it("fails on resume when the triggered run finished in a failed state", async () => { + testDag.task( + triggerDagRun({ dagId: "downstream", waitForCompletion: true, deferrable: true }), + { taskId: "trigger" }, + )(); + + const result = await driveSupervisor( + makeStartupDetails("trigger", "test_dag", "r1", resumedWith("failed")), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "TaskState", state: "failed" }); + }); + + it("fails on resume when the event cannot be read", async () => { + testDag.task( + triggerDagRun({ dagId: "downstream", waitForCompletion: true, deferrable: true }), + { taskId: "trigger" }, + )(); + + const result = await driveSupervisor( + makeStartupDetails("trigger", "test_dag", "r1", { + next_method: "execute_complete", + next_kwargs: { event: "garbage" }, + }), + ); + + expect(result.firstResponse!.body).toMatchObject({ type: "TaskState", state: "failed" }); + }); + + it("fails on resume through __fail__, when the trigger failed or timed out", async () => { + testDag.task( + triggerDagRun({ dagId: "downstream", waitForCompletion: true, deferrable: true }), + { taskId: "trigger" }, + )(); + + const result = await driveSupervisor( + makeStartupDetails("trigger", "test_dag", "r1", { + should_retry: true, + next_method: "__fail__", + next_kwargs: { error: "Trigger timeout", traceback: ["Traceback", "TimeoutError"] }, + }), + ); + + expect(result.firstResponse!.body).toMatchObject({ + type: "RetryTask", + retry_reason: "Trigger timeout", + }); + expect( + result.logRecords.some((r) => String(r["event"]).includes("Trigger failed:\nTraceback")), + ).toBe(true); + }); + }); }); diff --git a/ts-sdk/tests/sdk/trigger-dag-run.test.ts b/ts-sdk/tests/sdk/trigger-dag-run.test.ts new file mode 100644 index 0000000000000..d3c7ca7eb5a8f --- /dev/null +++ b/ts-sdk/tests/sdk/trigger-dag-run.test.ts @@ -0,0 +1,308 @@ +/*! + * 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 { serializeDag } from "../../src/coordinator/serde.js"; +import { Bundle, finalizeBundleDags, getBundleTrigger } from "../../src/sdk/bundle.js"; +import { + Dag, + finalizeDag, + getDagOrderEdges, + getDagTaskRecords, + type TaskRecord, +} from "../../src/sdk/dag.js"; +import { triggerDagRun, type TriggerDagRunSpec } from "../../src/sdk/trigger-dag-run.js"; + +type Json = Record; + +function serialize(dag: Dag) { + const serialized = serializeDag(dag, "", ".") as Json; + const tasks = serialized["tasks"] as { __var: Json }[]; + return { + serialized, + tasks: new Map(tasks.map(({ __var }) => [__var["task_id"] as string, __var])), + }; +} + +function triggered(spec: Partial = {}, taskSpec = {}) { + const dag = new Dag("d"); + const trigger = dag.task( + triggerDagRun({ dagId: "downstream_etl", ...spec } as TriggerDagRunSpec), + { taskId: "trigger_downstream", ...taskSpec }, + )(); + return { dag, trigger }; +} + +describe("triggerDagRun", () => { + afterEach(() => { + vi.unstubAllEnvs(); + }); + + it("is placed by calling the factory dag.task returns, like any task", () => { + const { dag, trigger } = triggered(); + + expect(trigger).toMatchObject({ dagId: "d", taskId: "trigger_downstream" }); + expect(dag.taskIds).toEqual(["trigger_downstream"]); + }); + + it("carries no handler, and its options with the operator's defaults", () => { + const { dag } = triggered(); + + const record = getDagTaskRecords(dag).get("trigger_downstream")!; + expect(record.fn).toBeUndefined(); + expect(record.trigger).toEqual({ + dagId: "downstream_etl", + runId: undefined, + conf: undefined, + resetDagRun: false, + waitForCompletion: false, + pokeInterval: 60, + allowedStates: ["success"], + failedStates: ["failed"], + skipWhenAlreadyExists: false, + failWhenDagIsPaused: false, + note: undefined, + deferrable: false, + }); + const bundle = new Bundle(dag); + expect(bundle.getTaskHandler("d", "trigger_downstream")).toBeUndefined(); + expect(getBundleTrigger(bundle, "d", "trigger_downstream")).toBe(record.trigger); + }); + + it("falls back on the default for empty allowedStates, not for empty failedStates", () => { + const { dag } = triggered({ allowedStates: [], failedStates: [] }); + + expect(getDagTaskRecords(dag).get("trigger_downstream")!.trigger).toMatchObject({ + allowedStates: ["success"], + failedStates: [], + }); + }); + + it.each([ + ["True", {}, true], + ["false", {}, false], + ["1", { deferrable: false }, false], + ["f", { deferrable: true }, true], + ])( + "with AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE=%s and %o, defers: %s", + (value, spec, deferrable) => { + vi.stubEnv("AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE", value); + + const { dag } = triggered(spec); + + expect(getDagTaskRecords(dag).get("trigger_downstream")!.trigger).toMatchObject({ + deferrable, + }); + }, + ); + + it("is complete once the factory is called", () => { + const { dag } = triggered(); + + expect(() => finalizeBundleDags(new Bundle(dag))).not.toThrow(); + }); + + it("keeps a task record to either a handler or a trigger", () => { + const { dag } = triggered(); + const record = getDagTaskRecords(dag).get("trigger_downstream")!; + + // @ts-expect-error -- a record cannot carry both a handler and a trigger. + const both: TaskRecord = { ...record, fn: async () => undefined }; + + expect(both.trigger).toBe(record.trigger); + }); + + it("takes its id positionally too", () => { + const dag = new Dag("d"); + dag.task("trigger_downstream", triggerDagRun({ dagId: "downstream_etl" }))(); + + expect(dag.taskIds).toEqual(["trigger_downstream"]); + }); + + it("serializes as a TypeScript task, drawn as TriggerDagRunOperator", () => { + const { dag } = triggered({ waitForCompletion: true, pokeInterval: 30 }); + + const task = serialize(dag).tasks.get("trigger_downstream")!; + expect(task).toMatchObject({ + task_id: "trigger_downstream", + task_type: "TypeScriptOperator", + _task_module: "airflow.sdk.coordinators.node", + language: "typescript", + is_stub: true, + template_fields: [], + _operator_name: "TriggerDagRunOperator", + ui_color: "#ffefeb", + _operator_extra_links: { "Triggered DAG": "_link_TriggerDagRunLink" }, + }); + expect(task).not.toHaveProperty("trigger_dag_id"); + expect(task).not.toHaveProperty("wait_for_completion"); + expect(task).not.toHaveProperty("_arg_bindings"); + }); + + it("records the triggered Dag as a dependency of this one", () => { + const { dag } = triggered(); + + expect(serialize(dag).serialized["dag_dependencies"]).toEqual([ + { + source: "d", + target: "downstream_etl", + label: "trigger_downstream", + dependency_type: "trigger", + dependency_id: "trigger_downstream", + }, + ]); + }); + + it("inherits the Dag's queue, so it runs where the Dag's other tasks run", () => { + const dag = new Dag("d", { queue: "typescript" }); + dag.task("extract", async () => undefined)(); + dag.task(triggerDagRun({ dagId: "downstream_etl" }), { taskId: "trigger_downstream" })(); + + const { tasks } = serialize(dag); + expect(tasks.get("extract")?.["queue"]).toBe("typescript"); + expect(tasks.get("trigger_downstream")?.["queue"]).toBe("typescript"); + }); + + it("still takes a task spec", () => { + const { dag } = triggered({}, { retries: 2, queue: "typescript_heavy" }); + + expect(serialize(dag).tasks.get("trigger_downstream")).toMatchObject({ + retries: 2, + queue: "typescript_heavy", + }); + }); + + it("participates in before and after like any task", () => { + const dag = new Dag("d"); + const loaded = dag.task("load", async () => undefined)(); + const trigger = dag.task(triggerDagRun({ dagId: "downstream_etl" }), { + taskId: "trigger_downstream", + })(); + + trigger.after(loaded); + + expect(getDagOrderEdges(dag)).toEqual([{ upstream: "load", downstream: "trigger_downstream" }]); + expect(serialize(dag).tasks.get("load")?.["downstream_task_ids"]).toEqual([ + "trigger_downstream", + ]); + }); + + it("sits in a task group like any task, prefix included", () => { + const dag = new Dag("d"); + dag + .taskGroup("publish") + .task(triggerDagRun({ dagId: "downstream_etl" }), { taskId: "trigger" })(); + + expect(dag.taskIds).toEqual(["publish.trigger"]); + }); + + describe("rejects", () => { + it("no taskId", () => { + const dag = new Dag("d"); + + expect(() => dag.task(triggerDagRun({ dagId: "downstream_etl" }), {})).toThrowError( + /A triggerDagRun task of Dag "d" has no taskId/, + ); + }); + + it("inputs given to a trigger task", () => { + const dag = new Dag("d"); + const factory = dag.task(triggerDagRun({ dagId: "downstream_etl" }), { + taskId: "trigger_downstream", + }) as unknown as (inputs: unknown) => unknown; + + expect(() => factory({ rows: 1 })).toThrowError( + /Task "trigger_downstream" of Dag "d" is a triggerDagRun task, which takes no inputs/, + ); + }); + + it("a trigger task left uncalled, when the Dag is read", () => { + const dag = new Dag("d"); + dag.task(triggerDagRun({ dagId: "downstream_etl" }), { taskId: "trigger_downstream" }); + + expect(() => finalizeDag(dag)).toThrowError(/trigger_downstream/); + }); + + it.each([ + ["no dagId", {}], + ["an empty dagId", { dagId: "" }], + ["a dagId that is not a string", { dagId: 7 }], + ])("%s", (_label, spec) => { + expect(() => triggerDagRun(spec as unknown as TriggerDagRunSpec)).toThrowError( + /triggerDagRun\(\.\.\.\) needs the dagId of the Dag to trigger/, + ); + }); + + it.each([ + ["null", null], + ["a string", "downstream_etl"], + ["an array", []], + ])("%s where the options belong", (_label, spec) => { + expect(() => triggerDagRun(spec as unknown as TriggerDagRunSpec)).toThrowError( + /triggerDagRun\(\.\.\.\) takes an options object/, + ); + }); + + it("an option the operator does not declare", () => { + expect(() => + triggerDagRun({ dagId: "d2", waitFor: true } as unknown as TriggerDagRunSpec), + ).toThrowError(/Unknown option "waitFor" for triggerDagRun\(\.\.\.\)/); + }); + + it.each([ + ["a Date", new Date("2026-01-01")], + ["a Map", new Map()], + ["a function", () => 1], + ])("%s inside conf, which the Dag JSON cannot carry", (_label, value) => { + const { dag } = triggered({ conf: { bad: value } as never }); + + expect(() => serialize(dag)).toThrowError( + /conf of task "trigger_downstream" of Dag "d" holds .*, which JSON cannot carry/, + ); + }); + + it.each([ + ["allowedStates", { allowedStates: ["done"] }], + ["failedStates", { failedStates: ["error"] }], + ])("a %s entry that is not a Dag run state", (name, spec) => { + expect(() => + triggerDagRun({ dagId: "d2", ...spec } as unknown as TriggerDagRunSpec), + ).toThrowError(new RegExp(`option "${name}" holds "\\w+", which is not a Dag run state`)); + }); + + it.each([ + ["a negative pokeInterval", { pokeInterval: -1 }, /"pokeInterval" must be a non-negative/], + ["a non-boolean deferrable", { deferrable: "yes" }, /"deferrable" must be a boolean/], + ["a non-string runId", { runId: 7 }, /"runId" must be a string/], + ["a conf that is an array", { conf: [1] }, /"conf" must be an object/], + ])("%s", (_label, spec, error) => { + expect(() => + triggerDagRun({ dagId: "d2", ...spec } as unknown as TriggerDagRunSpec), + ).toThrowError(error); + }); + + it("an AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE that is not a boolean", () => { + vi.stubEnv("AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE", "yes"); + + expect(() => triggered()).toThrowError( + 'AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE is "yes", which is not a boolean; use true or false', + ); + }); + }); +});