-
Notifications
You must be signed in to change notification settings - Fork 17.4k
add WasbDagBundle to load Dags from Azure Blob Storage #67016
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
Nishieee
wants to merge
8
commits into
apache:main
Choose a base branch
from
Nishieee:feature/azure-blob-dag-bundle
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
8 commits
Select commit
Hold shift + click to select a range
d8f7cd1
add WasbDagBundle to load Dags from Azure Blob Storage
Nishieee edf61fd
Use ValueError and NotADirectoryError in WasbDagBundle instead of Air…
Nishieee 230b1a2
Docs: Add WasbDagBundle guide to Microsoft Azure provider
Nishieee 2b5dacd
Update providers/microsoft/azure/src/airflow/providers/microsoft/azur…
Nishieee 50c1b92
Add WasbHook.check_for_container tests and fix WasbDagBundle view URL…
Nishieee 6baf8dc
Apply ruff-format to WasbHook tests
Nishieee 4070f0e
Use ContainerClient.exists() in WasbHook.check_for_container
Nishieee 51d73e9
Add sync_to_local_dir tests for WasbHook
Nishieee File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
| ####### | ||
|
|
||
| Dag bundles allow Airflow to load Dags from external sources. For a general overview see | ||
| :doc:`apache-airflow:administration-and-deployment/dag-bundles`. | ||
|
|
||
| WasbDagBundle | ||
| ============= | ||
|
|
||
| Use the :class:`~airflow.providers.microsoft.azure.bundles.wasb.WasbDagBundle` to configure an Azure Blob | ||
| Storage bundle in your Airflow's ``[dag_processor] dag_bundle_config_list``. The bundle does not support | ||
| versioning; tasks always use the latest blobs synced to the local bundle directory. | ||
|
|
||
| Example of using the WasbDagBundle: | ||
|
|
||
| **JSON format example**: | ||
|
|
||
| .. code-block:: bash | ||
|
|
||
| export AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST='[ | ||
| { | ||
| "name": "my-wasb-dags", | ||
| "classpath": "airflow.providers.microsoft.azure.bundles.wasb.WasbDagBundle", | ||
| "kwargs": { | ||
| "wasb_conn_id": "wasb_default", | ||
| "container_name": "airflow-dags", | ||
| "prefix": "dags/", | ||
| "refresh_interval": 60 | ||
| } | ||
| } | ||
| ]' | ||
|
|
||
| Authentication | ||
| -------------- | ||
|
|
||
| The bundle uses a ``wasb`` Connection (``wasb_conn_id``). Authentication is the same as for | ||
| :class:`~airflow.providers.microsoft.azure.hooks.wasb.WasbHook` — see :doc:`../connections/wasb`. On Azure-hosted | ||
| Airflow, managed identity via ``DefaultAzureCredential`` is typical. | ||
|
|
||
| Permissions | ||
| ----------- | ||
|
|
||
| The identity needs read access to list and download blobs in the target container. Assign | ||
| `Storage Blob Data Reader <https://learn.microsoft.com/en-us/azure/role-based-access-control/built-in-roles/storage#storage-blob-data-reader>`_ | ||
| at the storage account or container scope. | ||
|
|
||
| Container and prefix | ||
| -------------------- | ||
|
|
||
| Set ``container_name`` to the blob container that holds your Dag files. Use ``prefix`` for an optional | ||
| virtual folder inside the container. | ||
|
|
||
| Networking | ||
| ---------- | ||
|
|
||
| The Dag processor needs outbound HTTPS to the blob endpoint. Storage firewalls and private endpoints | ||
| must allow access from Airflow, as for any WASB client. | ||
|
|
||
| Reusing the Connection in Dags | ||
| ------------------------------ | ||
|
|
||
| You can use the same ``wasb`` Connection ID in ``wasb_conn_id`` for the bundle and for operators or sensors | ||
| that use ``WasbHook``. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
16 changes: 16 additions & 0 deletions
16
providers/microsoft/azure/src/airflow/providers/microsoft/azure/bundles/__init__.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,16 @@ | ||
| # 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. |
154 changes: 154 additions & 0 deletions
154
providers/microsoft/azure/src/airflow/providers/microsoft/azure/bundles/wasb.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 pathlib import Path | ||
|
|
||
| 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__) | ||
|
dabla marked this conversation as resolved.
|
||
| 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, | ||
| ) | ||
| self._wasb_hook: WasbHook | None = None | ||
|
|
||
| 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() | ||
|
|
||
| @property | ||
| def wasb_hook(self): | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This could simply become a cached_property I think, no need to keep self._wasb_hook. |
||
| if self._wasb_hook is None: | ||
| self._wasb_hook = WasbHook(wasb_conn_id=self.wasb_conn_id) | ||
| return self._wasb_hook | ||
|
|
||
| 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: | ||
|
Nishieee marked this conversation as resolved.
|
||
| """ | ||
| 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 | ||
| url = f"{account_url.rstrip('/')}/{self.container_name}" | ||
| if self.prefix: | ||
| url += f"/{self.prefix}" | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
| return url | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.