Skip to content

Commit bc4cde1

Browse files
committed
TS SDK: draw order-only edges with before and after
1 parent ea3d111 commit bc4cde1

4 files changed

Lines changed: 330 additions & 17 deletions

File tree

‎airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -301,6 +301,28 @@ The task id may be omitted, in which case it is the handler's function name:
301301
name of its own, such as an arrow function passed inline, has nothing to take an id from and needs
302302
one: either positionally or as ``taskId`` in its spec. Give it in one place only, not both.
303303

304+
Order-only edges
305+
~~~~~~~~~~~~~~~~
306+
307+
An edge that carries no value has no argument name to travel under, so it is drawn between the
308+
references themselves with ``before`` and ``after``, the TypeScript pair for Python's ``>>`` and
309+
``<<``:
310+
311+
.. code-block:: typescript
312+
313+
const loaded = load({ transformed });
314+
const cleaned = cleanup();
315+
316+
loaded.before(cleaned); // loaded >> cleaned
317+
cleaned.after(loaded, transformed); // [loaded, transformed] >> cleaned
318+
319+
Both take any number of references, so one call draws several edges, and drawing an edge that
320+
already exists changes nothing. Each returns the reference it was called on, so
321+
``loaded.before(cleaned).before(notified)`` draws both edges from ``loaded``.
322+
323+
Pass a value as an argument when the downstream task needs it, and use ``before`` or ``after`` when
324+
it only needs to run in order.
325+
304326
``new Dag`` and ``dag.task`` both take a trailing spec of Airflow options:
305327
``{ schedule: "@daily", tags: ["etl"] }`` for the Dag, ``{ retries: 2, retryDelay: 30 }`` for a task.
306328

‎ts-sdk/src/sdk/dag.ts‎

Lines changed: 118 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -128,9 +128,6 @@ declare const RETURN_TYPE: unique symbol;
128128
/**
129129
* A reference to the result of one task, returned by calling that task.
130130
*
131-
* Identity only: the handler and the value are deliberately not exposed. Pass a
132-
* reference as an input of a downstream task to make that task depend on it.
133-
*
134131
* `TReturn` is the handler's return type, so a construct that needs a
135132
* particular one can ask for it. A reference of a narrower type is usable
136133
* wherever a wider one is: a `TaskRef<number>` is a `TaskRef<unknown>`.
@@ -146,8 +143,48 @@ export interface TaskRef<TReturn = unknown> {
146143
readonly taskId: string;
147144
/** @internal Never set; see {@link RETURN_TYPE}. */
148145
readonly [RETURN_TYPE]?: TReturn;
146+
/**
147+
* Run this task before each of `downstream`, carrying no value — the
148+
* TypeScript spelling of Python's `>>`.
149+
*
150+
* ```ts
151+
* loaded.before(cleaned, notified); // loaded >> [cleanup, notify]
152+
* ```
153+
*
154+
* Variadic, so one call fans out, and it returns its own receiver rather
155+
* than its arguments: a fan-out has no single "next" reference to hand back.
156+
* Declaring an edge that already exists changes nothing.
157+
*/
158+
before(...downstream: readonly TaskRef[]): TaskRef<TReturn>;
159+
/**
160+
* Run this task after each of `upstream`, carrying no value — Python's `<<`.
161+
*
162+
* ```ts
163+
* cleaned.after(loaded, transformed); // [load, transform] >> cleanup
164+
* ```
165+
*
166+
* Fan-*in* that carries data is the wiring object instead
167+
* (`summarize({ north: extractNorth(), south: extractSouth() })`), so each
168+
* direction has an answer: named keys when values flow, `after` when only
169+
* order does.
170+
*/
171+
after(...upstream: readonly TaskRef[]): TaskRef<TReturn>;
172+
}
173+
174+
/**
175+
* An order-only edge of a Dag: upstream task ID, then downstream task ID.
176+
*
177+
* Kept apart from the wiring a factory call records, because an edge that
178+
* carries no value has no argument name to be recorded under.
179+
*/
180+
export interface OrderEdge {
181+
readonly upstream: string;
182+
readonly downstream: string;
149183
}
150184

185+
// A task id cannot hold a NUL, so a joined pair cannot collide with one.
186+
const EDGE_KEY_SEPARATOR = "\u0000";
187+
151188
/** Whether `value` is a TaskRef returned by any copy of this package. */
152189
function isTaskRef(value: unknown): value is TaskRef {
153190
return hasBrand(value, "TaskRef");
@@ -293,6 +330,7 @@ export type RecordedInputs = Readonly<Record<string, TaskRef | JsonValue>>;
293330
// Dag's private state without public accessors on the class.
294331
let taskRecordsOf: (dag: Dag) => ReadonlyMap<string, TaskRecord>;
295332
let inputsOf: (dag: Dag) => ReadonlyMap<string, RecordedInputs>;
333+
let orderEdgesOf: (dag: Dag) => readonly OrderEdge[];
296334
let finalizeOf: (dag: Dag) => void;
297335

298336
/** Internal: whether `value` is a Dag built by any copy of this package. */
@@ -326,11 +364,15 @@ export class Dag {
326364
readonly spec: DagSpec;
327365
readonly #tasks = new Map<string, TaskRecord>();
328366
readonly #inputs = new Map<string, RecordedInputs>();
367+
// Keyed by the two task ids, so declaring an edge twice records it once, and
368+
// insertion-ordered so the serialized Dag reads as written.
369+
readonly #orderEdges = new Map<string, OrderEdge>();
329370
#finalized = false;
330371

331372
static {
332373
taskRecordsOf = (dag) => dag.#tasks;
333374
inputsOf = (dag) => dag.#inputs;
375+
orderEdgesOf = (dag) => [...dag.#orderEdges.values()];
334376
finalizeOf = (dag) => dag.#finalize();
335377
}
336378

@@ -441,7 +483,7 @@ export class Dag {
441483
);
442484
}
443485
const spec = this.#taskSpecOf(taskId, options);
444-
const task = createTaskRef(this.dagId, taskId);
486+
const task = this.#createTaskRef(taskId);
445487
this.#tasks.set(taskId, {
446488
task,
447489
// The runtime dispatches every handler through one instantiation, as it
@@ -474,6 +516,73 @@ export class Dag {
474516
return spec as TaskSpec;
475517
}
476518

519+
#createTaskRef(taskId: string): TaskRef {
520+
const task: TaskRef = {
521+
dagId: this.dagId,
522+
taskId,
523+
before: (...downstream) => {
524+
for (const other of downstream) this.#addOrderEdge(task, other, "before");
525+
return task;
526+
},
527+
after: (...upstream) => {
528+
for (const other of upstream) this.#addOrderEdge(other, task, "after");
529+
return task;
530+
},
531+
};
532+
brand(task, "TaskRef");
533+
return Object.freeze(task);
534+
}
535+
536+
#addOrderEdge(upstream: TaskRef, downstream: TaskRef, verb: "before" | "after"): void {
537+
if (this.#finalized) {
538+
throw new Error(
539+
`An edge was drawn on Dag "${this.dagId}" after the Dag was read; ` +
540+
"declare every edge while the module is loading",
541+
);
542+
}
543+
// The argument is the one that can be foreign: the receiver is a reference
544+
// this Dag handed out, since it is what carries the method.
545+
const other = verb === "before" ? downstream : upstream;
546+
this.#validateOwnRef(other, verb);
547+
if (upstream.taskId === downstream.taskId) {
548+
throw new Error(
549+
`${verb}() cannot draw an edge from task "${upstream.taskId}" of Dag "${this.dagId}" to ` +
550+
"itself; an edge orders two different tasks",
551+
);
552+
}
553+
const key = `${upstream.taskId}${EDGE_KEY_SEPARATOR}${downstream.taskId}`;
554+
// Idempotent, so an edge drawn from both ends is one edge.
555+
if (!this.#orderEdges.has(key)) {
556+
this.#orderEdges.set(
557+
key,
558+
Object.freeze({ upstream: upstream.taskId, downstream: downstream.taskId }),
559+
);
560+
}
561+
}
562+
563+
#validateOwnRef(ref: TaskRef, verb: string): void {
564+
if (!isTaskRef(ref)) {
565+
throw new Error(
566+
`${verb}() on Dag "${this.dagId}" takes task references returned by calling a task, ` +
567+
"not arbitrary values",
568+
);
569+
}
570+
if (ref.dagId !== this.dagId) {
571+
throw new Error(
572+
`${verb}() cannot draw an edge to Dag "${ref.dagId}" task "${ref.taskId}" from Dag ` +
573+
`"${this.dagId}"; an edge joins two tasks of one Dag`,
574+
);
575+
}
576+
// Identity, not the ID pair: two Dag objects can carry the same dagId, and
577+
// a second resolved copy of this package brands its own references.
578+
if (this.#tasks.get(ref.taskId)?.task !== ref) {
579+
throw new Error(
580+
`${verb}() was given a reference to "${ref.taskId}" that this Dag did not hand out; ` +
581+
`it comes from another Dag object with the same ID, or ${DUPLICATE_COPY_HINT}`,
582+
);
583+
}
584+
}
585+
477586
#wire(taskId: string, inputs: unknown): void {
478587
if (this.#finalized) {
479588
throw new Error(
@@ -642,12 +751,6 @@ function destructuredKeys(handler: unknown): readonly string[] | undefined {
642751
return Object.freeze(keys as string[]);
643752
}
644753

645-
function createTaskRef(dagId: string, taskId: string): TaskRef {
646-
const task: TaskRef = { dagId, taskId };
647-
brand(task, "TaskRef");
648-
return Object.freeze(task);
649-
}
650-
651754
function validateDagSpec(dagId: string, spec: DagSpec): void {
652755
const value: unknown = spec;
653756
if (!isPlainRecord(value)) {
@@ -670,6 +773,11 @@ export function getDagTaskRecords(dag: Dag): ReadonlyMap<string, TaskRecord> {
670773
return taskRecordsOf(dag);
671774
}
672775

776+
/** Internal: the order-only edges of a Dag, in the order they were drawn. */
777+
export function getDagOrderEdges(dag: Dag): readonly OrderEdge[] {
778+
return orderEdgesOf(dag);
779+
}
780+
673781
/** Internal: what each task of a Dag was called with, keyed by task ID.
674782
* A task that has not been called is absent. */
675783
export function getDagTaskInputs(dag: Dag): ReadonlyMap<string, RecordedInputs> {

‎ts-sdk/tests/public-api.test.ts‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,8 +61,11 @@ describe("public API", () => {
6161
);
6262
const upstream = upstreamTask();
6363
const downstream = downstreamTask({ upstream });
64-
expect(upstream).toEqual({ dagId: "public_api_dag", taskId: "public_api_task" });
65-
expect(downstream).toEqual({ dagId: "public_api_dag", taskId: "public_api_downstream" });
64+
expect(upstream).toMatchObject({ dagId: "public_api_dag", taskId: "public_api_task" });
65+
expect(downstream).toMatchObject({
66+
dagId: "public_api_dag",
67+
taskId: "public_api_downstream",
68+
});
6669
expect(dag.taskIds).toEqual(["public_api_task", "public_api_downstream"]);
6770
// serve() hands the bundle to the runtime, which needs the supervisor's
6871
// socket addresses that Airflow puts on argv.
@@ -300,13 +303,26 @@ describe("public API", () => {
300303
});
301304

302305
it("keeps the Dag authoring signatures extensible via trailing specs", () => {
306+
// Identity, plus the two order-only edge verbs; the handler and the value
307+
// stay hidden.
308+
expectTypeOf<Extract<keyof TaskRef, string>>().toEqualTypeOf<
309+
"dagId" | "taskId" | "before" | "after"
310+
>();
303311
expectTypeOf<TaskRef["dagId"]>().toEqualTypeOf<string>();
304312
expectTypeOf<TaskRef["taskId"]>().toEqualTypeOf<string>();
305313
// A reference carries its handler's return type, so a construct that needs
306314
// a particular one can ask for it: a narrower reference is usable wherever
307315
// a wider one is, and not the other way round.
308316
expectTypeOf<TaskRef<boolean>>().toMatchTypeOf<TaskRef>();
309317
expectTypeOf<TaskRef>().not.toMatchTypeOf<TaskRef<boolean>>();
318+
// Variadic, and each returns its own receiver rather than its arguments,
319+
// return type included.
320+
expectTypeOf<TaskRef<boolean>["before"]>().toEqualTypeOf<
321+
(...downstream: readonly TaskRef[]) => TaskRef<boolean>
322+
>();
323+
expectTypeOf<TaskRef<boolean>["after"]>().toEqualTypeOf<
324+
(...upstream: readonly TaskRef[]) => TaskRef<boolean>
325+
>();
310326
// Wiring moved to the factory call, and the spec is the trailing argument
311327
// itself.
312328
expectTypeOf<TaskOptions>().toEqualTypeOf<TaskSpec>();

0 commit comments

Comments
 (0)