diff --git a/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py b/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py new file mode 100644 index 0000000000000..e95ffe52396f0 --- /dev/null +++ b/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py @@ -0,0 +1,82 @@ +# +# 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. +"""Bundle-metadata helpers shared by the subprocess coordinators.""" + +from __future__ import annotations + +import os +import pathlib +from typing import Any + +import attrs +import yaml + +from airflow.sdk.execution_time.schema import get_schema_version_migrator + + +def convert_roots( + value: None | os.PathLike[str] | pathlib.Path | list[os.PathLike[str] | pathlib.Path], +) -> list[pathlib.Path]: + """Normalize a coordinator's root-directories kwarg into a list of expanded paths.""" + if value is None: + return [] + if isinstance(value, (str, os.PathLike, pathlib.Path)): + return [pathlib.Path(value).expanduser()] + return [pathlib.Path(v).expanduser() for v in value] + + +def validate_schema_version(instance, _, value) -> str: + """Attrs validator resolving a bundle's supervisor schema version to a known one.""" + return get_schema_version_migrator().resolve_version(str(value)) + + +@attrs.define +class ResolvedBundle: + """A located bundle whose supervisor schema version has been resolved.""" + + path: pathlib.Path + schema_version: str = attrs.field(validator=validate_schema_version) + + +def parse_metadata_mapping(content: str | bytes, *, source: str) -> dict[str, Any]: + """ + Parse *content* as the ``airflow-metadata.yaml`` mapping. + + Raises ``ValueError`` on undecodable, unparsable, or non-mapping content; + *source* names the metadata's origin in the error message. + """ + try: + data = yaml.safe_load(content.decode("utf-8") if isinstance(content, bytes) else content) + except (UnicodeDecodeError, yaml.YAMLError) as exc: + raise ValueError(f"cannot parse {source}: {exc}") from exc + + if not isinstance(data, dict): + raise ValueError(f"{source} must contain a mapping") + return data + + +def extract_supervisor_schema_version(metadata: dict[str, Any]) -> str: + """Return ``sdk.supervisor_schema_version`` from bundle metadata, raising ``ValueError`` if absent.""" + sdk = metadata.get("sdk") + if not isinstance(sdk, dict): + raise ValueError("missing sdk metadata mapping") + + value = sdk.get("supervisor_schema_version") + if not isinstance(value, str) or not value: + raise ValueError("missing or invalid sdk.supervisor_schema_version") + return value diff --git a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py index 2275628d0b1c8..70ba5490450c3 100644 --- a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py @@ -29,10 +29,14 @@ import attrs import structlog -import yaml +from airflow.sdk.coordinators._bundle_metadata import ( + ResolvedBundle, + convert_roots, + extract_supervisor_schema_version, + parse_metadata_mapping, +) from airflow.sdk.coordinators._subprocess import SubprocessCoordinator -from airflow.sdk.execution_time.schema import get_schema_version_migrator if TYPE_CHECKING: from collections.abc import Iterable, Iterator, Sequence @@ -236,21 +240,11 @@ def _read_bundle_metadata(path: pathlib.Path) -> dict[str, Any] | None: return None try: - data = yaml.safe_load(metadata_bytes.decode("utf-8")) - except (UnicodeDecodeError, yaml.YAMLError) as exc: + return parse_metadata_mapping(metadata_bytes, source="bundle metadata") + except ValueError as exc: log.debug("Cannot decode bundle metadata; skipping", path=str(path), error=str(exc)) return None - if not isinstance(data, dict): - log.debug( - "Bundle metadata is not a mapping; skipping", - path=str(path), - type=type(data).__name__, - ) - return None - - return data - def _dag_ids(metadata: dict[str, Any]) -> set[str]: dags = metadata.get("dags") @@ -260,18 +254,6 @@ def _dag_ids(metadata: dict[str, Any]) -> set[str]: return set(dags.keys()) -def _supervisor_schema_version(metadata: dict[str, Any]) -> str | None: - sdk = metadata.get("sdk") - if not isinstance(sdk, dict): - return None - - value = sdk.get("supervisor_schema_version") - if not isinstance(value, str) or not value: - return None - - return value - - def _find_executables(items: Iterable[pathlib.Path]) -> Iterator[pathlib.Path]: """ Yield executable regular files under *items*, descending into directories. @@ -308,15 +290,8 @@ def _walk_executables( yield item -def _validate_schema_version(instance, _, value) -> str: - return get_schema_version_migrator().resolve_version(str(value)) - - @attrs.define -class _Bundle: - path: pathlib.Path - schema_version: str = attrs.field(validator=_validate_schema_version) - +class _Bundle(ResolvedBundle): @classmethod def find(cls, executables_root: Sequence[pathlib.Path], dag_id: str) -> Self: log.debug("Finding executable bundles recursively", roots=executables_root) @@ -328,12 +303,7 @@ def find(cls, executables_root: Sequence[pathlib.Path], dag_id: str) -> Self: continue try: - if (schema_version := _supervisor_schema_version(metadata)) is None: - reason = "missing or invalid sdk.supervisor_schema_version" - log.debug("Bundle metadata rejected; skipping", path=str(p), error=reason) - rejected.append((p.resolve(), reason)) - continue - return cls(path=p.resolve(), schema_version=schema_version) + return cls(path=p.resolve(), schema_version=extract_supervisor_schema_version(metadata)) except (TypeError, ValueError) as exc: log.debug("Bundle metadata rejected; skipping", path=str(p), error=str(exc)) rejected.append((p.resolve(), str(exc))) @@ -352,16 +322,6 @@ def find(cls, executables_root: Sequence[pathlib.Path], dag_id: str) -> Self: raise FileNotFoundError(tp.format(dag_id, resolved_paths, details)) -def _convert_executables_root( - value: None | os.PathLike[str] | pathlib.Path | list[os.PathLike[str] | pathlib.Path], -) -> list[pathlib.Path]: - if value is None: - return [] - if isinstance(value, (str, os.PathLike, pathlib.Path)): - return [pathlib.Path(value).expanduser()] - return [pathlib.Path(v).expanduser() for v in value] - - @attrs.define(kw_only=True) class ExecutableCoordinator(SubprocessCoordinator): """ @@ -385,7 +345,7 @@ class ExecutableCoordinator(SubprocessCoordinator): """ executables_root: list[pathlib.Path] = attrs.field( - converter=_convert_executables_root, + converter=convert_roots, validator=attrs.validators.min_len(1), ) diff --git a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py index da35d0f37b84e..0ebe089162ba1 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -29,8 +29,8 @@ import attrs import structlog +from airflow.sdk.coordinators._bundle_metadata import convert_roots, validate_schema_version from airflow.sdk.coordinators._subprocess import SubprocessCoordinator -from airflow.sdk.execution_time.schema import get_schema_version_migrator if TYPE_CHECKING: from collections.abc import Iterable, Iterator, Sequence @@ -109,14 +109,10 @@ def from_jar(cls, path: pathlib.Path) -> Self | None: return None -def _validate_schema_version(instance, _, value) -> str: - return get_schema_version_migrator().resolve_version(str(value)) - - @attrs.define class _JarInfo: main_class: str - schema_version: str = attrs.field(validator=_validate_schema_version) + schema_version: str = attrs.field(validator=validate_schema_version) @attrs.define class _Progress: @@ -156,16 +152,6 @@ def find(cls, roots: Sequence[pathlib.Path], main_class: str) -> _JarInfo: raise FileNotFoundError(tp.format(main_class, os.pathsep.join(os.fspath(p.resolve()) for p in roots))) -def _convert_jars_root( - value: None | os.PathLike[str] | pathlib.Path | list[os.PathLike[str] | pathlib.Path], -) -> list[pathlib.Path]: - if value is None: - return [] - if isinstance(value, (str, os.PathLike, pathlib.Path)): - return [pathlib.Path(value).expanduser()] - return [pathlib.Path(v).expanduser() for v in value] - - @attrs.define(kw_only=True) class JavaCoordinator(SubprocessCoordinator): """ @@ -213,7 +199,7 @@ class JavaCoordinator(SubprocessCoordinator): java_executable: str = "java" jvm_args: list[str] = attrs.field(factory=list) jars_root: list[pathlib.Path] = attrs.field( - converter=_convert_jars_root, + converter=convert_roots, validator=attrs.validators.min_len(1), ) main_class: str = "" diff --git a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py index c1b23615b89a1..41b07710f8f6f 100644 --- a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py @@ -25,10 +25,14 @@ import attrs import structlog -import yaml +from airflow.sdk.coordinators._bundle_metadata import ( + ResolvedBundle, + convert_roots, + extract_supervisor_schema_version, + parse_metadata_mapping, +) from airflow.sdk.coordinators._subprocess import SubprocessCoordinator -from airflow.sdk.execution_time.schema import get_schema_version_migrator if TYPE_CHECKING: from collections.abc import Sequence @@ -43,45 +47,19 @@ METADATA_FILENAME = "airflow-metadata.yaml" -def _validate_schema_version(instance, _, value) -> str: - return get_schema_version_migrator().resolve_version(str(value)) - - -@attrs.define -class _NodeBundle: - path: pathlib.Path - schema_version: str = attrs.field(validator=_validate_schema_version) - - def _read_bundle_metadata(metadata_path: pathlib.Path) -> dict[str, Any]: if not metadata_path.is_file(): raise ValueError(f"missing {METADATA_FILENAME}") try: - with metadata_path.open(encoding="utf-8") as metadata_file: - data = yaml.safe_load(metadata_file) + content = metadata_path.read_text(encoding="utf-8") except OSError as exc: raise ValueError(f"cannot read {METADATA_FILENAME}: {exc}") from exc - except yaml.YAMLError as exc: - raise ValueError(f"cannot parse {METADATA_FILENAME}: {exc}") from exc - - if not isinstance(data, dict): - raise ValueError(f"{METADATA_FILENAME} must contain a mapping") - return data + return parse_metadata_mapping(content, source=METADATA_FILENAME) -def _supervisor_schema_version(metadata: dict[str, Any]) -> str: - sdk = metadata.get("sdk") - if not isinstance(sdk, dict): - raise ValueError("missing sdk metadata mapping") - value = sdk.get("supervisor_schema_version") - if not isinstance(value, str) or not value: - raise ValueError("missing or invalid sdk.supervisor_schema_version") - return value - - -def _find_bundle(bundles_root: Sequence[pathlib.Path]) -> _NodeBundle: +def _find_bundle(bundles_root: Sequence[pathlib.Path]) -> ResolvedBundle: """ Locate the ``.mjs`` entry point in *bundles_root*. @@ -101,9 +79,9 @@ def _find_bundle(bundles_root: Sequence[pathlib.Path]) -> _NodeBundle: try: metadata = _read_bundle_metadata(root / METADATA_FILENAME) log.debug("Selected TypeScript bundle", path=candidate, root=root) - return _NodeBundle( + return ResolvedBundle( path=candidate, - schema_version=_supervisor_schema_version(metadata), + schema_version=extract_supervisor_schema_version(metadata), ) except (TypeError, ValueError) as exc: log.debug( @@ -123,16 +101,6 @@ def _find_bundle(bundles_root: Sequence[pathlib.Path]) -> _NodeBundle: raise FileNotFoundError(f"Cannot find {BUNDLE_FILENAME} in {searched}") -def _convert_bundles_root( - value: None | os.PathLike[str] | pathlib.Path | list[os.PathLike[str] | pathlib.Path], -) -> list[pathlib.Path]: - if value is None: - return [] - if isinstance(value, (str, os.PathLike, pathlib.Path)): - return [pathlib.Path(value).expanduser()] - return [pathlib.Path(v).expanduser() for v in value] - - @attrs.define(kw_only=True) class NodeCoordinator(SubprocessCoordinator): """ @@ -163,7 +131,7 @@ class NodeCoordinator(SubprocessCoordinator): node_executable: str = "node" bundles_root: list[pathlib.Path] = attrs.field( - converter=_convert_bundles_root, + converter=convert_roots, validator=attrs.validators.min_len(1), ) diff --git a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py index ec426a09ed741..5963e5f4892a7 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -354,9 +354,9 @@ def test_logs_when_metadata_is_not_a_mapping(self, tmp_path): _Bundle.find([tmp_path], "tutorial_dag") mock_log.debug.assert_any_call( - "Bundle metadata is not a mapping; skipping", + "Cannot decode bundle metadata; skipping", path=str(bundle_path), - type=mock.ANY, + error=mock.ANY, )