This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 275aa6959f3 Support `abfs://` and `abfss://` URIs in DataFusion Azure
Blob Storage (#74228)
275aa6959f3 is described below
commit 275aa6959f360b04c8405f1f23f25809c97ea360
Author: Pankaj Singh <[email protected]>
AuthorDate: Wed Oct 14 13:19:58 2026 +0530
Support `abfs://` and `abfss://` URIs in DataFusion Azure Blob Storage
(#74228)
PR #73374 added az:// support to the DataFusion object storage layer
but deliberately deferred abfs:// and abfss://, since they embed the
storage account in the URI itself ([email protected]/path)
rather than relying solely on the wasb connection, raising an open
question of which account should win on disagreement.
Datasources can now use either scheme. When the URI names an account
that disagrees with the one the wasb connection resolves to,
registering the datasource raises instead of silently picking one.
* Quote az:// and abfs(s):// in docstrings to fix docs spellcheck
Sphinx's spellchecker flattens unquoted docstring prose to plain text,
so the bare az:// and abfs(s):// in two docstrings were flagged as
misspelled words. Wrapping them in double backticks, the convention
already used elsewhere in this file, excludes them as RST literals.
* Pass host to register_object_store so bucket-style URIs resolve
DataFusion's object-store registry keys on (schema, host); omitting
host only matches an empty-authority URI like file:///path. Every
bucket/container-style scheme (s3://, gs://, az://, and now
abfs(s)://) therefore failed at query time with "No suitable object
store found", a pre-existing gap in the original az:// support
invisible to the test suite because it mocks SessionContext rather
than exercising DataFusion's real registry. Found by manually testing
abfs:// against a real storage account.
* Fix object-store registry collisions, LOCAL regression, and config
mutation
Addresses review feedback: the abfs(s) object-store registry key collided
for two containers on the same storage account, an explicit
storage_type=StorageType.LOCAL with a bare path regressed when get_scheme
was generalized to take a URI, and a connection resolving no account
silently fell back to AZURE_STORAGE_ACCOUNT_NAME instead of the account
named in the URI.
Fixing the first of those by folding the container alone into the
registry's host key turned out to trade one collision for another -- two
different accounts sharing a container name would then collide instead.
Folding both account and container into the host key avoids both.
Also stopped register_datasource from mutating the caller's
DataSourceConfig.uri in place: reusing the same config across a second
engine fed the already-normalized URI back through get_bucket, which
misread the account+container host token as the container name. The
canonical URI is now threaded through locally instead.
Co-authored-by: Kaxil Naik <[email protected]>
---
providers/common/sql/docs/operators.rst | 18 ++-
.../sql/src/airflow/providers/common/sql/config.py | 19 +--
.../providers/common/sql/datafusion/base.py | 9 ++
.../providers/common/sql/datafusion/engine.py | 23 +++-
.../sql/datafusion/object_storage_provider.py | 70 +++++++++-
.../unit/common/sql/datafusion/test_engine.py | 100 ++++++++++++--
.../sql/datafusion/test_object_storage_provider.py | 153 +++++++++++++++++++++
.../sql/tests/unit/common/sql/test_config.py | 2 +
8 files changed, 368 insertions(+), 26 deletions(-)
diff --git a/providers/common/sql/docs/operators.rst
b/providers/common/sql/docs/operators.rst
index f4d2c251db9..5b59b805d71 100644
--- a/providers/common/sql/docs/operators.rst
+++ b/providers/common/sql/docs/operators.rst
@@ -365,10 +365,22 @@ resolved in this order:
Azure Storage
-------------
-Use an ``az://`` URI with a ``conn_id`` pointing to a ``wasb`` connection.
-``abfs://`` and ``abfss://`` URIs are not recognized yet. The account name
+Use an ``az://``, ``abfs://``, or ``abfss://`` URI with a ``conn_id`` pointing
+to a ``wasb`` connection. ``az://container/path`` takes the storage account
+from the connection; when the connection resolves none (no ``host``, and no
``login``
+or a service principal), the binding falls back to
``AZURE_STORAGE_ACCOUNT_NAME`` on
+the worker.
+``abfs://[email protected]/path`` and ``abfss://...``
+additionally name the account in the URI itself -- if it disagrees with the
+account the connection resolves to, registering the datasource raises rather
+than silently picking one, and when the connection names no account at all, the
+URI's account is used instead of ``AZURE_STORAGE_ACCOUNT_NAME``. The URI host
+follows the same cloud rule as the connection host below: a host outside
+``*.core.windows.net`` raises unless
``AZURE_STORAGE_ENDPOINT``/``AZURE_ENDPOINT``
+is set. The account name
comes from ``host`` (its first DNS label) when set, falling back to
-``login`` only when ``host`` is empty; only the public
+``login`` only when ``host`` is empty and the connection is not a service
principal
+(``tenant_id`` set, where ``login`` is the client ID); only the public
``*.blob.core.windows.net`` cloud is supported, since DataFusion's binding
otherwise has no endpoint override -- unless the worker sets
``AZURE_STORAGE_ENDPOINT``/``AZURE_ENDPOINT`` for a real sovereign-cloud
diff --git a/providers/common/sql/src/airflow/providers/common/sql/config.py
b/providers/common/sql/src/airflow/providers/common/sql/config.py
index ad1d2d82eb5..28e76b71055 100644
--- a/providers/common/sql/src/airflow/providers/common/sql/config.py
+++ b/providers/common/sql/src/airflow/providers/common/sql/config.py
@@ -52,6 +52,14 @@ class StorageType(str, Enum):
LOCAL = "local"
+STORAGE_TYPE_SCHEMES: dict[StorageType, tuple[str, ...]] = {
+ StorageType.S3: ("s3://",),
+ StorageType.GCS: ("gs://",),
+ StorageType.AZURE: ("az://", "abfs://", "abfss://"),
+ StorageType.LOCAL: ("file://",),
+}
+
+
@dataclass
class DataSourceConfig:
"""
@@ -114,12 +122,7 @@ class DataSourceConfig:
@property
def _extract_storage_type(self) -> StorageType | None:
"""Extract storage type."""
- if self.uri.startswith("s3://"):
- return StorageType.S3
- if self.uri.startswith("gs://"):
- return StorageType.GCS
- if self.uri.startswith("az://"):
- return StorageType.AZURE
- if self.uri.startswith("file://"):
- return StorageType.LOCAL
+ for storage_type, schemes in STORAGE_TYPE_SCHEMES.items():
+ if self.uri.startswith(schemes):
+ return storage_type
raise ValueError(f"Unsupported storage type for URI: {self.uri}")
diff --git
a/providers/common/sql/src/airflow/providers/common/sql/datafusion/base.py
b/providers/common/sql/src/airflow/providers/common/sql/datafusion/base.py
index f09881e552a..bb35ff7e22c 100644
--- a/providers/common/sql/src/airflow/providers/common/sql/datafusion/base.py
+++ b/providers/common/sql/src/airflow/providers/common/sql/datafusion/base.py
@@ -50,6 +50,15 @@ class ObjectStorageProvider(LoggingMixin, ABC):
"""Return URL scheme for this storage type (e.g., 's3://', 'gs://')."""
raise NotImplementedError
+ def normalize_uri(self, uri: str) -> str:
+ """
+ Return the canonical ``uri`` to register under.
+
+ Default is identity; a provider overrides this when its URI shape
differs from
+ DataFusion's (schema, host) registry key.
+ """
+ return uri
+
def get_bucket(self, path: str) -> str | None:
"""Extract the bucket name from the given path."""
if path and path.startswith(self.get_scheme()):
diff --git
a/providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py
b/providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py
index 3eec4d0f9c7..e5cc70b0527 100644
--- a/providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py
+++ b/providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py
@@ -17,6 +17,7 @@
from __future__ import annotations
import os
+from dataclasses import replace
from typing import TYPE_CHECKING, Any
from urllib.parse import parse_qsl, urlsplit
@@ -64,14 +65,15 @@ class DataFusionEngine(LoggingMixin):
else:
connection_config =
self._get_connection_config(datasource_config.conn_id)
- self._register_object_store(datasource_config, connection_config)
+ canonical_uri = self._register_object_store(datasource_config,
connection_config)
+ datasource_config = replace(datasource_config, uri=canonical_uri)
self._register_data_source_format(datasource_config)
def _register_object_store(
self, datasource_config: DataSourceConfig, connection_config:
ConnectionConfig | None
- ):
- """Register object stores."""
+ ) -> str:
+ """Register the object store and return the canonical URI to register
the table under."""
if TYPE_CHECKING:
assert datasource_config.storage_type is not None
@@ -80,9 +82,14 @@ class DataFusionEngine(LoggingMixin):
object_store = storage_provider.create_object_store(
datasource_config.uri, connection_config=connection_config
)
+ canonical_uri =
storage_provider.normalize_uri(datasource_config.uri)
schema = storage_provider.get_scheme()
- self.session_context.register_object_store(schema=schema,
store=object_store)
- self.log.info("Registered object store for schema: %s", schema)
+ # abfs(s) folds account+container into one host token, so the
container-name host
+ # DataFusion defaults to when host is omitted would miss the
registered store.
+ host = urlsplit(canonical_uri).netloc
+ self.session_context.register_object_store(schema=schema,
store=object_store, host=host)
+ self.log.info("Registered object store for schema: %s host: %s",
schema, host)
+ return canonical_uri
except Exception as e:
raise ObjectStoreCreationException(
f"Error while creating object store for
{datasource_config.storage_type}: {e}"
@@ -240,6 +247,9 @@ class DataFusionEngine(LoggingMixin):
credentials.update(
{"client_id": conn.login, "client_secret":
conn.password, "tenant_id": tenant_id}
)
+ if not conn.host:
+ # login holds the client_id here, so it cannot name
the storage account.
+ credentials["account"] = None
credential_tier = "client_secret"
elif shared_access_key := _get_wasb_extra_field(extra_dejson,
"shared_access_key"):
# Checked ahead of sas_token to match WasbHook.get_conn's
precedence.
@@ -353,7 +363,8 @@ class DataFusionEngine(LoggingMixin):
From ``host`` (its netloc's first label) when set, falling back to
``login`` only when
``host`` is empty -- login holds the service-principal client_id
otherwise, not the
- account name. Returns ``None`` when neither is set, so the binding
falls back to
+ account name (the caller drops the login-derived account for service
principals).
+ Returns ``None`` when neither is set, so the binding falls back to
``AZURE_STORAGE_ACCOUNT_NAME`` rather than the literal string
``"None"``. Reimplemented
locally instead of importing
``airflow.providers.microsoft.azure.utils.parse_blob_account_url``, to
avoid pulling in
diff --git
a/providers/common/sql/src/airflow/providers/common/sql/datafusion/object_storage_provider.py
b/providers/common/sql/src/airflow/providers/common/sql/datafusion/object_storage_provider.py
index 9799cbe00e2..78e1e50e4aa 100644
---
a/providers/common/sql/src/airflow/providers/common/sql/datafusion/object_storage_provider.py
+++
b/providers/common/sql/src/airflow/providers/common/sql/datafusion/object_storage_provider.py
@@ -17,6 +17,8 @@
from __future__ import annotations
import json
+import os
+import re
import tempfile
from pathlib import Path
@@ -107,6 +109,15 @@ class GCSObjectStorageProvider(ObjectStorageProvider):
return "gs://"
+_ABFS_URI_RE = re.compile(
+
r"^abfss?://(?P<container>[^@/]+)@(?P<account>[^./]+)\.(?P<host>[^/]+)(?P<path>/.*)?$"
+)
+# MicrosoftAzure has no endpoint kwarg, so a non-public host only reaches the
right cloud when
+# AZURE_STORAGE_ENDPOINT/AZURE_ENDPOINT redirects the binding; same rule as
_resolve_wasb_account.
+_AZURE_PUBLIC_HOST_SUFFIX = ".core.windows.net"
+_ABFS_SCHEMES = ("abfs://", "abfss://")
+
+
class AzureObjectStorageProvider(ObjectStorageProvider):
"""Azure Object Storage Provider using DataFusion's MicrosoftAzure."""
@@ -115,6 +126,48 @@ class AzureObjectStorageProvider(ObjectStorageProvider):
"""Return the storage type."""
return StorageType.AZURE
+ def get_bucket(self, path: str) -> str | None:
+ """Extract the container name, from either the ``az://`` or
``abfs(s)://`` URI shape."""
+ if match := _ABFS_URI_RE.match(path):
+ return match.group("container")
+ if path.startswith(_ABFS_SCHEMES):
+ raise ValueError(
+ f"{path!r} does not match the required
abfs(s)://<container>@<account>.<host>/<path> shape"
+ )
+ return super().get_bucket(path)
+
+ def _validate_uri_host(self, path: str) -> None:
+ match = _ABFS_URI_RE.match(path)
+ if not match:
+ return
+ host = match.group("host").lower()
+ if host == _AZURE_PUBLIC_HOST_SUFFIX[1:] or
host.endswith(_AZURE_PUBLIC_HOST_SUFFIX):
+ return
+ if os.environ.get("AZURE_STORAGE_ENDPOINT") or
os.environ.get("AZURE_ENDPOINT"):
+ return
+ raise ValueError(
+ f"URI {path!r} host {match.group('host')!r} is not a public
*{_AZURE_PUBLIC_HOST_SUFFIX} "
+ "endpoint, and DataFusion's Azure Blob Storage binding has no
endpoint override, so the "
+ "request would go to the public cloud. A sovereign cloud is
supported once the "
+ "AZURE_STORAGE_ENDPOINT environment variable is set."
+ )
+
+ def _get_uri_account(self, path: str) -> str | None:
+ match = _ABFS_URI_RE.match(path)
+ return match.group("account") if match else None
+
+ def normalize_uri(self, path: str) -> str:
+ """
+ Rewrite ``abfs(s)://<container>@<account>.<host>/<path>`` to
``az://<account>.<container>/<path>``.
+
+ DataFusion's registry keys on (schema, host) alone, so container-only
would collide for
+ two different accounts sharing a container name. Account and container
names never
+ contain a dot, so joining on one is unambiguous. ``az://`` URIs pass
through unchanged.
+ """
+ if match := _ABFS_URI_RE.match(path):
+ return
f"az://{match.group('account')}.{match.group('container')}{match.group('path')
or ''}"
+ return path
+
def create_object_store(self, path: str, connection_config:
ConnectionConfig | None = None):
"""Create an Azure object store using DataFusion's MicrosoftAzure."""
if connection_config is None:
@@ -124,6 +177,21 @@ class AzureObjectStorageProvider(ObjectStorageProvider):
credentials = connection_config.credentials
container = self.get_bucket(path)
+ self._validate_uri_host(path)
+ uri_account = self._get_uri_account(path)
+ resolved_account = credentials.get("account")
+ if uri_account and resolved_account and uri_account.lower() !=
resolved_account.lower():
+ raise ValueError(
+ f"URI {path!r} names storage account {uri_account!r}, but
connection "
+ f"{connection_config.conn_id!r} resolves to account
{resolved_account!r}. Point "
+ "the URI and the connection at the same account, or omit
the account from one of "
+ "them."
+ )
+ if uri_account and not resolved_account:
+ # Pins the URI's account; MicrosoftAzure otherwise falls back
to AZURE_STORAGE_ACCOUNT_NAME,
+ # which can name a different account.
+ credentials = {**credentials, "account": uri_account}
+
azure_store = MicrosoftAzure(container_name=container,
**credentials)
self.log.info("Created Azure object store for container %s",
container)
@@ -134,7 +202,7 @@ class AzureObjectStorageProvider(ObjectStorageProvider):
raise ObjectStoreCreationException(f"Failed to create Azure object
store: {e}")
def get_scheme(self) -> str:
- """Return the scheme for Azure."""
+ """Return the scheme for Azure; ``normalize_uri`` rewrites
``abfs(s)://`` URIs onto it."""
return "az://"
diff --git
a/providers/common/sql/tests/unit/common/sql/datafusion/test_engine.py
b/providers/common/sql/tests/unit/common/sql/datafusion/test_engine.py
index c80434ef578..05e62da11f0 100644
--- a/providers/common/sql/tests/unit/common/sql/datafusion/test_engine.py
+++ b/providers/common/sql/tests/unit/common/sql/datafusion/test_engine.py
@@ -18,10 +18,12 @@ from __future__ import annotations
import os
import tempfile
+from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
from datafusion import SessionContext
+from datafusion.object_store import LocalFileSystem
from airflow.models import Connection
from airflow.providers.common.sql.config import ConnectionConfig,
DataSourceConfig
@@ -96,6 +98,7 @@ class TestDataFusionEngine:
mock_store = MagicMock()
mock_provider.create_object_store.return_value = mock_store
mock_provider.get_scheme.return_value = scheme
+ mock_provider.normalize_uri.side_effect = lambda uri: uri
mock_factory.return_value = mock_provider
engine = DataFusionEngine()
@@ -112,7 +115,9 @@ class TestDataFusionEngine:
mock_provider.create_object_store.assert_called_once_with(
f"{scheme}://bucket/path",
connection_config=mock_get_conn.return_value
)
-
engine.df_ctx.register_object_store.assert_called_once_with(schema=scheme,
store=mock_store)
+ engine.df_ctx.register_object_store.assert_called_once_with(
+ schema=scheme, store=mock_store, host="bucket"
+ )
if format == "parquet":
engine.df_ctx.register_parquet.assert_called_once_with("test_table",
f"{scheme}://bucket/path")
@@ -208,8 +213,10 @@ class TestDataFusionEngine:
with pytest.raises(QueryExecutionException, match="Error while
executing query"):
engine.execute_query("SELECT * FROM test_table")
+ @pytest.mark.parametrize("uri_prefix", ["file://", ""], ids=["file-uri",
"bare-path"])
@patch.object(DataFusionEngine, "_get_connection_config")
- def test_execute_query_with_local_csv(self, mock_get_conn):
+ def test_execute_query_with_local_csv(self, mock_get_conn, uri_prefix):
+ """A bare path (no file:// prefix) is accepted with an explicit
storage_type=StorageType.LOCAL."""
mock_get_conn.return_value = None
with tempfile.NamedTemporaryFile(mode="w", suffix=".csv",
delete=False) as f:
@@ -220,7 +227,7 @@ class TestDataFusionEngine:
engine = DataFusionEngine()
datasource_config = DataSourceConfig(
table_name="test_csv",
- uri=f"file://{csv_path}",
+ uri=f"{uri_prefix}{csv_path}",
format="csv",
storage_type="local",
conn_id="",
@@ -230,11 +237,87 @@ class TestDataFusionEngine:
result = engine.execute_query("SELECT * FROM test_csv ORDER BY
name")
- expected = {"name": ["Alice", "Bob"], "age": [30, 25]}
- assert result == expected
+ assert result == {"name": ["Alice", "Bob"], "age": [30, 25]}
+ finally:
+ os.unlink(csv_path)
+
+ @patch.object(DataFusionEngine, "_get_connection_config")
+ def test_execute_query_with_abfs_uri_matches_real_registry(self,
mock_get_conn):
+ """Only object-store construction is mocked; normalize_uri and
DataFusion's real
+ registry run live, so a missing store surfaces as "No suitable object
store found"."""
+ mock_get_conn.return_value = None
+
+ with tempfile.NamedTemporaryFile(mode="w", suffix=".csv",
delete=False) as f:
+ f.write("name,age\nAlice,30\nBob,25\n")
+ csv_path = f.name
+
+ try:
+ engine = DataFusionEngine()
+ datasource_config = DataSourceConfig(
+ table_name="abfs_csv",
+
uri=f"abfss://[email protected]{csv_path}",
+ format="csv",
+ conn_id="wasb_default",
+ )
+
+ with patch(
+
"airflow.providers.common.sql.datafusion.object_storage_provider"
+ ".AzureObjectStorageProvider.create_object_store",
+ autospec=True,
+ return_value=LocalFileSystem(),
+ ):
+ engine.register_datasource(datasource_config)
+
+ result = engine.execute_query("SELECT * FROM abfs_csv ORDER BY
name")
+
+ assert result == {"name": ["Alice", "Bob"], "age": [30, 25]}
+ assert datasource_config.uri ==
f"abfss://[email protected]{csv_path}"
+ assert engine.registered_tables["abfs_csv"] ==
f"az://myaccount.demo-container{csv_path}"
finally:
os.unlink(csv_path)
+ @pytest.mark.parametrize(
+ ("uri_a", "uri_b"),
+ [
+ pytest.param(
+ "abfss://[email protected]/data.csv",
+ "abfss://[email protected]/data.csv",
+ id="same-account-different-containers",
+ ),
+ pytest.param(
+ "abfss://[email protected]/data.csv",
+ "abfss://[email protected]/data.csv",
+ id="same-container-different-accounts",
+ ),
+ ],
+ )
+ @patch.object(DataFusionEngine, "_get_connection_config")
+ def test_execute_query_with_shared_abfs_identifiers_reads_each_store(self,
mock_get_conn, uri_a, uri_b):
+ mock_get_conn.return_value = None
+
+ with (
+ tempfile.TemporaryDirectory() as dir_a,
+ tempfile.TemporaryDirectory() as dir_b,
+ ):
+ (Path(dir_a) / "data.csv").write_text("name,age\nAlice,30\n")
+ (Path(dir_b) / "data.csv").write_text("name,age\nBob,25\n")
+
+ engine = DataFusionEngine()
+ config_a = DataSourceConfig(table_name="table_a", uri=uri_a,
format="csv", conn_id="wasb_a")
+ config_b = DataSourceConfig(table_name="table_b", uri=uri_b,
format="csv", conn_id="wasb_b")
+
+ with patch(
+
"airflow.providers.common.sql.datafusion.object_storage_provider"
+ ".AzureObjectStorageProvider.create_object_store",
+ autospec=True,
+ side_effect=[LocalFileSystem(prefix=dir_a),
LocalFileSystem(prefix=dir_b)],
+ ):
+ engine.register_datasource(config_a)
+ engine.register_datasource(config_b)
+
+ assert engine.execute_query("SELECT * FROM table_a") == {"name":
["Alice"], "age": [30]}
+ assert engine.execute_query("SELECT * FROM table_b") == {"name":
["Bob"], "age": [25]}
+
@patch("airflow.providers.common.sql.datafusion.engine.get_object_storage_provider",
autospec=True)
@patch.object(DataFusionEngine, "_get_connection_config")
def test_register_datasource_with_options(self, mock_get_conn,
mock_factory):
@@ -243,6 +326,7 @@ class TestDataFusionEngine:
mock_store = MagicMock()
mock_provider.create_object_store.return_value = mock_store
mock_provider.get_scheme.return_value = "s3"
+ mock_provider.normalize_uri.side_effect = lambda uri: uri
mock_factory.return_value = mock_provider
engine = DataFusionEngine()
@@ -263,7 +347,9 @@ class TestDataFusionEngine:
mock_provider.create_object_store.assert_called_once_with(
"s3://bucket/path/", connection_config=mock_get_conn.return_value
)
-
engine.df_ctx.register_object_store.assert_called_once_with(schema="s3",
store=mock_store)
+ engine.df_ctx.register_object_store.assert_called_once_with(
+ schema="s3", store=mock_store, host="bucket"
+ )
engine.df_ctx.register_parquet.assert_called_once_with(
"test_table",
@@ -430,7 +516,6 @@ class TestDataFusionEngine:
credentials, extra_config = engine._get_credentials(mock_conn)
assert credentials == {
- "account": "client-id",
"client_id": "client-id",
"client_secret": "client-secret",
"tenant_id": "tenant-id",
@@ -748,7 +833,6 @@ class TestDataFusionEngine:
else:
credentials, extra_config = engine._get_credentials(mock_conn)
assert credentials == {
- "account": "client-id",
"client_id": "client-id",
"client_secret": "client-secret",
"tenant_id": "tenant-id",
diff --git
a/providers/common/sql/tests/unit/common/sql/datafusion/test_object_storage_provider.py
b/providers/common/sql/tests/unit/common/sql/datafusion/test_object_storage_provider.py
index 308879489f7..a6c8ad7298b 100644
---
a/providers/common/sql/tests/unit/common/sql/datafusion/test_object_storage_provider.py
+++
b/providers/common/sql/tests/unit/common/sql/datafusion/test_object_storage_provider.py
@@ -159,6 +159,136 @@ class TestObjectStorageProvider:
assert provider.get_storage_type == StorageType.AZURE
assert provider.get_scheme() == "az://"
+ @pytest.mark.parametrize(
+ "uri",
+ [
+ "abfs://[email protected]/path",
+ "abfss://[email protected]/path",
+ "abfss://[email protected]/path",
+ "abfss://[email protected]/path",
+ ],
+ )
+
@patch("airflow.providers.common.sql.datafusion.object_storage_provider.MicrosoftAzure")
+ def test_azure_provider_success_with_abfs_uri(self, mock_azure, uri,
monkeypatch):
+ # Without these, a public-cloud host check that is too strict would
still pass.
+ monkeypatch.delenv("AZURE_STORAGE_ENDPOINT", raising=False)
+ monkeypatch.delenv("AZURE_ENDPOINT", raising=False)
+ provider = AzureObjectStorageProvider()
+ connection_config = ConnectionConfig(
+ conn_id="wasb_default",
+ credentials={"account": "myaccount", "access_key": "fake_key"},
+ )
+
+ store = provider.create_object_store(uri, connection_config)
+
+ mock_azure.assert_called_once_with(
+ container_name="demo-container", account="myaccount",
access_key="fake_key"
+ )
+ assert store == mock_azure.return_value
+ assert provider.get_bucket(uri) == "demo-container"
+
+
@patch("airflow.providers.common.sql.datafusion.object_storage_provider.MicrosoftAzure")
+ def test_azure_provider_abfs_uri_account_mismatch_raises(self, mock_azure):
+ provider = AzureObjectStorageProvider()
+ connection_config = ConnectionConfig(
+ conn_id="wasb_default",
+ credentials={"account": "connection-account", "access_key":
"fake_key"},
+ )
+
+ with pytest.raises(ObjectStoreCreationException, match="names storage
account 'uri-account'"):
+ provider.create_object_store(
+
"abfss://[email protected]/path",
connection_config
+ )
+
+ mock_azure.assert_not_called()
+
+ @pytest.mark.parametrize(
+ "uri",
+ [
+ "abfs://demo-container/path", # missing @account entirely
+ "abfss://demo-container@/path", # empty account
+ ],
+ )
+
@patch("airflow.providers.common.sql.datafusion.object_storage_provider.MicrosoftAzure")
+ def test_azure_provider_malformed_abfs_uri_raises_clear_error(self,
mock_azure, uri):
+ provider = AzureObjectStorageProvider()
+ connection_config = ConnectionConfig(
+ conn_id="wasb_default",
+ credentials={"account": "myaccount", "access_key": "fake_key"},
+ )
+
+ with pytest.raises(ObjectStoreCreationException, match="does not match
the required"):
+ provider.create_object_store(uri, connection_config)
+
+ mock_azure.assert_not_called()
+
+
@patch("airflow.providers.common.sql.datafusion.object_storage_provider.MicrosoftAzure")
+ def
test_azure_provider_abfs_uri_account_used_when_connection_resolves_none(self,
mock_azure):
+ provider = AzureObjectStorageProvider()
+ connection_config = ConnectionConfig(
+ conn_id="wasb_default",
+ credentials={"access_key": "fake_key"},
+ )
+
+ provider.create_object_store(
+ "abfss://[email protected]/path",
connection_config
+ )
+
+ mock_azure.assert_called_once_with(
+ container_name="demo-container", account="uri-account",
access_key="fake_key"
+ )
+
+ @pytest.mark.parametrize("account", [None, "usgovacct"])
+
@patch("airflow.providers.common.sql.datafusion.object_storage_provider.MicrosoftAzure")
+ def test_azure_provider_abfs_uri_sovereign_host_without_endpoint_raises(
+ self, mock_azure, account, monkeypatch
+ ):
+ monkeypatch.delenv("AZURE_STORAGE_ENDPOINT", raising=False)
+ monkeypatch.delenv("AZURE_ENDPOINT", raising=False)
+ credentials = {"access_key": "fake_key"}
+ if account:
+ credentials["account"] = account
+ connection_config = ConnectionConfig(conn_id="wasb_default",
credentials=credentials)
+
+ with pytest.raises(ObjectStoreCreationException, match="no endpoint
override"):
+ AzureObjectStorageProvider().create_object_store(
+
"abfss://[email protected]/path",
connection_config
+ )
+
+ mock_azure.assert_not_called()
+
+ @pytest.mark.parametrize("env_var", ["AZURE_STORAGE_ENDPOINT",
"AZURE_ENDPOINT"])
+
@patch("airflow.providers.common.sql.datafusion.object_storage_provider.MicrosoftAzure")
+ def test_azure_provider_abfs_uri_sovereign_host_with_endpoint_env_passes(
+ self, mock_azure, env_var, monkeypatch
+ ):
+ monkeypatch.setenv(env_var,
"https://usgovacct.blob.core.usgovcloudapi.net";)
+ connection_config = ConnectionConfig(conn_id="wasb_default",
credentials={"access_key": "fake_key"})
+
+ AzureObjectStorageProvider().create_object_store(
+
"abfss://[email protected]/path",
connection_config
+ )
+
+ mock_azure.assert_called_once_with(
+ container_name="demo-container", account="usgovacct",
access_key="fake_key"
+ )
+
+
@patch("airflow.providers.common.sql.datafusion.object_storage_provider.MicrosoftAzure")
+ def test_azure_provider_abfs_uri_account_matches_case_insensitively(self,
mock_azure):
+ provider = AzureObjectStorageProvider()
+ connection_config = ConnectionConfig(
+ conn_id="wasb_default",
+ credentials={"account": "MyAccount", "access_key": "fake_key"},
+ )
+
+ provider.create_object_store(
+ "abfss://[email protected]/path",
connection_config
+ )
+
+ mock_azure.assert_called_once_with(
+ container_name="demo-container", account="MyAccount",
access_key="fake_key"
+ )
+
def test_azure_provider_failure(self):
provider = AzureObjectStorageProvider()
connection_config = ConnectionConfig(conn_id="wasb_default")
@@ -176,6 +306,29 @@ class TestObjectStorageProvider:
with pytest.raises(ValueError, match="connection_config must be
provided for azure"):
provider.create_object_store("az://demo-container/path")
+ @pytest.mark.parametrize(
+ ("uri", "expected"),
+ [
+ (
+
"abfs://[email protected]/path/to/data.parquet",
+ "az://myaccount.demo-container/path/to/data.parquet",
+ ),
+ (
+
"abfss://[email protected]/path/to/data.parquet",
+ "az://myaccount.demo-container/path/to/data.parquet",
+ ),
+ ("abfss://[email protected]",
"az://myaccount.demo-container"),
+ ("az://demo-container/path", "az://demo-container/path"),
+ # Same container name on different accounts maps to distinct hosts.
+ ("abfss://[email protected]/data.csv",
"az://devacct.data/data.csv"),
+ ("abfss://[email protected]/data.csv",
"az://prodacct.data/data.csv"),
+ ],
+ )
+ def test_azure_provider_normalize_uri(self, uri, expected):
+ provider = AzureObjectStorageProvider()
+
+ assert provider.normalize_uri(uri) == expected
+
def test_azure_provider_partial_service_principal_raises_clear_error(self):
"""Uses the real MicrosoftAzure binding, not a mock, since it's the
one that panics
on a partial client_id/client_secret/tenant_id combination."""
diff --git a/providers/common/sql/tests/unit/common/sql/test_config.py
b/providers/common/sql/tests/unit/common/sql/test_config.py
index 359721681fd..a2628c4d730 100644
--- a/providers/common/sql/tests/unit/common/sql/test_config.py
+++ b/providers/common/sql/tests/unit/common/sql/test_config.py
@@ -35,6 +35,8 @@ class TestDataSourceConfig:
("s3://bucket/path", StorageType.S3),
("gs://bucket/path", StorageType.GCS),
("az://container/path", StorageType.AZURE),
+ ("abfs://[email protected]/path",
StorageType.AZURE),
+ ("abfss://[email protected]/path",
StorageType.AZURE),
("file:///path/to/file", StorageType.LOCAL),
],
)
