Extract shared bundle metadata logic from coordinators

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

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),
     )
 

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 = ""

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),
     )
 

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,
         )