jason810496 commented on code in PR #72046: URL: https://github.com/apache/airflow/pull/72046#discussion_r3974978467
########## task-sdk/docs/ts-bundle-spec.rst: ########## @@ -0,0 +1,177 @@ + .. 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. + +TypeScript Bundle Format +======================== + +This document specifies the ``bundle.mjs`` format produced by +``airflow-ts-pack`` and consumed by +:class:`~airflow.sdk.coordinators.node.NodeCoordinator`. + +Container +--------- + +The bundle remains an ECMAScript module that runs directly with +``node bundle.mjs``. It has three regions: Review Comment: ```suggestion This document specifies the ``bundle.mjs`` format produced by ``airflow-ts-pack`` and consumed by :class:`~airflow.sdk.coordinators.node.NodeCoordinator`. Container --------- The bundle remains an ECMAScript module that runs directly with ``node bundle.mjs``. It has three regions: ``` ########## task-sdk/tests/task_sdk/coordinators/node/test_bundle_reader.py: ########## @@ -0,0 +1,417 @@ +# +# 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. + +from __future__ import annotations + +import base64 +import io +import json +import os +import pathlib +from unittest import mock + +import pytest +from task_sdk.coordinators.node._bundle_test_utils import ( + METADATA_PREFIX, + OFFSET_WIDTH, + SCHEMA_VERSION, + metadata_json as _metadata_json, + mutate_byte as _mutate_byte, + read_layout as _read_layout, + replace_layout_payload as _replace_layout_payload, + rewrite_layout as _rewrite_layout, + write_bundle, +) + +from airflow.sdk.coordinators.node import _bundle_reader as _reader +from airflow.sdk.coordinators.node._bundle_reader import _digest_cache, _hash_region, read_bundle + +from tests_common.test_utils.paths import AIRFLOW_ROOT_PATH + +TYPESCRIPT_V1_FIXTURE = AIRFLOW_ROOT_PATH / "ts-sdk" / "tests" / "cli" / "fixtures" / "bundle-v1.mjs" + + [email protected](autouse=True) +def clear_digest_cache(): + _digest_cache.clear() + + +class TestBundleReader: + def test_reader_returns_verified_bundle_metadata(self): + metadata = read_bundle(TYPESCRIPT_V1_FIXTURE) + + assert metadata.dag_ids == frozenset({"test_dag"}) + assert metadata.supervisor_schema_version == SCHEMA_VERSION + + def test_reads_bundle_produced_by_typescript_encoder(self, tmp_path): + bundle = tmp_path / "bundle.mjs" + bundle.write_bytes(TYPESCRIPT_V1_FIXTURE.read_bytes()) + + metadata = read_bundle(bundle) + + assert metadata.dag_ids == frozenset({"test_dag"}) + assert metadata.supervisor_schema_version == SCHEMA_VERSION Review Comment: Let's combine the test. ########## task-sdk/src/airflow/sdk/coordinators/node/coordinator.py: ########## @@ -45,80 +40,46 @@ log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.node") BUNDLE_FILENAME = "bundle.mjs" -EMBEDDED_METADATA_MARKER = b"//# airflowMetadata=" -EMBEDDED_METADATA_MAX_BYTES = 1024 * 1024 -def _read_embedded_metadata(bundle_path: pathlib.Path) -> dict[str, Any]: - """ - Read the manifest ``airflow-ts-pack`` embeds in the bundle itself. - - The packer prepends the metadata as a leading - ``//# airflowMetadata=<base64>`` line comment, keeping bundle and metadata - a single artifact. Raises ``ValueError`` when the bundle has no such marker. - """ - try: - with bundle_path.open("rb") as bundle_file: - line = bundle_file.readline(EMBEDDED_METADATA_MAX_BYTES + 1) - except OSError as exc: - raise ValueError(f"cannot read {bundle_path.name}: {exc}") from exc - - if not line.startswith(EMBEDDED_METADATA_MARKER): - raise ValueError(f"{bundle_path.name} has no embedded airflow metadata; rebuild with airflow-ts-pack") - if len(line) > EMBEDDED_METADATA_MAX_BYTES: - raise ValueError( - f"embedded airflow metadata exceeds {EMBEDDED_METADATA_MAX_BYTES} bytes; " - f"rebuild {bundle_path.name} with airflow-ts-pack" - ) - - payload = line[len(EMBEDDED_METADATA_MARKER) :].strip() - try: - decoded = base64.b64decode(payload, validate=True) - except ValueError as exc: - raise ValueError(f"cannot parse embedded airflow metadata: {exc}") from exc - return parse_metadata_mapping(decoded, source="embedded airflow metadata") - - -def _find_bundle(bundles_root: Sequence[pathlib.Path]) -> ResolvedBundle: - """ - Locate the ``.mjs`` entry point in *bundles_root*. - - Scans each configured directory for ``bundle.mjs`` and reads the bundle's - supervisor schema version from the metadata embedded in the bundle. - - This is an ordered fallback search, not Dag/task-aware multi-bundle - routing. The first bundle found wins. A future version can use the - metadata's ``dags`` section together with ``TaskInstance.dag_id`` and - ``TaskInstance.task_id`` to select the bundle that owns a specific task. - """ +def _select_bundle(bundles_root: Sequence[pathlib.Path], dag_id: str) -> ResolvedBundle: + """Return the first verified configured bundle that declares *dag_id*.""" rejected: list[tuple[pathlib.Path, str]] = [] for root in bundles_root: candidate = root / BUNDLE_FILENAME - if not candidate.is_file(): - continue try: - metadata = _read_embedded_metadata(candidate) - log.debug("Selected TypeScript bundle", path=candidate, root=root) - return ResolvedBundle( - path=candidate, - schema_version=extract_supervisor_schema_version(metadata), - ) - except (TypeError, ValueError) as exc: + if not candidate.is_file(): + continue + metadata = read_bundle(candidate) + if dag_id not in metadata.dag_ids: + log.debug( + "TypeScript bundle does not contain requested Dag; skipping", + path=candidate, + root=root, + dag_id=dag_id, + ) + continue Review Comment: Bundles that pass full integrity verification but declare a different `dag_id` are dropped without being recorded in `rejected`. ```suggestion rejected.append((candidate, f"verified bundle declares dag_ids={sorted(metadata.dag_ids)!r}")) continue ``` ########## task-sdk/src/airflow/sdk/coordinators/node/coordinator.py: ########## @@ -45,80 +40,46 @@ log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.node") BUNDLE_FILENAME = "bundle.mjs" -EMBEDDED_METADATA_MARKER = b"//# airflowMetadata=" -EMBEDDED_METADATA_MAX_BYTES = 1024 * 1024 -def _read_embedded_metadata(bundle_path: pathlib.Path) -> dict[str, Any]: - """ - Read the manifest ``airflow-ts-pack`` embeds in the bundle itself. - - The packer prepends the metadata as a leading - ``//# airflowMetadata=<base64>`` line comment, keeping bundle and metadata - a single artifact. Raises ``ValueError`` when the bundle has no such marker. - """ - try: - with bundle_path.open("rb") as bundle_file: - line = bundle_file.readline(EMBEDDED_METADATA_MAX_BYTES + 1) - except OSError as exc: - raise ValueError(f"cannot read {bundle_path.name}: {exc}") from exc - - if not line.startswith(EMBEDDED_METADATA_MARKER): - raise ValueError(f"{bundle_path.name} has no embedded airflow metadata; rebuild with airflow-ts-pack") - if len(line) > EMBEDDED_METADATA_MAX_BYTES: - raise ValueError( - f"embedded airflow metadata exceeds {EMBEDDED_METADATA_MAX_BYTES} bytes; " - f"rebuild {bundle_path.name} with airflow-ts-pack" - ) - - payload = line[len(EMBEDDED_METADATA_MARKER) :].strip() - try: - decoded = base64.b64decode(payload, validate=True) - except ValueError as exc: - raise ValueError(f"cannot parse embedded airflow metadata: {exc}") from exc - return parse_metadata_mapping(decoded, source="embedded airflow metadata") - - -def _find_bundle(bundles_root: Sequence[pathlib.Path]) -> ResolvedBundle: - """ - Locate the ``.mjs`` entry point in *bundles_root*. - - Scans each configured directory for ``bundle.mjs`` and reads the bundle's - supervisor schema version from the metadata embedded in the bundle. - - This is an ordered fallback search, not Dag/task-aware multi-bundle - routing. The first bundle found wins. A future version can use the - metadata's ``dags`` section together with ``TaskInstance.dag_id`` and - ``TaskInstance.task_id`` to select the bundle that owns a specific task. - """ Review Comment: No sure would introduce the `_Bundle` attrs be more readable? https://github.com/apache/airflow/blob/1952520a9eedb63c530be89eade6863ea975bb2e/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py#L293-L323 ########## task-sdk/src/airflow/sdk/coordinators/node/coordinator.py: ########## @@ -45,80 +40,46 @@ log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.node") BUNDLE_FILENAME = "bundle.mjs" -EMBEDDED_METADATA_MARKER = b"//# airflowMetadata=" -EMBEDDED_METADATA_MAX_BYTES = 1024 * 1024 -def _read_embedded_metadata(bundle_path: pathlib.Path) -> dict[str, Any]: - """ - Read the manifest ``airflow-ts-pack`` embeds in the bundle itself. - - The packer prepends the metadata as a leading - ``//# airflowMetadata=<base64>`` line comment, keeping bundle and metadata - a single artifact. Raises ``ValueError`` when the bundle has no such marker. - """ - try: - with bundle_path.open("rb") as bundle_file: - line = bundle_file.readline(EMBEDDED_METADATA_MAX_BYTES + 1) - except OSError as exc: - raise ValueError(f"cannot read {bundle_path.name}: {exc}") from exc - - if not line.startswith(EMBEDDED_METADATA_MARKER): - raise ValueError(f"{bundle_path.name} has no embedded airflow metadata; rebuild with airflow-ts-pack") - if len(line) > EMBEDDED_METADATA_MAX_BYTES: - raise ValueError( - f"embedded airflow metadata exceeds {EMBEDDED_METADATA_MAX_BYTES} bytes; " - f"rebuild {bundle_path.name} with airflow-ts-pack" - ) - - payload = line[len(EMBEDDED_METADATA_MARKER) :].strip() - try: - decoded = base64.b64decode(payload, validate=True) - except ValueError as exc: - raise ValueError(f"cannot parse embedded airflow metadata: {exc}") from exc - return parse_metadata_mapping(decoded, source="embedded airflow metadata") - - -def _find_bundle(bundles_root: Sequence[pathlib.Path]) -> ResolvedBundle: - """ - Locate the ``.mjs`` entry point in *bundles_root*. - - Scans each configured directory for ``bundle.mjs`` and reads the bundle's - supervisor schema version from the metadata embedded in the bundle. - - This is an ordered fallback search, not Dag/task-aware multi-bundle - routing. The first bundle found wins. A future version can use the - metadata's ``dags`` section together with ``TaskInstance.dag_id`` and - ``TaskInstance.task_id`` to select the bundle that owns a specific task. - """ +def _select_bundle(bundles_root: Sequence[pathlib.Path], dag_id: str) -> ResolvedBundle: + """Return the first verified configured bundle that declares *dag_id*.""" rejected: list[tuple[pathlib.Path, str]] = [] for root in bundles_root: candidate = root / BUNDLE_FILENAME - if not candidate.is_file(): - continue try: - metadata = _read_embedded_metadata(candidate) - log.debug("Selected TypeScript bundle", path=candidate, root=root) - return ResolvedBundle( - path=candidate, - schema_version=extract_supervisor_schema_version(metadata), - ) - except (TypeError, ValueError) as exc: + if not candidate.is_file(): + continue + metadata = read_bundle(candidate) + if dag_id not in metadata.dag_ids: + log.debug( + "TypeScript bundle does not contain requested Dag; skipping", + path=candidate, + root=root, + dag_id=dag_id, + ) + continue + bundle = ResolvedBundle(path=candidate, schema_version=metadata.supervisor_schema_version) + except (OSError, TypeError, ValueError) as exc: log.debug( - "TypeScript bundle metadata rejected; skipping", + "TypeScript bundle rejected; skipping", path=candidate, root=root, - exc_info=True, + reason=str(exc), Review Comment: Let's keep the `exc_info` ```suggestion reason=str(exc), exc_info=True, ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
