Skip to content

Commit 1b7e067

Browse files
committed
TS SDK: answer the Dag-parsing request from bundle.serve
1 parent 67fb9f2 commit 1b7e067

3 files changed

Lines changed: 185 additions & 15 deletions

File tree

‎ts-sdk/src/coordinator/runtime.ts‎

Lines changed: 68 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,10 @@ import {
5555
type StartupDetails,
5656
} from "./protocol.js";
5757
import { getArgNames } from "../sdk/arg-names.js";
58-
import { bundleDagTaskIds, type Bundle } from "../sdk/bundle.js";
58+
import { bundleDagTaskIds, listBundleNativeDags, type Bundle } from "../sdk/bundle.js";
59+
import type { Dag } from "../sdk/dag.js";
60+
import { SERIALIZATION_VERSION } from "../generated/dag-schema-fields.js";
61+
import { computeRelativeFileloc, serializeDag } from "./serde.js";
5962
import { runInTaskScope, type TaskContext } from "../sdk/task.js";
6063
import type { JsonValue } from "../sdk/client-types.js";
6164

@@ -262,22 +265,77 @@ export function createRuntimeAbort(
262265
};
263266
}
264267

268+
/**
269+
* Answer a parse request with the Dags this bundle declared in TypeScript.
270+
*
271+
* No handler body runs: a `TaskRef` is inert, so reading a Dag only walks what
272+
* its module already built. Reading it is also what enforces that every task
273+
* was called exactly once, which is why a Dag that is not fully laid out
274+
* surfaces here.
275+
*
276+
* A Dag that cannot be serialized becomes an import error against this file,
277+
* as a Python Dag file that raises does, rather than failing the whole parse:
278+
* one broken Dag must not take out the others a bundle serves.
279+
*/
265280
function handleParse(
266281
request: { file: string; bundle_path: string },
267282
bundle: Bundle,
268283
logs: LogChannel,
269284
): RuntimeDagFileParsingResult {
270-
// TypeScript-native Dag parsing is not yet supported.
271-
// Respond with an empty result so the Python-stub-Dag workflow works.
272-
logs.info("Parse-mode response (TS Dag parsing not yet supported)", {
273-
registered_tasks: Object.fromEntries(bundleDagTaskIds(bundle)),
285+
const fileloc = request.file;
286+
const relativeFileloc = computeRelativeFileloc(fileloc, request.bundle_path);
287+
const serializedDags: { data: Record<string, unknown> }[] = [];
288+
// Airflow keys an import error by the bundle-relative path and holds one row
289+
// per file (`DagFileProcessorManager.update_import_errors`), so every failure
290+
// in this bundle is reported under that one key, naming its Dag in the
291+
// message. An absolute path, or one with a Dag id appended, would give a row
292+
// the UI cannot tie back to the file, and would leave the file itself looking
293+
// healthy while its Dags had vanished.
294+
const failures: string[] = [];
295+
296+
let dags: Dag[];
297+
try {
298+
dags = listBundleNativeDags(bundle);
299+
} catch (err) {
300+
// Reading the bundle closes every native Dag in it, so a fault in one is
301+
// reported against the file rather than leaving the request unanswered.
302+
const detail = err instanceof Error ? err.message : String(err);
303+
logs.error("Bundle could not be read for parsing", { fileloc, detail });
304+
return {
305+
type: "DagFileParsingResult",
306+
fileloc,
307+
serialized_dags: [],
308+
import_errors: { [relativeFileloc]: detail },
309+
};
310+
}
311+
312+
for (const dag of dags) {
313+
try {
314+
serializedDags.push({
315+
data: {
316+
__version: SERIALIZATION_VERSION,
317+
dag: serializeDag(dag, fileloc, relativeFileloc),
318+
},
319+
});
320+
} catch (err) {
321+
const detail = err instanceof Error ? err.message : String(err);
322+
logs.error("Dag could not be serialized", { dag_id: dag.dagId, detail });
323+
failures.push(`Dag "${dag.dagId}": ${detail}`);
324+
}
325+
}
326+
327+
logs.info("Parse-mode response", {
328+
fileloc,
329+
dag_ids: dags.map((dag) => dag.dagId),
330+
serialized: serializedDags.length,
331+
import_errors: failures.length,
274332
});
275-
const response: RuntimeDagFileParsingResult = {
333+
return {
276334
type: "DagFileParsingResult",
277-
fileloc: request.file,
278-
serialized_dags: [],
279-
};
280-
return response;
335+
fileloc,
336+
serialized_dags: serializedDags,
337+
...(failures.length > 0 && { import_errors: { [relativeFileloc]: failures.join("\n") } }),
338+
} as RuntimeDagFileParsingResult;
281339
}
282340

283341
async function handleTask(

‎ts-sdk/src/sdk/bundle.ts‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -239,3 +239,21 @@ export function bundleDagTaskIds(bundle: Bundle): Map<string, string[]> {
239239
}
240240
return byDag;
241241
}
242+
243+
/**
244+
* Internal: the Dags this bundle declared in TypeScript, in registration order.
245+
*
246+
* What a parse request answers with, and what closes each of them: a Dag known
247+
* only through task handlers is not one of them, because its graph belongs to
248+
* the Python Dag file that declares it, and serializing it here would register
249+
* a second Dag with the same `dag_id` from a different `fileloc`.
250+
*
251+
* Supersedes {@link finalizeBundleDags} for a caller that also wants the Dags.
252+
*/
253+
export function listBundleNativeDags(bundle: Bundle): Dag[] {
254+
const dags = [...dagsOf(bundle).values()];
255+
for (const dag of dags) {
256+
finalizeDag(dag);
257+
}
258+
return dags;
259+
}

‎ts-sdk/tests/coordinator/integration.test.ts‎

Lines changed: 99 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -798,18 +798,112 @@ describe("coordinator runtime integration", () => {
798798
expect(calledSecondDag).toBe(false);
799799
});
800800

801-
it("returns empty serialized_dags for DagFileParseRequest", async () => {
801+
describe("DagFileParseRequest", () => {
802802
const parseRequest = {
803803
type: "DagFileParseRequest",
804804
file: "/dags/test.mjs",
805805
bundle_path: "/dags",
806806
};
807807

808-
const result = await driveSupervisor(parseRequest);
808+
async function parse(): Promise<Record<string, unknown>> {
809+
const result = await driveSupervisor(parseRequest);
810+
return result.firstResponse!.body as Record<string, unknown>;
811+
}
812+
813+
it("answers with the Dags the bundle declared in TypeScript", async () => {
814+
testDag.task("extract", async () => undefined)();
815+
otherDag.task("stage", async () => undefined)();
816+
817+
const body = await parse();
818+
819+
expect(body.type).toBe("DagFileParsingResult");
820+
expect(body.fileloc).toBe("/dags/test.mjs");
821+
const dags = body.serialized_dags as { data: Record<string, unknown> }[];
822+
expect(dags.map((entry) => (entry.data.dag as Record<string, unknown>).dag_id)).toEqual([
823+
"test_dag",
824+
"other_dag",
825+
]);
826+
expect(dags[0]!.data.__version).toBe(3);
827+
expect((dags[0]!.data.dag as Record<string, unknown>).relative_fileloc).toBe("test.mjs");
828+
expect(body.import_errors).toBeUndefined();
829+
});
830+
831+
it("runs no handler body while parsing", async () => {
832+
let ran = false;
833+
testDag.task("extract", async () => {
834+
ran = true;
835+
})();
836+
otherDag.task("stage", async () => undefined)();
837+
838+
await parse();
839+
840+
expect(ran).toBe(false);
841+
});
842+
843+
it("answers with nothing when the bundle only binds handlers to Python Dags", async () => {
844+
// A Python Dag's graph belongs to the Python file that declares it, so a
845+
// bundle of task handlers has no Dag of its own to serialize.
846+
bundle = new Bundle(new TaskHandler("py_dag", "transform", async () => undefined));
847+
848+
const body = await parse();
849+
850+
expect(body.serialized_dags).toEqual([]);
851+
expect(body.import_errors).toBeUndefined();
852+
});
853+
854+
it("reports an uncalled task as an import error rather than failing the parse", async () => {
855+
testDag.task("extract", async () => undefined)();
856+
testDag.task("orphan", async () => undefined);
809857

810-
const body = result.firstResponse!.body as Record<string, unknown>;
811-
expect(body.type).toBe("DagFileParsingResult");
812-
expect(body.serialized_dags).toEqual([]);
858+
const body = await parse();
859+
860+
expect(body.serialized_dags).toEqual([]);
861+
// Keyed by the bundle-relative path, which is how Airflow ties the row
862+
// to the file it came from.
863+
expect(body.import_errors).toEqual({
864+
"test.mjs": expect.stringContaining('Task "orphan" of Dag "test_dag" is never called'),
865+
});
866+
});
867+
868+
it("merges several failing Dags into the one row Airflow keeps per file", async () => {
869+
testDag.task("extract", async () => undefined)();
870+
for (const dagId of ["broken_a", "broken_b"]) {
871+
const broken = new Dag(dagId, { schedule: "" });
872+
broken.task("t", async () => undefined)();
873+
bundle.register(broken);
874+
}
875+
876+
const body = await parse();
877+
878+
const errors = body.import_errors as Record<string, string>;
879+
expect(Object.keys(errors)).toEqual(["test.mjs"]);
880+
expect(errors["test.mjs"]).toContain('Dag "broken_a"');
881+
expect(errors["test.mjs"]).toContain('Dag "broken_b"');
882+
});
883+
884+
it("keeps the Dags it can serialize when one of them cannot be", async () => {
885+
testDag.task("extract", async () => undefined)();
886+
// An empty schedule is rejected by the serializer, and only that Dag is
887+
// lost: the rest of the bundle still parses.
888+
const broken = new Dag("broken_dag", { schedule: "" });
889+
broken.task("t", async () => undefined)();
890+
bundle.register(broken);
891+
892+
const body = await parse();
893+
894+
const dags = body.serialized_dags as { data: Record<string, unknown> }[];
895+
expect(dags.map((entry) => (entry.data.dag as Record<string, unknown>).dag_id)).toEqual([
896+
"test_dag",
897+
"other_dag",
898+
]);
899+
// One row per file, so the failing Dag is named in the message rather
900+
// than appended to the key.
901+
expect(body.import_errors).toEqual({
902+
"test.mjs": expect.stringContaining(
903+
'Dag "broken_dag": schedule for Dag "broken_dag" is empty',
904+
),
905+
});
906+
});
813907
});
814908

815909
it("auto-pushes return_value XCom when handler returns a value", async () => {

0 commit comments

Comments
 (0)