apache-airflow-providers-common-sql 2.3.0__tar.gz → 2.4.0__tar.gz

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: apache-airflow-providers-common-sql
-Version: 2.3.0
+Version: 2.4.0
 Summary: Provider package apache-airflow-providers-common-sql for Apache Airflow
 Keywords: airflow-provider,common.sql,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -23,7 +23,7 @@ Classifier: Topic :: System :: Monitoring
 License-File: LICENSE
 License-File: NOTICE
 Requires-Dist: apache-airflow>=2.11.0
-Requires-Dist: apache-airflow-providers-common-compat>=1.14.1
+Requires-Dist: apache-airflow-providers-common-compat>=1.23.0
 Requires-Dist: sqlparse>=0.5.1
 Requires-Dist: more-itertools>=9.0.0
 Requires-Dist: methodtools>=0.4.7
@@ -38,8 +38,8 @@ Requires-Dist: polars>=1.26.0 ; extra == "polars"
 Requires-Dist: pyiceberg-core>=0.8.0 ; extra == "pyiceberg-core"
 Requires-Dist: sqlalchemy>=1.4.54,!=2.1.0 ; extra == "sqlalchemy"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.3.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.3.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.4.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.4.0
 Project-URL: Mastodon, https://fosstodon.org/@airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
@@ -82,7 +82,7 @@ Import-Namespace: airflow.providers.common
 
 Package ``apache-airflow-providers-common-sql``
 
-Release: ``2.3.0``
+Release: ``2.4.0``
 
 
 `Common SQL Provider <https://en.wikipedia.org/wiki/SQL>`__
@@ -95,7 +95,7 @@ This is a provider package for ``common.sql`` provider. All classes for this pro
 are in ``airflow.providers.common.sql`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.3.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.4.0/>`_.
 
 Installation
 ------------
@@ -113,7 +113,7 @@ Requirements
 PIP package                                 Version required
 ==========================================  ==================
 ``apache-airflow``                          ``>=2.11.0``
-``apache-airflow-providers-common-compat``  ``>=1.14.1``
+``apache-airflow-providers-common-compat``  ``>=1.23.0``
 ``sqlparse``                                ``>=0.5.1``
 ``more-itertools``                          ``>=9.0.0``
 ``methodtools``                             ``>=0.4.7``
@@ -157,5 +157,5 @@ Extra               Dependencies
 ==================  =======================================================================================================================================================================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.3.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.4.0/changelog.html>`_.
 

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/README.rst

@@ -23,7 +23,7 @@
 
 Package ``apache-airflow-providers-common-sql``
 
-Release: ``2.3.0``
+Release: ``2.4.0``
 
 
 `Common SQL Provider <https://en.wikipedia.org/wiki/SQL>`__
@@ -36,7 +36,7 @@ This is a provider package for ``common.sql`` provider. All classes for this pro
 are in ``airflow.providers.common.sql`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.3.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.4.0/>`_.
 
 Installation
 ------------
@@ -54,7 +54,7 @@ Requirements
 PIP package                                 Version required
 ==========================================  ==================
 ``apache-airflow``                          ``>=2.11.0``
-``apache-airflow-providers-common-compat``  ``>=1.14.1``
+``apache-airflow-providers-common-compat``  ``>=1.23.0``
 ``sqlparse``                                ``>=0.5.1``
 ``more-itertools``                          ``>=9.0.0``
 ``methodtools``                             ``>=0.4.7``
@@ -98,4 +98,4 @@ Extra               Dependencies
 ==================  =======================================================================================================================================================================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.3.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.4.0/changelog.html>`_.

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/docs/changelog.rst

@@ -25,6 +25,18 @@
 Changelog
 ---------
 
+2.4.0
+.....
+
+Features
+~~~~~~~~
+
+* ``Support 'abfs://' and 'abfss://' URIs in DataFusion Azure Blob Storage (#74228)``
+
+.. Below changes are excluded from the changelog. Move them to
+   appropriate section above if needed. Do not delete the lines(!):
+   * ``Remove unused mypy suppressions in common providers (#75257)``
+
 2.3.0
 .....
 

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/docs/index.rst

@@ -79,7 +79,7 @@ apache-airflow-providers-common-sql package
 `Common SQL Provider <https://en.wikipedia.org/wiki/SQL>`__
 
 
-Release: 2.3.0
+Release: 2.4.0
 
 Provider package
 ----------------
@@ -103,7 +103,7 @@ The minimum Apache Airflow version supported by this provider distribution is ``
 PIP package                                 Version required
 ==========================================  ==================
 ``apache-airflow``                          ``>=2.11.0``
-``apache-airflow-providers-common-compat``  ``>=1.14.1``
+``apache-airflow-providers-common-compat``  ``>=1.23.0``
 ``sqlparse``                                ``>=0.5.1``
 ``more-itertools``                          ``>=9.0.0``
 ``methodtools``                             ``>=0.4.7``
@@ -160,5 +160,5 @@ Downloading official packages
 You can download officially released packages and verify their checksums and signatures from the
 `Official Apache Download site <https://downloads.apache.org/airflow/providers/>`_
 
-* `The apache-airflow-providers-common-sql 2.3.0 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-2.3.0.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-2.3.0.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-2.3.0.tar.gz.sha512>`__)
-* `The apache-airflow-providers-common-sql 2.3.0 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-2.3.0-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-2.3.0-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-2.3.0-py3-none-any.whl.sha512>`__)
+* `The apache-airflow-providers-common-sql 2.4.0 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-2.4.0.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-2.4.0.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-2.4.0.tar.gz.sha512>`__)
+* `The apache-airflow-providers-common-sql 2.4.0 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-2.4.0-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-2.4.0-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_common_sql-2.4.0-py3-none-any.whl.sha512>`__)

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/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://container@account.dfs.core.windows.net/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

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/provider.yaml

@@ -23,12 +23,13 @@ description: |
 
 state: ready
 lifecycle: production
-source-date-epoch: 1791328846
+source-date-epoch: 1792489528
 # Note that those versions are maintained by release manager - do not update them manually
 # with the exception of case where other provider in sources has >= new provider version.
 # In such case adding >= NEW_VERSION and bumping to NEW_VERSION in a provider have
 # to be done in the same PR
 versions:
+  - 2.4.0
   - 2.3.0
   - 2.2.0
   - 2.1.1

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/pyproject.toml

@@ -25,7 +25,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-common-sql"
-version = "2.3.0"
+version = "2.4.0"
 description = "Provider package apache-airflow-providers-common-sql for Apache Airflow"
 readme = "README.rst"
 license = "Apache-2.0"
@@ -59,7 +59,7 @@ requires-python = ">=3.11"
 # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
 dependencies = [
     "apache-airflow>=2.11.0",
-    "apache-airflow-providers-common-compat>=1.14.1",
+    "apache-airflow-providers-common-compat>=1.23.0",
     "sqlparse>=0.5.1",
     "more-itertools>=9.0.0",
     # The methodtools dependency is necessary since the introduction of dialects:
@@ -147,8 +147,8 @@ apache-airflow-providers-common-sql = {workspace = true}
 apache-airflow-providers-standard = {workspace = true}
 
 [project.urls]
-"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.3.0"
-"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.3.0/changelog.html"
+"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.4.0"
+"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-common-sql/2.4.0/changelog.html"
 "Bug Tracker" = "https://github.com/apache/airflow/issues"
 "Source Code" = "https://github.com/apache/airflow"
 "Slack Chat" = "https://s.apache.org/airflow-slack"

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/src/airflow/providers/common/sql/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "2.3.0"
+__version__ = "2.4.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.11.0"

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/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}")

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/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()):

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/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

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/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://"
 
 

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/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://demo-container@myaccount.dfs.core.windows.net{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://demo-container@myaccount.dfs.core.windows.net{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://bronze@myaccount.dfs.core.windows.net/data.csv",
+                "abfss://silver@myaccount.dfs.core.windows.net/data.csv",
+                id="same-account-different-containers",
+            ),
+            pytest.param(
+                "abfss://data@devacct.dfs.core.windows.net/data.csv",
+                "abfss://data@prodacct.dfs.core.windows.net/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",

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/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://demo-container@myaccount.dfs.core.windows.net/path",
+            "abfss://demo-container@myaccount.dfs.core.windows.net/path",
+            "abfss://demo-container@myaccount.blob.core.windows.net/path",
+            "abfss://demo-container@myaccount.DFS.Core.Windows.Net/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://demo-container@uri-account.dfs.core.windows.net/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://demo-container@uri-account.dfs.core.windows.net/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://demo-container@usgovacct.dfs.core.usgovcloudapi.net/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://demo-container@usgovacct.dfs.core.usgovcloudapi.net/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://demo-container@myaccount.dfs.core.windows.net/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://demo-container@myaccount.dfs.core.windows.net/path/to/data.parquet",
+                "az://myaccount.demo-container/path/to/data.parquet",
+            ),
+            (
+                "abfss://demo-container@myaccount.dfs.core.windows.net/path/to/data.parquet",
+                "az://myaccount.demo-container/path/to/data.parquet",
+            ),
+            ("abfss://demo-container@myaccount.dfs.core.windows.net", "az://myaccount.demo-container"),
+            ("az://demo-container/path", "az://demo-container/path"),
+            # Same container name on different accounts maps to distinct hosts.
+            ("abfss://data@devacct.dfs.core.windows.net/data.csv", "az://devacct.data/data.csv"),
+            ("abfss://data@prodacct.dfs.core.windows.net/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."""

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/tests/unit/common/sql/operators/test_analytics.py

@@ -126,7 +126,7 @@ class TestAnalyticsOperator:
             datasource_configs=[datasource_config],
             queries=["SELECT * FROM users_data"],
             engine=mock_engine,
-            result_output_format=["invalid"],  # type: ignore
+            result_output_format=["invalid"],
         )
 
         with pytest.raises(ValueError, match="Unsupported output format"):

{apache_airflow_providers_common_sql-2.3.0 → apache_airflow_providers_common_sql-2.4.0}/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://container@account.dfs.core.windows.net/path", StorageType.AZURE),
+            ("abfss://container@account.dfs.core.windows.net/path", StorageType.AZURE),
             ("file:///path/to/file", StorageType.LOCAL),
         ],
     )