Skip to content

Commit a8a6257

Browse files
TS SDK: embed one source region per native Dag file (#73723)
* TS SDK: embed one source region per native Dag file A bundle whose native Dags live in more than one file only carried its entrypoint's source in the airflowSource region. The Code tab would show the entrypoint for every Dag, which is the file for the ones declared there and unrelated for the ones declared in imported modules. Multi-Dag bundles were effectively single-source in the UI. Every Dag file that declares at least one native Dag now ships its own source region, mapped to its Dags by dag_id. The `Dag` constructor records which file it ran in via a cross-realm global slot that `airflow-ts-pack`'s onLoad plugin writes at the top of each author-owned source file. `Function.prototype.toString` and stack-trace parsing stay out — ES module hoisting is what makes the slot the current file for every top-level `new Dag(...)` call, and the constructor just reads it. * Rename sourceFiles, drop the entrypoint fallback Per-file source-size cap, back to the same 1 MB limit each source region already had. The bundle grows with the number of Dags, not with any aggregate cap on their sum. `EMBEDDED_SOURCE_OPEN_PREFIX`/`_SUFFIX` collapse to a single `EMBEDDED_SOURCE_MARKER`; the trailing newline is a formatting detail of the concatenation, not a named constant. `entrypointName` and its `--source` CLI option are gone. The metadata `entrypoint` field they filled is only useful as a fallback for "no per-Dag source", and each native Dag now carries its own source path. Bundles that serve nothing still fail at pack time; mixed-lang bundles still work with an empty `dag_source_paths`. * Reject block-comment terminators in source paths, and utf-8 the header A source path holding `*/` would close the marker's own block comment early, splicing the rest of the region into executable position; the packer now rejects such paths (and stray newlines) rather than trying to escape them. `renderHeader` was writing the layout line with `Buffer.from(..., "ascii")`, which silently truncated a non-ASCII path like `données.ts` and broke the metadata-to-region round trip. Encode the header as utf-8. Drops the whitespace/syntax/keepNames minify split back to plain `minify: true`: nothing in this PR depends on identifier preservation, so the earlier split needed a rationale it did not have. * Save and restore the module-source slot around each module body The tag used to only write. When A imported B, B's write overwrote A's, and any `new Dag(...)` in A after the import was attributed to B. The tag now captures the previous value on entry and restores it after the module's top-level code returns, so nested imports and top-level `await` within a single module hand the slot back on the way out. Two scoping limits remain and are documented in module-source.ts: a Dag constructed in a scheduled callback (setTimeout, microtask) fires after the module's epilog has restored, so it lands under whichever module touched the slot most recently; and two modules genuinely evaluated in parallel through Promise.all([import, import]) can interleave. Both are non-standard authoring patterns. * Read the multi-source TypeScript bundle layout in the Task SDK The airflow-ts-pack encoder now writes a `sources` array — one embedded region per native Dag file, keyed by `dag_source_paths`, and empty for a bundle whose Dags are all authored in another language — in place of the single `source` section. The Task SDK bundle reader still demanded that singular section, so every bundle the current packer produces was rejected with "bundle layout is missing the source section" and no TypeScript task could launch. Teach the reader only the new layout; no released bundle uses the old one, so this is a clean break rather than carrying both formats. * Embed the entrypoint source as a per-Dag fallback With one source region per native Dag, a Dag the packer cannot attribute to a file — one constructed dynamically in a callback, or a mixed-language Dag owned by Python — had no embedded source at all, so reading its source failed. Restore the earlier guarantee that the entry file is always embedded: record its path as entrypoint_path in the metadata and pack its contents as a source region, so a reader falls back to it whenever a Dag has no attributed file of its own. * Return no source for a Dag the bundle does not define in TypeScript read_bundle_source fell back to the bundle entrypoint for any Dag it could not map to a source file. That is wrong for a mixed-language Dag: the Dag is defined in Python, which owns its Code tab, and the TypeScript bundle only supplies task handlers — showing the TypeScript entrypoint for it would be misleading. Return None instead, so the caller can show the owning language's source. The entrypoint stays embedded and recorded as entrypoint_path, a bundle-level reference a caller can fall back to for a native TypeScript Dag that has no source of its own. * Let the caller choose when a bundle serves a Dag's TypeScript source A mixed-language Dag is defined in Python, which owns its Code tab, so the TypeScript bundle should offer no source for it. The caller signals that by omitting dag_id, and the reader returns None instead of an unrelated file. When a dag_id is given the caller is asking for a Dag it treats as TypeScript-defined, so a Dag missing from dag_source_paths — one built dynamically that the packer could not attribute — falls back to the embedded entrypoint source rather than nothing.
1 parent 5799660 commit a8a6257

13 files changed

Lines changed: 989 additions & 208 deletions

File tree

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -435,9 +435,12 @@ The code is minified because an integrity digest is only worth taking over an ar
435435
read or edit in place. Function names are kept through minification, since a task id defaults to its
436436
handler's name. The ``/*! */`` license banners of bundled dependencies are kept.
437437

438-
Because the shipped code is not the code anyone wrote, the packer also embeds the entry module verbatim in a
439-
``/*# airflowSource ... #*/`` block comment, verified by its own digest, so Airflow has something readable to
440-
display for the Dag. Only the entry module is embedded, not the modules it imports.
438+
Because the shipped code is not the code anyone wrote, the packer also embeds each Dag-defining source
439+
file verbatim in its own ``/*# airflowSource:<path> ... #*/`` block comment, verified by its own digest,
440+
so Airflow has something readable to display for each Dag. Each native Dag's file is embedded — the file
441+
its ``new Dag(...)`` constructor ran in — so a bundle that declares its Dags across several files gets one
442+
source region per file, mapped to their ``dag_id`` by the ``dag_source_paths`` field in the manifest. Files
443+
that only supply utilities or types are not embedded.
441444

442445
``esbuild`` is an optional peer dependency: packing is build-time only, so the runtime install of
443446
``apache-airflow-ts-sdk`` skips it, and it must be installed separately before running ``airflow-ts-pack``.

‎task-sdk/src/airflow/sdk/coordinators/node/_bundle_reader.py‎

Lines changed: 154 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,9 @@
1818
"""
1919
Read and verify TypeScript Dag bundles.
2020
21-
File order: layout comment, metadata comment, source block comment, then executable JavaScript.
22-
Read the headers, verify the section digests, and return coordinator metadata or the source.
21+
File order: layout comment, metadata comment, one source block comment per native Dag file,
22+
then executable JavaScript. Read the headers, verify the section digests, and return coordinator
23+
metadata or the source of one Dag.
2324
"""
2425

2526
from __future__ import annotations
@@ -44,8 +45,8 @@
4445
_MAX_LAYOUT_LINE_BYTES = 4096
4546
_METADATA_COMMENT_PREFIX = b"//# airflowMetadata="
4647
_MAX_METADATA_LINE_BYTES = 1024 * 1024
47-
# A block comment, because the entrypoint spans more than one line.
48-
_SOURCE_COMMENT_OPEN = b"/*# airflowSource\n"
48+
# Each source is its own block comment; the marker is followed by the region's path and a newline.
49+
_SOURCE_COMMENT_OPEN_PREFIX = b"/*# airflowSource:"
4950
_SOURCE_COMMENT_CLOSE = b"\n#*/\n"
5051
_MAX_SOURCE_BYTES = 1024 * 1024
5152
_SUPPORTED_BUNDLE_MAJOR_VERSION = 1
@@ -67,12 +68,23 @@ class _DeclaredSection:
6768
sha256: bytes
6869

6970

71+
@attrs.define(frozen=True)
72+
class _DeclaredSourceRegion:
73+
"""One source region's author path plus its declared byte range and expected digest."""
74+
75+
path: str
76+
start: int
77+
end: int
78+
sha256: bytes
79+
80+
7081
@attrs.define(frozen=True)
7182
class _BundleLayout:
7283
"""The metadata, source, and executable sections described by the layout comment."""
7384

7485
metadata: _DeclaredSection
75-
source: _DeclaredSection
86+
# One region per native Dag file, in file order; empty for a mixed-language bundle.
87+
sources: tuple[_DeclaredSourceRegion, ...]
7688
code: _DeclaredSection
7789

7890

@@ -82,7 +94,7 @@ class _DigestCacheKey:
8294

8395
path: str
8496
metadata: _DeclaredSection
85-
source: _DeclaredSection
97+
sources: tuple[_DeclaredSourceRegion, ...]
8698
code: _DeclaredSection
8799
device: int
88100
inode: int
@@ -96,7 +108,7 @@ class _ComputedDigests:
96108
"""SHA-256 digests calculated from the actual section bytes."""
97109

98110
metadata: bytes
99-
source: bytes
111+
sources: tuple[bytes, ...]
100112
code: bytes
101113

102114

@@ -119,17 +131,37 @@ def read_bundle(bundle_path: pathlib.Path) -> BundleMetadata:
119131
return _parse_bundle_metadata(_read_verified_payloads(bundle_path).metadata)
120132

121133

122-
def read_bundle_source(bundle_path: pathlib.Path) -> str:
123-
"""Return the entrypoint source ``airflow-ts-pack`` embedded in *bundle_path*."""
124-
return _decode_source(_read_verified_payloads(bundle_path).source)
134+
def read_bundle_source(bundle_path: pathlib.Path, dag_id: str | None = None) -> str | None:
135+
"""
136+
Return the author source ``airflow-ts-pack`` embedded for *dag_id*, or ``None``.
137+
138+
Pass *dag_id* only for a Dag the bundle defines in TypeScript. Omit it for a Dag owned by
139+
another language (Python), which shows its own source; the result is then ``None`` rather than
140+
an unrelated TypeScript file. When given, a Dag mapped in ``dag_source_paths`` returns its own
141+
file; a Dag that is not (one constructed dynamically, which the packer could not attribute)
142+
falls back to the always-embedded entrypoint source.
143+
"""
144+
if dag_id is None:
145+
return None
146+
payloads = _read_verified_payloads(bundle_path)
147+
source_path = _resolve_source_path(payloads.metadata, dag_id)
148+
if source_path is None:
149+
return None
150+
try:
151+
payload = payloads.sources[source_path]
152+
except KeyError:
153+
raise ValueError(
154+
f"bundle declares no source region at path {source_path!r} for dag_id {dag_id!r}"
155+
) from None
156+
return _decode_source(payload)
125157

126158

127159
@attrs.define(frozen=True)
128160
class _VerifiedPayloads:
129-
"""The metadata and source payloads of a bundle whose digests have been checked."""
161+
"""The metadata and per-path source payloads of a bundle whose digests have been checked."""
130162

131163
metadata: bytes
132-
source: bytes
164+
sources: dict[str, bytes]
133165

134166

135167
def _read_verified_payloads(bundle_path: pathlib.Path) -> _VerifiedPayloads:
@@ -198,6 +230,17 @@ def _parse_offset(section: dict[str, Any], field: str) -> int:
198230
return int(value, 16)
199231

200232

233+
def _parse_digest(section: dict[str, Any], name: str) -> bytes:
234+
sha256 = section.get("sha256")
235+
if (
236+
not isinstance(sha256, str)
237+
or len(sha256) != 64
238+
or any(character not in _LOWER_HEX_DIGITS for character in sha256)
239+
):
240+
raise ValueError(f"bundle layout {name}.sha256 must be 64 lowercase hexadecimal digits")
241+
return bytes.fromhex(sha256)
242+
243+
201244
def _parse_section(layout: dict[str, Any], name: str) -> _DeclaredSection:
202245
section = layout.get(name)
203246
if not isinstance(section, dict):
@@ -206,14 +249,32 @@ def _parse_section(layout: dict[str, Any], name: str) -> _DeclaredSection:
206249
end = _parse_offset(section, "end")
207250
if start >= end:
208251
raise ValueError(f"bundle layout {name} section must contain at least one byte")
209-
sha256 = section.get("sha256")
210-
if (
211-
not isinstance(sha256, str)
212-
or len(sha256) != 64
213-
or any(character not in _LOWER_HEX_DIGITS for character in sha256)
214-
):
215-
raise ValueError(f"bundle layout {name}.sha256 must be 64 lowercase hexadecimal digits")
216-
return _DeclaredSection(start=start, end=end, sha256=bytes.fromhex(sha256))
252+
return _DeclaredSection(start=start, end=end, sha256=_parse_digest(section, name))
253+
254+
255+
def _parse_source_region(entry: Any) -> _DeclaredSourceRegion:
256+
if not isinstance(entry, dict):
257+
raise ValueError("bundle layout sources entries must be mappings")
258+
path = entry.get("path")
259+
if not isinstance(path, str) or not path:
260+
raise ValueError("bundle layout source path must be a non-empty string")
261+
start = _parse_offset(entry, "start")
262+
end = _parse_offset(entry, "end")
263+
if start >= end:
264+
raise ValueError("bundle layout source section must contain at least one byte")
265+
return _DeclaredSourceRegion(path=path, start=start, end=end, sha256=_parse_digest(entry, "source"))
266+
267+
268+
def _parse_sources(layout: dict[str, Any]) -> tuple[_DeclaredSourceRegion, ...]:
269+
# A native-TypeScript bundle declares one region per Dag file; a mixed-language bundle that only
270+
# ships handlers for Dags authored elsewhere declares an empty list. A missing key is neither.
271+
sources = layout.get("sources")
272+
if not isinstance(sources, list):
273+
raise ValueError("bundle layout is missing the sources section")
274+
regions = tuple(_parse_source_region(entry) for entry in sources)
275+
if len({region.path for region in regions}) != len(regions):
276+
raise ValueError("bundle layout declares duplicate source paths")
277+
return regions
217278

218279

219280
def _parse_layout(payload: bytes) -> _BundleLayout:
@@ -225,7 +286,7 @@ def _parse_layout(payload: bytes) -> _BundleLayout:
225286
raise ValueError("embedded airflow bundle layout must contain a mapping")
226287
return _BundleLayout(
227288
metadata=_parse_section(layout, "metadata"),
228-
source=_parse_section(layout, "source"),
289+
sources=_parse_sources(layout),
229290
code=_parse_section(layout, "code"),
230291
)
231292

@@ -301,8 +362,9 @@ def _compute_stable_digests(
301362
metadata=_hash_region(
302363
bundle_file, start=layout.metadata.start, end=layout.metadata.end, path=path, section="metadata"
303364
),
304-
source=_hash_region(
305-
bundle_file, start=layout.source.start, end=layout.source.end, path=path, section="source"
365+
sources=tuple(
366+
_hash_region(bundle_file, start=region.start, end=region.end, path=path, section="source")
367+
for region in layout.sources
306368
),
307369
code=_hash_region(
308370
bundle_file, start=layout.code.start, end=layout.code.end, path=path, section="code"
@@ -328,7 +390,7 @@ def _verify_integrity(
328390
cache_key = _DigestCacheKey(
329391
path=os.fspath(path),
330392
metadata=layout.metadata,
331-
source=layout.source,
393+
sources=layout.sources,
332394
code=layout.code,
333395
device=initial_file_info.st_dev,
334396
inode=initial_file_info.st_ino,
@@ -344,11 +406,15 @@ def _verify_integrity(
344406
)
345407
_digest_cache.put(cache_key, computed_digests)
346408

347-
for section, computed_digest, declared_digest in (
409+
checks = [
348410
("metadata", computed_digests.metadata, layout.metadata.sha256),
349-
("source", computed_digests.source, layout.source.sha256),
411+
*(
412+
(f"source {region.path}", computed, region.sha256)
413+
for region, computed in zip(layout.sources, computed_digests.sources)
414+
),
350415
("code", computed_digests.code, layout.code.sha256),
351-
):
416+
]
417+
for section, computed_digest, declared_digest in checks:
352418
if computed_digest != declared_digest:
353419
raise ValueError(f"{path.name} {section} SHA-256 mismatch")
354420

@@ -385,6 +451,31 @@ def _parse_bundle_metadata(payload: bytes) -> BundleMetadata:
385451
)
386452

387453

454+
def _resolve_source_path(payload: bytes, dag_id: str) -> str | None:
455+
"""
456+
Return the embedded source path for *dag_id*, or ``None``.
457+
458+
A native TypeScript Dag maps to its own file through ``dag_source_paths``. A Dag that is not
459+
mapped (one constructed dynamically, so the packer could not attribute a file) falls back to
460+
``entrypoint_path``. Absent both, there is nothing to show and the result is ``None``.
461+
"""
462+
try:
463+
metadata = json.loads(payload.decode("utf-8"))
464+
except (UnicodeDecodeError, json.JSONDecodeError, RecursionError) as exc:
465+
raise ValueError(f"cannot parse embedded airflow metadata: {exc}") from exc
466+
if not isinstance(metadata, dict):
467+
raise ValueError("embedded airflow metadata must contain a mapping")
468+
dag_source_paths = metadata.get("dag_source_paths")
469+
if isinstance(dag_source_paths, dict):
470+
mapped = dag_source_paths.get(dag_id)
471+
if isinstance(mapped, str):
472+
return mapped
473+
entrypoint_path = metadata.get("entrypoint_path")
474+
if isinstance(entrypoint_path, str):
475+
return entrypoint_path
476+
return None
477+
478+
388479
def _read_bundle_headers(
389480
bundle_file: BinaryIO, *, path: pathlib.Path, file_size: int
390481
) -> tuple[_BundleLayout, _VerifiedPayloads]:
@@ -411,37 +502,44 @@ def _read_bundle_headers(
411502
metadata_end = metadata_start + len(metadata_payload)
412503
if (layout.metadata.start, layout.metadata.end) != (metadata_start, metadata_end):
413504
raise ValueError("bundle layout metadata offsets do not match the metadata section")
414-
source_payload = _read_source_region(
415-
bundle_file, path=path, start=metadata_end + 1 + len(_SOURCE_COMMENT_OPEN), layout=layout
416-
)
417-
code_start = layout.source.end + len(_SOURCE_COMMENT_CLOSE)
505+
# The source regions follow the metadata line's terminating newline, back to back.
506+
sources, code_start = _read_source_regions(bundle_file, path=path, start=metadata_end + 1, layout=layout)
418507
if (layout.code.start, layout.code.end) != (code_start, file_size):
419508
raise ValueError("bundle layout code offsets do not match the executable section")
420-
return layout, _VerifiedPayloads(metadata=metadata_payload, source=source_payload)
509+
return layout, _VerifiedPayloads(metadata=metadata_payload, sources=sources)
421510

422511

423-
def _read_source_region(
512+
def _read_source_regions(
424513
bundle_file: BinaryIO, *, path: pathlib.Path, start: int, layout: _BundleLayout
425-
) -> bytes:
426-
"""Read the source block comment and check that it frames the declared range exactly."""
427-
# The length is declared, not derivable, so the comment markers pin the range to real bytes.
428-
if layout.source.start != start:
429-
raise ValueError("bundle layout source offsets do not match the source section")
430-
length = layout.source.end - start
431-
if length > _MAX_SOURCE_BYTES:
432-
raise ValueError(f"embedded airflow source exceeds {_MAX_SOURCE_BYTES} bytes")
433-
try:
434-
opener = bundle_file.read(len(_SOURCE_COMMENT_OPEN))
435-
payload = bundle_file.read(length)
436-
closer = bundle_file.read(len(_SOURCE_COMMENT_CLOSE))
437-
except OSError as exc:
438-
raise OSError(f"cannot read {path.name}: {exc}") from exc
439-
if opener != _SOURCE_COMMENT_OPEN:
440-
raise ValueError(f"{path.name} has no embedded airflow source after its metadata")
441-
if len(payload) != length or closer != _SOURCE_COMMENT_CLOSE:
442-
raise ValueError("embedded airflow source is not closed by its block comment")
443-
# An unescaped ``*/`` would end the comment where Node reads the file,
444-
# so bytes the layout calls source would execute while both digests still match.
445-
if b"*/" in payload:
446-
raise ValueError("embedded airflow source contains an unescaped block comment terminator")
447-
return payload
514+
) -> tuple[dict[str, bytes], int]:
515+
"""Read every source block comment in order and return the payloads and the code start offset."""
516+
payloads: dict[str, bytes] = {}
517+
cursor = start
518+
for region in layout.sources:
519+
# The marker carries the region's own path, so a bundle cannot silently swap regions.
520+
open_marker = _SOURCE_COMMENT_OPEN_PREFIX + region.path.encode("utf-8") + b"\n"
521+
# The length is declared, not derivable, so the comment markers pin the range to real bytes.
522+
if region.start != cursor + len(open_marker):
523+
raise ValueError("bundle layout source offsets do not match the source section")
524+
length = region.end - region.start
525+
if length > _MAX_SOURCE_BYTES:
526+
raise ValueError(f"embedded airflow source exceeds {_MAX_SOURCE_BYTES} bytes")
527+
try:
528+
opener = bundle_file.read(len(open_marker))
529+
payload = bundle_file.read(length)
530+
closer = bundle_file.read(len(_SOURCE_COMMENT_CLOSE))
531+
except OSError as exc:
532+
raise OSError(f"cannot read {path.name}: {exc}") from exc
533+
if not opener.startswith(_SOURCE_COMMENT_OPEN_PREFIX):
534+
raise ValueError(f"{path.name} has no embedded airflow source after its metadata")
535+
if opener != open_marker:
536+
raise ValueError("embedded airflow source marker does not match its declared path")
537+
if len(payload) != length or closer != _SOURCE_COMMENT_CLOSE:
538+
raise ValueError("embedded airflow source is not closed by its block comment")
539+
# An unescaped ``*/`` would end the comment where Node reads the file,
540+
# so bytes the layout calls source would execute while both digests still match.
541+
if b"*/" in payload:
542+
raise ValueError("embedded airflow source contains an unescaped block comment terminator")
543+
payloads[region.path] = payload
544+
cursor = region.end + len(_SOURCE_COMMENT_CLOSE)
545+
return payloads, cursor

0 commit comments

Comments
 (0)