Copilot commented on code in PR #70805:
URL: https://github.com/apache/airflow/pull/70805#discussion_r3794989862
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -385,23 +432,146 @@ class SubprocessCoordinator(BaseCoordinator):
:param task_startup_timeout: Maximum time the coordinator waits for the
subprocess to connect to both servers, in seconds. The default is 10
seconds.
+ :param dag_bundle_name: Locate artifacts through a configured Dag bundle
rather
+ than an explicit root. Mutually exclusive with the subclass's explicit
root;
+ if neither is set, the task's own bundle is used. A named bundle
resolves to
+ the version current when the task starts; the task's own bundle uses
the
+ run's version. Either way the resolved version is pinned for the whole
task.
"""
task_startup_timeout: float = 10.0
+ dag_bundle_name: str | None = None
+
+ _artifact_source: _ArtifactSource = attrs.field(init=False)
+ # The subclass's explicit root, recorded at construction so the base can
+ # resolve roots without knowing the subclass field name.
+ _configured_roots: list[pathlib.Path] = attrs.field(init=False,
factory=list)
+ _active_bundle_info: BundleInfo | None = attrs.field(init=False,
default=None)
+ _active_scan_roots: tuple[pathlib.Path, ...] | None =
attrs.field(init=False, default=None)
+
+ @property
+ def _explicit_artifact_roots(self) -> tuple[str, Sequence[pathlib.Path]]:
+ """
+ The subclass's explicit-root kwarg name and its configured value.
+
+ The name is only used in error messages. An empty value — the default,
for a
+ subclass that does not override this — selects task-bundle mode rather
than
+ failing at execute time.
+ """
+ return "root", ()
+
+ def __attrs_post_init__(self) -> None:
+ self._classify_artifact_source()
+
+ def _classify_artifact_source(self) -> None:
+ """
+ Classify and validate how this coordinator locates artifacts
(construction time).
+
+ Rejects setting both an explicit root and ``dag_bundle_name``, fails
fast
+ when ``dag_bundle_name`` names a bundle that is not configured, and
records
+ the resulting :class:`_ArtifactSource` and explicit root.
+ """
+ root_kwarg, configured = self._explicit_artifact_roots
+ if configured and self.dag_bundle_name is not None:
+ raise ValueError(
+ f"Set at most one of {root_kwarg!r} or 'dag_bundle_name':
{root_kwarg!r} for an "
+ f"explicit path, 'dag_bundle_name' for a configured Dag
bundle, or leave both "
+ f"unset to scan the task's own bundle."
+ )
+ if configured:
+ source = _ArtifactSource.EXPLICIT_ROOT
+ self._configured_roots = list(configured)
+ elif self.dag_bundle_name is not None:
+ source = _ArtifactSource.NAMED_BUNDLE
+ if not
DagBundlesManager.is_bundle_configured(self.dag_bundle_name):
+ raise ValueError(
+ f"Coordinator 'dag_bundle_name' references unconfigured
Dag bundle "
+ f"{self.dag_bundle_name!r}."
+ )
+ else:
+ source = _ArtifactSource.TASK_BUNDLE
+
+ self._artifact_source = source
+ log.debug(
+ "Coordinator artifact source selected",
+ mode=source.name,
+ dag_bundle_name=self.dag_bundle_name,
+ configured_roots=[str(root) for root in self._configured_roots],
+ )
+
+ def _init_root_source(
+ self, logger: FilteringBoundLogger
+ ) -> tuple[list[pathlib.Path], BaseDagBundle | None]:
+ """
+ Resolve the directories to scan for artifacts, dispatched on the
classified mode.
+
+ Returns ``(roots, bundle)``: an explicit root yields no bundle
(``None``);
+ a Dag-bundle mode returns the materialized path and the resolved
bundle so
+ :meth:`execute_task` can hold a version lock over it. *logger* is the
task
+ logger, so materialization failures surface in the task log.
+ """
+ if self._artifact_source is _ArtifactSource.EXPLICIT_ROOT:
+ return self._configured_roots, None
+
+ if self._artifact_source is _ArtifactSource.NAMED_BUNDLE:
+ # NAMED_BUNDLE implies dag_bundle_name is set.
+ target = BundleInfo(name=cast("str", self.dag_bundle_name))
+ else:
Review Comment:
`typing.cast()` expects a type object as its first argument. Passing the
string "str" defeats static checking and will likely fail mypy/ruff-type rules;
here `dag_bundle_name` is known to be non-None in the NAMED_BUNDLE branch, so
cast should use `str` (the type), not "str".
##########
task-sdk/src/airflow/sdk/execution_time/bundles.py:
##########
@@ -0,0 +1,92 @@
+# 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.
+"""
+Materialize a task instance's Dag bundle on the worker.
+
+Kept apart from :mod:`airflow.sdk.execution_time.task_runner` so the supervisor
+side — which needs a bundle on disk before it can launch a language-SDK
+subprocess — does not have to import the task runner and everything it pulls
in.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING
+
+from airflow.dag_processing.bundles.manager import DagBundlesManager # noqa:
SDK002
+from airflow.sdk.execution_time.tracing import detail_span
+
+if TYPE_CHECKING:
+ from airflow.dag_processing.bundles.base import BaseDagBundle # noqa:
SDK002
+ from airflow.sdk.api.datamodels._generated import BundleInfo
+
+__all__ = ["initialize_ti_bundle", "verify_bundle_access"]
+
+
+def initialize_ti_bundle(bundle_info: BundleInfo) -> BaseDagBundle:
+ """
+ Resolve, initialize, and access-check the Dag bundle for a task instance.
+
+ Shared by :func:`~airflow.sdk.execution_time.task_runner.parse` (Python
task
+ path) and the subprocess coordinators (language-SDK path), which both need
a
+ task instance's bundle materialized on disk before use. Returns the
+ initialized bundle so callers can read ``bundle.path``.
+ """
+ bundle_instance = DagBundlesManager().get_bundle(
+ name=bundle_info.name,
+ version=bundle_info.version,
+ version_data=bundle_info.version_data,
+ )
+ bundle_instance.initialize()
+ verify_bundle_access(bundle_instance)
+ return bundle_instance
+
+
+@detail_span("verify_bundle_access")
+def verify_bundle_access(bundle_instance: BaseDagBundle) -> None:
+ """
+ Verify bundle is accessible by the current user.
+
+ This is called after user impersonation (if any) to ensure the bundle
+ is actually accessible. Uses os.access() which works with any permission
+ scheme (standard Unix permissions, ACLs, SELinux, etc.).
+
+ :param bundle_instance: The bundle instance to check
+ :raises AirflowException: if bundle is not accessible
+ """
+ from getpass import getuser
+
+ from airflow.sdk.exceptions import AirflowException
+
Review Comment:
Imports inside `verify_bundle_access()` (e.g. `getuser`, `AirflowException`)
violate the project rule that imports should be at module top-level unless
there is a specific need (circular import avoidance / lazy loading /
TYPE_CHECKING). These imports look safe to hoist, and keeping them inside the
function makes linting and dependency analysis harder.
--
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]
