dabla commented on code in PR #67016: URL: https://github.com/apache/airflow/pull/67016#discussion_r4116383365
########## providers/microsoft/azure/docs/bundles/index.rst: ########## @@ -0,0 +1,80 @@ + .. 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. + +Bundles Review Comment: **[nit]** The page title `Bundles` has no underline, so it renders as a plain paragraph instead of a heading. The Amazon provider's `docs/bundles/index.rst` uses: ```rst Bundles ####### ``` Without the underline the document has no top-level title and the `WasbDagBundle` section becomes the only heading on the page. ########## providers/microsoft/azure/src/airflow/providers/microsoft/azure/bundles/wasb.py: ########## @@ -0,0 +1,154 @@ +# 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 os +from functools import cached_property +from pathlib import Path +from urllib.parse import urlsplit, urlunsplit + +import structlog + +from airflow.dag_processing.bundles.base import BaseDagBundle +from airflow.providers.microsoft.azure.hooks.wasb import WasbHook + + +class WasbDagBundle(BaseDagBundle): + """ + WASB Dag bundle - exposes a directory in Azure Blob Storage as a Dag bundle. + + This allows Airflow to load Dags directly from an Azure Blob Storage container. + + :param wasb_conn_id: Airflow connection ID for Azure Blob Storage. Defaults to WasbHook.default_conn_name. + :param container_name: The name of the blob container containing the Dag files. + :param prefix: Optional subdirectory within the container where the Dags are stored. + If empty, Dags are assumed to be at the root of the container. + """ + + supports_versioning = False + + def __init__( + self, + *, + wasb_conn_id: str = WasbHook.default_conn_name, + container_name: str, + prefix: str = "", + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.wasb_conn_id = wasb_conn_id + self.container_name = container_name + self.prefix = prefix + self.wasb_dags_dir: Path = self.base_dir + + log = structlog.get_logger(__name__) + self._log = log.bind( + bundle_name=self.name, + version=self.version, + container_name=self.container_name, + prefix=self.prefix, + wasb_conn_id=self.wasb_conn_id, + ) + + def _initialize(self): + with self.lock(): + if not self.wasb_dags_dir.exists(): + self._log.info("Creating local Dags directory: %s", self.wasb_dags_dir) + os.makedirs(self.wasb_dags_dir) + + if not self.wasb_dags_dir.is_dir(): + raise NotADirectoryError(f"Local Dags path: {self.wasb_dags_dir} is not a directory.") + + if not self.wasb_hook.check_for_container(container_name=self.container_name): + raise ValueError(f"WASB container '{self.container_name}' does not exist.") + + if self.prefix: + if not self.wasb_hook.check_for_prefix( Review Comment: **[nit]** The prefix existence check is not anchored the way the sync now is, so a sibling prefix can make initialization pass while the bundle syncs nothing. `WasbHook.check_for_prefix` is a raw `name_starts_with` match (it delegates to `get_blobs_list`, which calls `walk_blobs(name_starts_with=prefix)`). After the earlier fix, `sync_to_local_dir` lists `f"{prefix.rstrip('/')}/"`. With `prefix="dags"` and a container that only has `dags-archive/…`, the check succeeds, the sync downloads zero blobs, and the bundle silently comes up empty instead of raising the "prefix does not exist" error. `S3Hook.check_for_prefix` appends the delimiter before checking; the cheapest equivalent here is to pass the anchored prefix: ```python if not self.wasb_hook.check_for_prefix( container_name=self.container_name, prefix=f"{self.prefix.rstrip('/')}/", delimiter="/" ): ``` ########## providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/wasb.py: ########## @@ -463,6 +475,99 @@ def download( # TODO: rework the interface as it might also return Awaitable return blob_client.download_blob(offset=offset, length=length, **kwargs) # type: ignore[return-value] + def _sync_to_local_dir_delete_stale_local_files( + self, current_wasb_objects: list[Path], local_dir: Path + ) -> None: + current_wasb_keys = {key.resolve() for key in current_wasb_objects} + + for item in local_dir.rglob("*"): + if item.is_file() and item.resolve() not in current_wasb_keys: + self.log.debug("Deleting stale local file: %s", item) + item.unlink() + for root, dirs, _ in os.walk(local_dir, topdown=False): + for d in dirs: + dir_path = os.path.join(root, d) + if not os.listdir(dir_path): + self.log.debug("Deleting stale empty directory: %s", dir_path) + os.rmdir(dir_path) + + def _sync_to_local_dir_if_changed( + self, container_name: str, blob: BlobProperties, local_target_path: Path + ) -> None: + should_download = False + download_logs: list[str] = [] + download_log_params: list[Any] = [] + + if not local_target_path.exists(): + should_download = True + download_logs.append("Local file %s does not exist.") + download_log_params.append(local_target_path) + else: + local_stats = local_target_path.stat() + if blob.size != local_stats.st_size: + should_download = True + download_logs.append("Blob size (%s) and local file size (%s) differ.") + download_log_params.extend([blob.size, local_stats.st_size]) + + blob_last_modified = blob.last_modified + if blob_last_modified and local_stats.st_mtime < blob_last_modified.timestamp(): + should_download = True + download_logs.append("Blob last modified (%s) and local file last modified (%s) differ.") + download_log_params.extend([blob_last_modified.timestamp(), local_stats.st_mtime]) + + if should_download: + self.get_file( + file_path=str(local_target_path), + container_name=container_name, + blob_name=blob.name, + ) + download_logs.append("Downloaded %s to %s") + download_log_params.extend([blob.name, local_target_path.as_posix()]) + self.log.debug(" ".join(download_logs), *download_log_params) + else: + self.log.debug( + "Local file %s is up-to-date with blob %s. Skipping download.", + local_target_path.as_posix(), + blob.name, + ) + + def sync_to_local_dir( + self, + container_name: str, + local_dir: Path, + prefix: str | None = None, + delete_stale: bool = True, + ) -> None: + """Download files from an Azure Blob Storage container to a local directory.""" + self.log.debug("Downloading data from wasb://%s/%s to %s", container_name, prefix, local_dir) + + local_wasb_objects: list[Path] = [] + container = self._get_container_client(container_name) + self.check_for_variable_type("container", container, ContainerClient) + container = cast("ContainerClient", container) + + list_prefix = f"{prefix.rstrip('/')}/" if prefix else prefix + for blob in container.list_blobs(name_starts_with=list_prefix): + if blob.name.endswith("/"): + continue + blob_path = Path(blob.name) + if prefix: + local_target_path = local_dir.joinpath(blob_path.relative_to(prefix)) Review Comment: **[blocker]** `sync_to_local_dir` has no path-traversal guard, so a blob named `dags/../../outside.py` is written outside `local_dir`. `Path("dags/../../outside.py").relative_to("dags")` is `../../outside.py`, and `local_dir.joinpath(...)` happily resolves that above the bundle directory. Both sibling implementations already guard against this ([CWE-22](https://cwe.mitre.org/data/definitions/22.html)): `S3Hook.sync_to_local_dir` raises `S3HookPathTraversalError` since #67847 (merged 2026-06-17, after this PR was opened), and `GCSHook` checks `is_relative_to(local_dir_resolved)`. A new bundle should not ship without the same hardening; anyone with write access to the container would otherwise get arbitrary file writes on the Dag processor host. Mirror the S3 shape: ```python local_dir_resolved = local_dir.resolve() ... if not local_target_path.resolve().is_relative_to(local_dir_resolved): raise ValueError( f"Blob name {blob.name!r} resolves outside local directory {local_dir}" ) ``` (A dedicated exception in `providers/microsoft/azure/exceptions.py` would match what Amazon did; `ValueError` is fine too.) Please also add the equivalent of `test_sync_to_local_dir_rejects_key_path_traversal` from `test_s3.py`, asserting the exception is raised and that `tmp_path / "outside.py"` does not exist afterwards. ########## providers/microsoft/azure/src/airflow/providers/microsoft/azure/bundles/wasb.py: ########## @@ -0,0 +1,154 @@ +# 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 os +from functools import cached_property +from pathlib import Path +from urllib.parse import urlsplit, urlunsplit + +import structlog + +from airflow.dag_processing.bundles.base import BaseDagBundle +from airflow.providers.microsoft.azure.hooks.wasb import WasbHook + + +class WasbDagBundle(BaseDagBundle): + """ + WASB Dag bundle - exposes a directory in Azure Blob Storage as a Dag bundle. + + This allows Airflow to load Dags directly from an Azure Blob Storage container. + + :param wasb_conn_id: Airflow connection ID for Azure Blob Storage. Defaults to WasbHook.default_conn_name. + :param container_name: The name of the blob container containing the Dag files. + :param prefix: Optional subdirectory within the container where the Dags are stored. + If empty, Dags are assumed to be at the root of the container. + """ + + supports_versioning = False + + def __init__( + self, + *, + wasb_conn_id: str = WasbHook.default_conn_name, + container_name: str, + prefix: str = "", + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.wasb_conn_id = wasb_conn_id + self.container_name = container_name + self.prefix = prefix + self.wasb_dags_dir: Path = self.base_dir + + log = structlog.get_logger(__name__) + self._log = log.bind( + bundle_name=self.name, + version=self.version, + container_name=self.container_name, + prefix=self.prefix, + wasb_conn_id=self.wasb_conn_id, + ) + + def _initialize(self): + with self.lock(): + if not self.wasb_dags_dir.exists(): + self._log.info("Creating local Dags directory: %s", self.wasb_dags_dir) + os.makedirs(self.wasb_dags_dir) + + if not self.wasb_dags_dir.is_dir(): + raise NotADirectoryError(f"Local Dags path: {self.wasb_dags_dir} is not a directory.") + + if not self.wasb_hook.check_for_container(container_name=self.container_name): + raise ValueError(f"WASB container '{self.container_name}' does not exist.") + + if self.prefix: + if not self.wasb_hook.check_for_prefix( + container_name=self.container_name, prefix=self.prefix, delimiter="/" + ): + raise ValueError( + f"WASB prefix 'wasb://{self.container_name}/{self.prefix}' does not exist." + ) + self.refresh() + + def initialize(self) -> None: + self._initialize() + super().initialize() + + @cached_property + def wasb_hook(self) -> WasbHook: + return WasbHook(wasb_conn_id=self.wasb_conn_id) + + def __repr__(self): + return ( + f"<WasbDagBundle(" + f"name={self.name!r}, " + f"container_name={self.container_name!r}, " + f"prefix={self.prefix!r}, " + f"version={self.version!r}" + f")>" + ) + + def get_current_version(self) -> str | None: + """Return the current version of the Dag bundle. Currently not supported.""" + return None + + @property + def path(self) -> Path: + """Return the local path to the Dag files.""" + return self.wasb_dags_dir + + def refresh(self) -> None: + """Refresh the Dag bundle by re-downloading the Dags from Azure Blob Storage.""" + if self.version: + raise ValueError("Refreshing a specific version is not supported") + + with self.lock(): + self._log.debug( + "Downloading Dags from wasb://%s/%s to %s", + self.container_name, + self.prefix, + self.wasb_dags_dir, + ) + self.wasb_hook.sync_to_local_dir( + container_name=self.container_name, + prefix=self.prefix, + local_dir=self.wasb_dags_dir, + delete_stale=True, + ) + + def view_url(self, version: str | None = None) -> str | None: + """ + Return a URL for viewing the Dags in Azure Blob Storage. Currently, versioning is not supported. + + This method is deprecated and will be removed when the minimum supported Airflow version is 3.1. + Use `view_url_template` instead. + """ + return self.view_url_template() + + def view_url_template(self) -> str | None: + """Return a URL for viewing the Dags in Azure Blob Storage. Currently, versioning is not supported.""" + if self.version: + raise ValueError("WASB url with version is not supported") + if hasattr(self, "_view_url_template") and self._view_url_template: + return self._view_url_template + account_url = self.wasb_hook.blob_service_client.url + scheme, netloc, path, query, fragment = urlsplit(account_url) + path = f"{path.rstrip('/')}/{self.container_name}" + if self.prefix: + path = f"{path}/{self.prefix}" + return urlunsplit((scheme, netloc, path, query, fragment)) Review Comment: **[blocker]** The view URL still embeds the SAS token: the `urlsplit`/`urlunsplit` change fixed the malformed path @aaron-y-chen pointed out, but it deliberately re-attaches `query`, which is exactly where the SAS signature lives. I checked this against `azure-storage-blob`: when the connection's `sas_token` extra is an `https://…` URL (the hook passes it straight in as `account_url` on lines 205–206 of `hooks/wasb.py`), `blob_service_client.url` comes back as `https://<account>.blob.core.windows.net/?sv=…&sig=…`. This method then returns `https://<account>.blob.core.windows.net/<container>/<prefix>?sv=…&sig=…`. That string is stored by the bundle manager as the bundle's template URL and rendered as a link in the UI, so every user who can see the Dag gets a working SAS credential for the storage account ([CWE-522](https://cwe.mitre.org/data/definitions/522.html)). Neither `S3DagBundle` nor `GCSDagBundle` put credentials in their view URL. Drop the query and fragment when rebuilding the URL: ```python scheme, netloc, path, _, _ = urlsplit(account_url) path = f"{path.rstrip('/')}/{self.container_name}" if self.prefix: path = f"{path}/{self.prefix}" return urlunsplit((scheme, netloc, path, "", "")) ``` And add a test in `bundles/test_wasb.py` where `mock_hook.blob_service_client.url` carries a `?sv=…&sig=…` query and assert the returned URL contains neither `sig=` nor `?`. The two existing `test_view_url*` tests only cover a bare account URL, which is why this slipped through. ########## providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_wasb.py: ########## @@ -565,6 +597,101 @@ def test_download(self, mocked_blob_service_client): blob_client.assert_called_once_with(container="mycontainer", blob="myblob") blob_client.return_value.download_blob.assert_called_once_with(offset=2, length=4) + def test_sync_to_local_dir_behaviour(self, mocked_blob_service_client, tmp_path): Review Comment: **[warning]** The sync test only ever calls `sync_to_local_dir` with `prefix=""`, so the prefix branch, including the anchoring fix from the earlier thread, has no coverage. Every call in this test passes `prefix=""`, which means `list_prefix = f"{prefix.rstrip('/')}/"` and `blob_path.relative_to(prefix)` are never executed. Please add a case with `prefix="project1/dags"` (and one with a trailing slash) that asserts: - `mock_container.list_blobs` is called with `name_starts_with="project1/dags/"`, and - a blob named `project1/dags/dag_a.py` lands at `local_dir / "dag_a.py"`, not at `local_dir / "project1/dags/dag_a.py"`. The traversal test suggested in comment [4] naturally uses a prefix as well, so the two can share a setup. -- 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]
