pankajastro commented on code in PR #73374:
URL: https://github.com/apache/airflow/pull/73374#discussion_r4081237562
##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -196,6 +197,45 @@ def _fetch_extra_configs(keys: list[str]) -> dict[str,
Any]:
key_path = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
credentials = self._remove_none_values({"key_path": key_path,
"keyfile_dict": keyfile_dict})
+ case "wasb":
+ extra_dejson = conn.extra_dejson
+ for unsupported_field in (
+ "connection_string",
+ "managed_identity_client_id",
+ "workload_identity_tenant_id",
+ ):
+ if extra_dejson.get(unsupported_field):
+ raise ValueError(
+ f"Connection field {unsupported_field!r} is not
supported for DataFusion "
+ "Azure Blob Storage access; only
tenant_id+login+password (service "
+ "principal), sas_token,
shared_access_key/account_key/password, or ambient "
+ "credentials (AZURE_* environment variables,
managed identity, workload "
+ "identity, or az login) are used."
+ )
+ credentials = {"account":
self._resolve_wasb_account(conn.host, conn.login)}
+ tenant_id = extra_dejson.get("tenant_id")
+ sas_token = extra_dejson.get("sas_token")
+ if tenant_id and conn.login and conn.password:
+ # client_id/client_secret/tenant_id must all be set
together, or not at all --
+ # DataFusion's binding panics on a partial combination.
+ credentials.update(
Review Comment:
Good catch, thanks — the binding has no way to skip `from_env()`, so added a
guard that raises when the connection supplies an explicit credential while any
of those five env vars is set on the worker, rather than silently
authenticating as the wrong identity.
---
Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting
##########
providers/common/sql/docs/operators.rst:
##########
@@ -362,6 +363,28 @@ resolved in this order:
:start-after: [START howto_analytics_operator_with_gcs]
:end-before: [END howto_analytics_operator_with_gcs]
+Azure Storage
+-------------
+Use an ``az://`` URI with a ``conn_id`` pointing to a ``wasb`` connection.
+``abfs://`` and ``abfss://`` URIs are not recognized yet. Credentials are
+resolved in this order:
+
+1. Azure AD service principal -- ``tenant_id`` extra, with ``login`` as the
+ client ID and ``password`` as the client secret
+2. SAS token -- ``sas_token`` extra, as a query string
+3. Shared key -- ``password``, or the ``shared_access_key``/``account_key``
extra
+4. Ambient credentials -- ``AZURE_*`` environment variables, managed identity,
Review Comment:
Fixed, thanks — docs now say the environment can outrank the connection, and
that the Azure CLI needs `AZURE_USE_AZURE_CLI=true` (default ambient path is
IMDS).
---
Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting
##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -205,6 +245,28 @@ def _remove_none_values(params: dict[str, Any]) ->
dict[str, Any]:
"""Filter out None values from the dictionary."""
return {k: v for k, v in params.items() if v is not None}
+ @staticmethod
+ def _resolve_wasb_account(host: str | None, login: str | None) -> str:
+ """
+ Return the storage account name the way WasbHook resolves it.
+
+ From ``host`` when set (its netloc's first label), falling back to
``login`` only when
+ ``host`` is empty -- login holds the service-principal client_id in
that auth mode, not
+ the account name. Reimplemented locally rather than importing
+ ``airflow.providers.microsoft.azure.utils.parse_blob_account_url``, to
avoid pulling the
+ microsoft-azure provider's full Azure SDK dependency stack into
common-sql for one string
+ operation that only needs the stdlib.
+ """
+ netloc = urlsplit(host if host else
f"https://{login}.blob.core.windows.net/";).netloc
+ if not netloc:
+ # No scheme was given (e.g. a bare DNS name); urlsplit put it all
in the path instead.
+ netloc = urlsplit(f"https://{host}";).netloc
+ if "." not in netloc:
+ # Only an Active Directory ID was given, not a full URL or DNS
name.
+ netloc = f"{login}.blob.core.windows.net"
+ # Azure storage account names are capped at 24 characters.
+ return netloc.split(".", 1)[0][:24]
Review Comment:
Confirmed both repro cases, thanks. Rather than try to support arbitrary
endpoints, it now raises when the resolved host isn't
`*.blob.core.windows.net`, pointing at `AZURE_STORAGE_ENDPOINT` — matches this
PR's existing "raise instead of silently redirect" approach.
---
Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting
##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -196,6 +197,45 @@ def _fetch_extra_configs(keys: list[str]) -> dict[str,
Any]:
key_path = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
credentials = self._remove_none_values({"key_path": key_path,
"keyfile_dict": keyfile_dict})
+ case "wasb":
+ extra_dejson = conn.extra_dejson
+ for unsupported_field in (
+ "connection_string",
+ "managed_identity_client_id",
+ "workload_identity_tenant_id",
+ ):
+ if extra_dejson.get(unsupported_field):
+ raise ValueError(
+ f"Connection field {unsupported_field!r} is not
supported for DataFusion "
+ "Azure Blob Storage access; only
tenant_id+login+password (service "
+ "principal), sas_token,
shared_access_key/account_key/password, or ambient "
+ "credentials (AZURE_* environment variables,
managed identity, workload "
+ "identity, or az login) are used."
+ )
+ credentials = {"account":
self._resolve_wasb_account(conn.host, conn.login)}
+ tenant_id = extra_dejson.get("tenant_id")
+ sas_token = extra_dejson.get("sas_token")
+ if tenant_id and conn.login and conn.password:
Review Comment:
Good point, thanks — now raises naming whichever of `login`/`password` is
missing, instead of falling through to ambient auth or sending the client
secret as a shared key. Flipped the test to `pytest.raises` as suggested.
---
Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting
##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -196,6 +197,45 @@ def _fetch_extra_configs(keys: list[str]) -> dict[str,
Any]:
key_path = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
credentials = self._remove_none_values({"key_path": key_path,
"keyfile_dict": keyfile_dict})
+ case "wasb":
+ extra_dejson = conn.extra_dejson
+ for unsupported_field in (
+ "connection_string",
+ "managed_identity_client_id",
+ "workload_identity_tenant_id",
+ ):
+ if extra_dejson.get(unsupported_field):
+ raise ValueError(
+ f"Connection field {unsupported_field!r} is not
supported for DataFusion "
+ "Azure Blob Storage access; only
tenant_id+login+password (service "
+ "principal), sas_token,
shared_access_key/account_key/password, or ambient "
+ "credentials (AZURE_* environment variables,
managed identity, workload "
+ "identity, or az login) are used."
+ )
+ credentials = {"account":
self._resolve_wasb_account(conn.host, conn.login)}
Review Comment:
Real regression from my earlier fix, thanks for catching it —
`_resolve_wasb_account` now returns `None` when both `host` and `login` are
empty, so it's dropped from the dict like before instead of sending the literal
string "None".
---
Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting
##########
providers/common/sql/tests/unit/common/sql/datafusion/test_engine.py:
##########
@@ -372,6 +373,145 @@ def
test_get_credentials_gcs_rejects_unsupported_identity_fields(self, unsupport
with pytest.raises(ValueError, match=f"{unsupported_field!r} is not
supported"):
engine._get_credentials(mock_conn)
+ def test_get_credentials_azure_with_shared_key(self):
+ mock_conn = MagicMock()
+ mock_conn.conn_type = "wasb"
+ mock_conn.host = None
+ mock_conn.login = "myaccount"
+ mock_conn.password = "mykey"
+ mock_conn.extra_dejson = {}
+ engine = DataFusionEngine()
+
+ credentials, extra_config = engine._get_credentials(mock_conn)
+
+ assert credentials == {"account": "myaccount", "access_key": "mykey"}
+ assert extra_config == {}
+
+ def test_get_credentials_azure_with_shared_access_key_extra(self):
+ mock_conn = MagicMock()
+ mock_conn.conn_type = "wasb"
+ mock_conn.host = None
+ mock_conn.login = "myaccount"
+ mock_conn.password = None
+ mock_conn.extra_dejson = {"shared_access_key": "extra-key"}
+ engine = DataFusionEngine()
+
+ credentials, extra_config = engine._get_credentials(mock_conn)
+
+ assert credentials == {"account": "myaccount", "access_key":
"extra-key"}
+ assert extra_config == {}
+
+ def test_get_credentials_azure_with_service_principal(self):
+ mock_conn = MagicMock()
+ mock_conn.conn_type = "wasb"
+ mock_conn.host = None
+ mock_conn.login = "client-id"
+ mock_conn.password = "client-secret"
+ mock_conn.extra_dejson = {"tenant_id": "tenant-id"}
+ engine = DataFusionEngine()
+
+ 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",
+ }
+ assert extra_config == {}
+
+ def
test_get_credentials_azure_with_service_principal_and_host_prefers_host_account(self):
+ mock_conn = MagicMock()
+ mock_conn.conn_type = "wasb"
+ mock_conn.host = "realaccount.blob.core.windows.net"
+ mock_conn.login = "11111111-2222-3333-4444-555555555555"
+ mock_conn.password = "client-secret"
+ mock_conn.extra_dejson = {"tenant_id": "tenant-id"}
+ engine = DataFusionEngine()
+
+ credentials, extra_config = engine._get_credentials(mock_conn)
+
+ assert credentials == {
+ "account": "realaccount",
+ "client_id": "11111111-2222-3333-4444-555555555555",
+ "client_secret": "client-secret",
+ "tenant_id": "tenant-id",
+ }
+ assert extra_config == {}
+
+ def test_get_credentials_azure_tenant_id_without_login_falls_back(self):
+ """A partial service-principal config (tenant_id alone) must not be
forwarded --
+ DataFusion's binding panics on a partial
client_id/client_secret/tenant_id combination."""
+ mock_conn = MagicMock()
+ mock_conn.conn_type = "wasb"
+ mock_conn.host = None
+ mock_conn.login = None
+ mock_conn.password = None
+ mock_conn.extra_dejson = {"tenant_id": "tenant-id"}
+ engine = DataFusionEngine()
+
+ credentials, extra_config = engine._get_credentials(mock_conn)
+
+ assert "tenant_id" not in credentials
Review Comment:
Added a full-dict assertion for the empty case, plus parametrized tests
calling `_resolve_wasb_account` directly for the URL/path, sovereign-suffix,
no-dot, and 24-char-cap cases, and one connection with both `sas_token` and
`shared_access_key` set to pin the precedence.
---
Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting
##########
providers/common/sql/docs/operators.rst:
##########
@@ -362,6 +363,28 @@ resolved in this order:
:start-after: [START howto_analytics_operator_with_gcs]
:end-before: [END howto_analytics_operator_with_gcs]
+Azure Storage
+-------------
+Use an ``az://`` URI with a ``conn_id`` pointing to a ``wasb`` connection.
+``abfs://`` and ``abfss://`` URIs are not recognized yet. Credentials are
+resolved in this order:
+
+1. Azure AD service principal -- ``tenant_id`` extra, with ``login`` as the
+ client ID and ``password`` as the client secret
+2. SAS token -- ``sas_token`` extra, as a query string
+3. Shared key -- ``password``, or the ``shared_access_key``/``account_key``
extra
+4. Ambient credentials -- ``AZURE_*`` environment variables, managed identity,
+ workload identity, or the Azure CLI
+
+``connection_string``, ``managed_identity_client_id``,
``workload_identity_tenant_id``,
+and a URL-form ``sas_token`` are not supported.
Review Comment:
Added both, thanks — docs now say the account comes from `host` first, else
`login`, and that `client_secret_auth_config` isn't read here.
---
Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting
--
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]
