Skip to content

Commit 5487c08

Browse files
pankajastroclaude
andauthored
Add Azure Blob Storage support to the DataFusion object storage layer (#73374)
* Add Azure Blob Storage support to the DataFusion object storage layer DataFusion credential resolution only handled AWS and GCS, so a wasb-backed DataSourceConfig raised Unknown connection type instead of working, even though DataFusion's storage layer supports Azure. Add a StorageType.AZURE, an AzureObjectStorageProvider using DataFusion's MicrosoftAzure, and a wasb credential branch resolving the connection's shared key, SAS token, or Azure AD service principal, falling back to ambient auth when none are set. The service principal fields (client_id, client_secret, tenant_id) must all be set together or not at all -- DataFusion's binding panics on a partial combination -- so they're only forwarded when the connection has all three (tenant_id extra, login, and password). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Remove unused WasbHook install guard from DataFusion Azure credentials The wasb credential branch never called WasbHook -- it only imported it to detect whether apache-airflow-providers-microsoft-azure was installed, unlike the aws branch which actually calls AwsGenericHook. Since Azure credentials come entirely from the generic Connection object, this guard and its extra dependency are unnecessary, matching how the GCS branch already works without a provider-specific import. * Reject unsupported Azure identity fields in DataFusion credential resolution DataFusion's Azure binding has no equivalent for a raw connection string, a specific managed/workload identity, or a URL-form SAS token, so these connection extra fields were silently dropped and the code fell back to whatever ambient credentials were available. That can authenticate as a different identity than the one the connection was configured for. Raise a clear error instead, matching the same guard already used for unsupported GCS identity fields. * Fix ambiguous object-storage list in LLMSchemaCompareOperator docs Missing comma made "Azure Blob Storage Parquet" read as one item instead of two separate examples in the supported source list. * Document Azure Blob Storage support in the Analytics operator docs The Analytics operator docs listed Azure as supported in the intro line but not in the Supported Storage Systems list, still flagged it under "not yet supported", and had no usage example or credential docs, unlike S3 and GCS. * Fix DataFusion Azure account resolution ignoring connection host The wasb credential branch took the storage account name unconditionally from the connection's login field. For service-principal auth, login holds the client_id, not the account name, so a connection whose host correctly named the storage account still had the object store target the client_id GUID instead. A connection with no login at all (key/SAS auth) silently fell back to the AZURE_STORAGE_ACCOUNT_NAME environment variable. Resolve the account from host first, matching WasbHook's own behavior, and fall back to login only when host is empty. Also format the storage type name consistently with .value (avoids drifting str-Enum formatting across Python versions) and move the sas_token parsing import to the top of the file. * Document that DataFusion's Azure storage only recognizes az:// URIs abfs:// and abfss:// embed the storage account in the URI itself, a different shape from az:// that engine.py does not parse; using them currently fails with an unhelpful "Unsupported storage type" error with no hint that Azure support exists under a different scheme. * Fix silent identity/endpoint mismatches in DataFusion Azure credentials Four issues surfaced by review, each verified against the real datafusion binding and object_store source: - An empty connection (no host or login -- the shape of the wasb_default connection airflow db creates) sent the literal string "None" as the account instead of omitting it, breaking the fallback to AZURE_STORAGE_ACCOUNT_NAME that worked before the host-priority fix. - A partial service-principal config (tenant_id set but login or password missing) silently fell through to ambient auth, or sent the client secret as a shared key, instead of raising -- authenticating as a different identity than the one configured. - The account name kept only the first DNS label, so a sovereign-cloud host (*.chinacloudapi.cn, *.usgovcloudapi.net) or the Azurite emulator's loopback URL was silently redirected to the public *.blob.core.windows.net account of the same name; the binding has no endpoint override to correct this after the fact. - DataFusion's MicrosoftAzure always reads AZURE_* worker environment variables before overlaying the connection's credential, and checks an environment-derived access key or workload-identity token ahead of the connection's SAS token or client secret. In the common AKS workload-identity deployment, this meant the connection's credential was silently ignored in favor of the pod's identity. The binding has no way to skip its own environment read, so this can only be caught, not fixed. * Read the legacy extra__<conn_type>__ prefix for wasb and GCS credentials WasbHook and GoogleBaseHook still fall back to fields written as extra__wasb__<name> / extra__google_cloud_platform__<name> -- the spelling older Airflow connection UIs used for custom extra fields. This code only looked at the bare key, so a connection created that way had its sas_token or connection_string silently invisible: the unsupported-field guard never saw a prefixed connection_string, and a prefixed sas_token fell through to ambient auth instead of being used. * Trim two overlong inline comments in Azure credential resolution Same facts, fewer words -- the mechanism and consequence were spelled out twice in slightly different phrasing. * Simplify the Azure env-precedence docs paragraph Same facts, plainer words -- "outrank", "ahead of", and "consulted last" weren't adding anything a simpler verb wouldn't. * Attribute AZURE_USE_AZURE_CLI to object_store, not Airflow Names the actual library so it's clear this is DataFusion's underlying binding's behavior, not an Airflow-defined setting, and contrasts it directly with WasbHook's automatic az login to explain why it's worth calling out at all. * Fix docs spell-check failure on British spelling "honours" The docs spellcheck job uses an American-English dictionary; "honors" is already the spelling used elsewhere in this same doc tree. * Simplify wasb credential resolution per review Collapse the legacy-field lookup into one dict.get() call, and use the walrus operator so tenant_id/sas_token are only computed where they're actually checked, instead of unconditionally ahead of the branch. Also consolidate the three shared-key test variants (password, shared_access_key extra, account_key extra) into one parametrized test, picking up account_key coverage that was previously untested. * Fix Azure Blob Storage credential precedence gaps in DataFusion The env-var precedence guard was missing two real object_store aliases (AZURE_STORAGE_MASTER_KEY and the AZURE_CLIENT_ID/SECRET/TENANT_ID triple, including the AKS workload-identity webhook's own env vars) that silently outrank a connection's SAS token, while AZURE_STORAGE_SAS_KEY never can and was raising for nothing. The account-resolution check also rejected a non-public-cloud host even when the worker had already set AZURE_STORAGE_ENDPOINT to fix exactly that, so the documented workaround never actually worked. Finally, the SAS-token branch was checked ahead of the shared-key extra, the opposite of WasbHook.get_conn's order -- a connection with a current shared_access_key but an expired SAS token would work through the hook and fail here. Co-Authored-By: Claude <noreply@anthropic.com> * Move Azure env-var precedence constants next to their only caller They were declared next to the unrelated _AZURE_PUBLIC_SUFFIX constant, ~180 lines from _get_credentials, the only place that reads them -- inconsistent with the file's existing convention of keeping a constant next to its sole user. Co-Authored-By: Claude <noreply@anthropic.com> * Make the Azure credential env-var guard aware of which tier it protects The flat guard list raised for env vars that could never actually outrank a connection's own credential in object_store's precedence chain -- most notably AZURE_FEDERATED_TOKEN_FILE against a plain shared-key connection, which blocks the single most common Azure Kubernetes deployment shape (the workload-identity webhook injects it into every labelled pod) even though the connection's own key always wins there. Also closes a gap in the sovereign-cloud/endpoint-override check: an Azurite-style host (bare hostname or host:port, with no login) fell through an unrelated Active-Directory-ID fallback and silently resolved to the literal string "None" instead of raising. Co-Authored-By: Claude <noreply@anthropic.com> * Consolidate Azure env-precedence tests into one parametrized case per tier Each credential tier (access_key, service_principal, SAS) had a "rejects when outranked" test and a separate "ignores when it can't be outranked" test that only differed in which env vars were set and whether the call was expected to raise. Folding both into a single should_raise-parametrized test per tier keeps the same coverage with far less repetition. Co-Authored-By: Claude <noreply@anthropic.com> * De-nest the endpoint-override check in _resolve_wasb_account Left over from incrementally adding the emulator-host colon check; the suffix check and the endpoint-override check were always meant to be one combined condition, not a nested if. Co-Authored-By: Claude <noreply@anthropic.com> * Fix inaccurate "dead" claim about AZURE_STORAGE_SAS_KEY in test docstring object_store's resolve_sas_token() does read this variable into self.sas_key; it's just checked after sas_query_pairs within that same step, so it never wins once the connection sets sas_query_pairs. "Dead" overstated that as unused entirely, which could mislead a future reader into thinking the safety argument doesn't depend on precedence order. Co-Authored-By: Claude <noreply@anthropic.com> --------- Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
1 parent 4083290 commit 5487c08

11 files changed

Lines changed: 793 additions & 6 deletions

File tree

‎providers/common/ai/docs/operators/llm_schema_compare.rst‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ With Object Storage or a Database Table
7070

7171
Use ``data_sources`` with
7272
:class:`~airflow.providers.common.sql.config.DataSourceConfig` to include
73-
object-storage sources (S3 or GCS Parquet, CSV, Iceberg, etc.) in the comparison.
73+
object-storage sources (S3, GCS, Azure Blob Storage, Parquet, CSV, Iceberg, etc.) in the comparison.
7474
These can be freely combined with ``db_conn_ids``. Whether an entry is
7575
introspected via ``DbApiHook`` or DataFusion depends on what its ``conn_id``
7676
resolves to, not on its ``uri``/``format`` fields: a ``DataSourceConfig``

‎providers/common/ai/docs/operators/llm_sql.rst‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ With Object Storage
6565
-------------------
6666

6767
Use ``datasource_config`` to generate queries for data stored in object storage
68-
(e.g., S3, GCS, local filesystem) via `DataFusion <https://datafusion.apache.org/>`_.
68+
(e.g., S3, GCS, Azure Blob Storage, local filesystem) via `DataFusion <https://datafusion.apache.org/>`_.
6969
The operator uses :class:`~airflow.providers.common.sql.config.DataSourceConfig`
7070
to register the object storage source as a table so the LLM can include it in
7171
the schema context.

‎providers/common/ai/docs/toolsets/datafusion.rst‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ Files with DataFusion: ``DataFusionToolset``
2626
Curated toolset wrapping
2727
:class:`~airflow.providers.common.sql.datafusion.engine.DataFusionEngine`
2828
with three tools (``list_tables``, ``get_schema``, and ``query``) for
29-
querying files on object stores (S3, GCS, local filesystem, Iceberg) via Apache DataFusion.
29+
querying files on object stores (S3, GCS, Azure Blob Storage, local filesystem, Iceberg) via Apache DataFusion.
3030

3131
.. list-table::
3232
:header-rows: 1

‎providers/common/sql/docs/operators.rst‎

Lines changed: 59 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -297,10 +297,11 @@ Supported Storage Systems
297297
-------------------------
298298
- S3
299299
- GCS
300+
- Azure Blob Storage
300301
- Local File System
301302

302303
.. note::
303-
Azure, HTTP, Delta are not yet supported but will be added in the future.
304+
HTTP, Delta are not yet supported but will be added in the future.
304305

305306

306307

@@ -362,6 +363,63 @@ resolved in this order:
362363
:start-after: [START howto_analytics_operator_with_gcs]
363364
:end-before: [END howto_analytics_operator_with_gcs]
364365

366+
Azure Storage
367+
-------------
368+
Use an ``az://`` URI with a ``conn_id`` pointing to a ``wasb`` connection.
369+
``abfs://`` and ``abfss://`` URIs are not recognized yet. The account name
370+
comes from ``host`` (its first DNS label) when set, falling back to
371+
``login`` only when ``host`` is empty; only the public
372+
``*.blob.core.windows.net`` cloud is supported, since DataFusion's binding
373+
otherwise has no endpoint override -- unless the worker sets
374+
``AZURE_STORAGE_ENDPOINT``/``AZURE_ENDPOINT`` for a real sovereign-cloud
375+
hostname. A ``host:port`` address (the Azurite emulator's shape, e.g.
376+
``azurite:10000``) is never a real DNS hostname and always raises instead,
377+
even with those variables set: put the account name in ``login`` with
378+
``host`` left empty for Azurite, alongside those same variables and
379+
``AZURE_ALLOW_HTTP=true``. ``client_secret_auth_config`` (the authority
380+
override ``WasbHook`` honors) is not read here.
381+
382+
The connection supplies one of the following credentials, checked in this
383+
order (matching ``WasbHook.get_conn``):
384+
385+
1. Azure AD service principal -- ``tenant_id`` extra, with ``login`` as the
386+
client ID and ``password`` as the client secret (both required together)
387+
2. Shared key -- the ``shared_access_key`` extra
388+
3. SAS token -- ``sas_token`` extra, as a query string
389+
4. Shared key -- ``password``, or the ``account_key`` extra
390+
5. None of the above -- ambient auth (see below)
391+
392+
**A worker environment variable can override the connection, but only from a
393+
higher-priority tier.** DataFusion reads ``AZURE_*`` environment variables
394+
first, and resolves credentials in this priority order:
395+
396+
1. Bearer token
397+
2. Access key -- the shared-key connection's tier
398+
3. Workload identity
399+
4. Client secret -- the service-principal connection's tier
400+
5. SAS -- the SAS connection's tier
401+
402+
Only a complete credential from an earlier tier can override the
403+
connection -- a single bearer or access-key variable, or every variable a
404+
multi-field tier needs. This raises when found, naming the variables,
405+
instead of silently using the wrong identity.
406+
407+
With no explicit credential, authentication falls back to ``AZURE_*``
408+
environment variables, managed identity, or workload identity. Unlike
409+
``WasbHook`` (which tries ``az login`` automatically), DataFusion's
410+
underlying ``object_store`` binding only tries the Azure CLI if
411+
``AZURE_USE_AZURE_CLI=true`` is set; otherwise it defaults straight to
412+
IMDS managed identity.
413+
414+
``connection_string``, ``managed_identity_client_id``, ``workload_identity_tenant_id``,
415+
and a URL-form ``sas_token`` are not supported.
416+
417+
.. exampleinclude:: /../../sql/src/airflow/providers/common/sql/example_dags/example_analytics.py
418+
:language: python
419+
:dedent: 4
420+
:start-after: [START howto_analytics_operator_with_azure]
421+
:end-before: [END howto_analytics_operator_with_azure]
422+
365423
Local File System Storage
366424
-------------------------
367425
.. exampleinclude:: /../../sql/src/airflow/providers/common/sql/example_dags/example_analytics.py

‎providers/common/sql/src/airflow/providers/common/sql/config.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ class StorageType(str, Enum):
4848

4949
S3 = "s3"
5050
GCS = "gcs"
51+
AZURE = "azure"
5152
LOCAL = "local"
5253

5354

@@ -117,6 +118,8 @@ def _extract_storage_type(self) -> StorageType | None:
117118
return StorageType.S3
118119
if self.uri.startswith("gs://"):
119120
return StorageType.GCS
121+
if self.uri.startswith("az://"):
122+
return StorageType.AZURE
120123
if self.uri.startswith("file://"):
121124
return StorageType.LOCAL
122125
raise ValueError(f"Unsupported storage type for URI: {self.uri}")

‎providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py‎

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import os
2020
from typing import TYPE_CHECKING, Any
21+
from urllib.parse import parse_qsl, urlsplit
2122

2223
from datafusion import SessionContext
2324

@@ -161,6 +162,12 @@ def _get_gcp_extra_field(extra_dejson: dict[str, Any], field_name: str) -> Any:
161162
return extra_dejson[field_name]
162163
return extra_dejson.get(f"extra__google_cloud_platform__{field_name}")
163164

165+
def _get_wasb_extra_field(extra_dejson: dict[str, Any], field_name: str) -> Any:
166+
# Older Airflow connection UIs wrote custom extra fields as
167+
# extra__wasb__<field_name> instead of the bare key; WasbHook still reads that
168+
# legacy spelling as a fallback, so this must too.
169+
return extra_dejson.get(field_name, extra_dejson.get(f"extra__wasb__{field_name}"))
170+
164171
match conn.conn_type:
165172
case "aws":
166173
try:
@@ -204,6 +211,66 @@ def _get_gcp_extra_field(extra_dejson: dict[str, Any], field_name: str) -> Any:
204211
key_path = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
205212
credentials = self._remove_none_values({"key_path": key_path, "keyfile_dict": keyfile_dict})
206213

214+
case "wasb":
215+
extra_dejson = conn.extra_dejson
216+
for unsupported_field in (
217+
"connection_string",
218+
"managed_identity_client_id",
219+
"workload_identity_tenant_id",
220+
):
221+
if _get_wasb_extra_field(extra_dejson, unsupported_field):
222+
raise ValueError(
223+
f"Connection field {unsupported_field!r} is not supported for DataFusion "
224+
"Azure Blob Storage access; only tenant_id+login+password (service "
225+
"principal), sas_token, shared_access_key/account_key/password, or ambient "
226+
"credentials (AZURE_* environment variables, managed identity, workload "
227+
"identity, or az login) are used."
228+
)
229+
credentials = {"account": self._resolve_wasb_account(conn.host, conn.login)}
230+
credential_tier: str | None = None
231+
if tenant_id := _get_wasb_extra_field(extra_dejson, "tenant_id"):
232+
if not conn.login or not conn.password:
233+
# Falling through here would silently switch identity (ambient auth, or
234+
# the client secret sent as a shared key) instead of failing clearly.
235+
missing = "login (client_id)" if not conn.login else "password (client_secret)"
236+
raise ValueError(
237+
f"Connection extra 'tenant_id' is set for DataFusion Azure Blob Storage "
238+
f"service-principal auth, but {missing} is not."
239+
)
240+
credentials.update(
241+
{"client_id": conn.login, "client_secret": conn.password, "tenant_id": tenant_id}
242+
)
243+
credential_tier = "client_secret"
244+
elif shared_access_key := _get_wasb_extra_field(extra_dejson, "shared_access_key"):
245+
# Checked ahead of sas_token to match WasbHook.get_conn's precedence.
246+
credentials["access_key"] = shared_access_key
247+
credential_tier = "access_key"
248+
elif sas_token := _get_wasb_extra_field(extra_dejson, "sas_token"):
249+
if sas_token.startswith("http"):
250+
raise ValueError(
251+
"A URL-form `sas_token` is not supported for DataFusion Azure Blob Storage "
252+
"access; provide the SAS token as a query string instead."
253+
)
254+
credentials["sas_query_pairs"] = parse_qsl(sas_token.lstrip("?"))
255+
credential_tier = "sas"
256+
else:
257+
access_key = conn.password or _get_wasb_extra_field(extra_dejson, "account_key")
258+
if access_key:
259+
credentials["access_key"] = access_key
260+
credential_tier = "access_key"
261+
262+
if credential_tier is not None:
263+
conflicting_env_vars = self._find_conflicting_azure_env_vars(credential_tier)
264+
if conflicting_env_vars:
265+
raise ValueError(
266+
f"Worker environment variable(s) {', '.join(conflicting_env_vars)} would "
267+
"silently take precedence over this connection's explicit credential in "
268+
"DataFusion's Azure Blob Storage binding. Unset them on the worker, or "
269+
"remove the explicit credential from this connection to rely on the "
270+
"environment instead."
271+
)
272+
credentials = self._remove_none_values(credentials)
273+
207274
case _:
208275
raise ValueError(f"Unknown connection type {conn.conn_type}")
209276
return credentials, extra_config
@@ -213,6 +280,124 @@ def _remove_none_values(params: dict[str, Any]) -> dict[str, Any]:
213280
"""Filter out None values from the dictionary."""
214281
return {k: v for k, v in params.items() if v is not None}
215282

283+
_AZURE_PUBLIC_SUFFIX = ".blob.core.windows.net"
284+
285+
# object_store's build() precedence, high to low: bearer token > access key > workload
286+
# identity (client_id+tenant_id+federated_token_file) > client secret
287+
# (client_id+client_secret+tenant_id) > SAS. Used by _find_conflicting_azure_env_vars.
288+
_AZURE_ENV_BEARER_VARS = ("AZURE_STORAGE_TOKEN",)
289+
_AZURE_ENV_ACCESS_KEY_VARS = (
290+
"AZURE_STORAGE_ACCOUNT_KEY",
291+
"AZURE_STORAGE_ACCESS_KEY",
292+
"AZURE_STORAGE_MASTER_KEY",
293+
)
294+
_AZURE_ENV_CLIENT_ID_VARS = ("AZURE_STORAGE_CLIENT_ID", "AZURE_CLIENT_ID")
295+
_AZURE_ENV_CLIENT_SECRET_VARS = ("AZURE_STORAGE_CLIENT_SECRET", "AZURE_CLIENT_SECRET")
296+
_AZURE_ENV_TENANT_ID_VARS = (
297+
"AZURE_STORAGE_TENANT_ID",
298+
"AZURE_STORAGE_AUTHORITY_ID",
299+
"AZURE_TENANT_ID",
300+
"AZURE_AUTHORITY_ID",
301+
)
302+
_AZURE_ENV_FEDERATED_TOKEN_FILE_VAR = "AZURE_FEDERATED_TOKEN_FILE"
303+
304+
@classmethod
305+
def _find_conflicting_azure_env_vars(cls, credential_tier: str) -> list[str]:
306+
"""
307+
Return worker env vars that would silently outrank the connection's own credential.
308+
309+
The binding always calls `from_env()` with no way to skip it, so this can only be
310+
caught here, not avoided. The connection's fields overwrite the same fields from
311+
`from_env()`, so only a tier above the connection's own, or an unoccupied tier
312+
(workload identity or client secret, for SAS) with *all* its fields present, can
313+
actually take over.
314+
"""
315+
316+
def env_set(*var_groups: tuple[str, ...]) -> list[str]:
317+
return [var for group in var_groups for var in group if os.environ.get(var)]
318+
319+
conflicting = env_set(cls._AZURE_ENV_BEARER_VARS)
320+
if credential_tier == "access_key":
321+
return conflicting
322+
conflicting += env_set(cls._AZURE_ENV_ACCESS_KEY_VARS)
323+
if credential_tier == "client_secret":
324+
# client_id and tenant_id are already the connection's own; only the federated
325+
# token file is left for env to complete the workload-identity triple with.
326+
if os.environ.get(cls._AZURE_ENV_FEDERATED_TOKEN_FILE_VAR):
327+
conflicting.append(cls._AZURE_ENV_FEDERATED_TOKEN_FILE_VAR)
328+
return conflicting
329+
# SAS occupies none of these fields, so each whole triple must come from env.
330+
if (
331+
any(os.environ.get(var) for var in cls._AZURE_ENV_CLIENT_ID_VARS)
332+
and any(os.environ.get(var) for var in cls._AZURE_ENV_TENANT_ID_VARS)
333+
and os.environ.get(cls._AZURE_ENV_FEDERATED_TOKEN_FILE_VAR)
334+
):
335+
conflicting += env_set(cls._AZURE_ENV_CLIENT_ID_VARS, cls._AZURE_ENV_TENANT_ID_VARS)
336+
conflicting.append(cls._AZURE_ENV_FEDERATED_TOKEN_FILE_VAR)
337+
if (
338+
any(os.environ.get(var) for var in cls._AZURE_ENV_CLIENT_ID_VARS)
339+
and any(os.environ.get(var) for var in cls._AZURE_ENV_CLIENT_SECRET_VARS)
340+
and any(os.environ.get(var) for var in cls._AZURE_ENV_TENANT_ID_VARS)
341+
):
342+
conflicting += env_set(
343+
cls._AZURE_ENV_CLIENT_ID_VARS,
344+
cls._AZURE_ENV_CLIENT_SECRET_VARS,
345+
cls._AZURE_ENV_TENANT_ID_VARS,
346+
)
347+
return list(dict.fromkeys(conflicting))
348+
349+
@classmethod
350+
def _resolve_wasb_account(cls, host: str | None, login: str | None) -> str | None:
351+
"""
352+
Return the storage account name the way WasbHook resolves it.
353+
354+
From ``host`` (its netloc's first label) when set, falling back to ``login`` only when
355+
``host`` is empty -- login holds the service-principal client_id otherwise, not the
356+
account name. Returns ``None`` when neither is set, so the binding falls back to
357+
``AZURE_STORAGE_ACCOUNT_NAME`` rather than the literal string ``"None"``. Reimplemented
358+
locally instead of importing
359+
``airflow.providers.microsoft.azure.utils.parse_blob_account_url``, to avoid pulling in
360+
the microsoft-azure provider's Azure SDK dependency for one stdlib string operation.
361+
362+
Only the public ``*.blob.core.windows.net`` cloud is supported, unless
363+
``AZURE_STORAGE_ENDPOINT``/``AZURE_ENDPOINT`` is set for a real sovereign-cloud hostname
364+
(DataFusion's binding otherwise has no endpoint override). A ``host:port`` netloc (the
365+
Azurite emulator's shape) is never a real hostname, so it always raises instead --
366+
checked before the dotless-host fallback below, which would otherwise mask it.
367+
"""
368+
if not host and not login:
369+
return None
370+
netloc = urlsplit(host if host else f"https://{login}.blob.core.windows.net/").netloc
371+
if not netloc:
372+
# No scheme was given (e.g. a bare DNS name); urlsplit put it all in the path instead.
373+
netloc = urlsplit(f"https://{host}").netloc
374+
if host and ":" in netloc:
375+
raise ValueError(
376+
f"Connection host {host!r} looks like an emulator address (host:port), which "
377+
"DataFusion's Azure Blob Storage binding cannot resolve an account name from -- "
378+
"even with AZURE_STORAGE_ENDPOINT set. Put the account name in `login` with "
379+
"`host` empty instead, alongside AZURE_STORAGE_ENDPOINT and AZURE_ALLOW_HTTP=true."
380+
)
381+
if "." not in netloc:
382+
if not login:
383+
raise ValueError(
384+
f"Connection host {host!r} is not a full URL or DNS name, and no `login` was "
385+
"given to resolve it as an Active Directory ID instead."
386+
)
387+
# Only an Active Directory ID was given, not a full URL or DNS name.
388+
netloc = f"{login}.blob.core.windows.net"
389+
if not netloc.endswith(cls._AZURE_PUBLIC_SUFFIX) and not (
390+
os.environ.get("AZURE_STORAGE_ENDPOINT") or os.environ.get("AZURE_ENDPOINT")
391+
):
392+
raise ValueError(
393+
f"Connection host {host!r} does not resolve to the public {cls._AZURE_PUBLIC_SUFFIX} "
394+
"cloud, which is the only one DataFusion's Azure Blob Storage binding can target (it "
395+
"has no endpoint override). A sovereign cloud is supported once the "
396+
"AZURE_STORAGE_ENDPOINT environment variable is set."
397+
)
398+
# Azure storage account names are capped at 24 characters.
399+
return netloc.split(".", 1)[0][:24]
400+
216401
def get_schema(self, table_name: str):
217402
"""Get the schema of a table."""
218403
schema = str(self.session_context.table(table_name).schema())

‎providers/common/sql/src/airflow/providers/common/sql/datafusion/object_storage_provider.py‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
import tempfile
2121
from pathlib import Path
2222

23-
from datafusion.object_store import AmazonS3, GoogleCloud, LocalFileSystem
23+
from datafusion.object_store import AmazonS3, GoogleCloud, LocalFileSystem, MicrosoftAzure
2424

2525
from airflow.providers.common.sql.config import ConnectionConfig, StorageType
2626
from airflow.providers.common.sql.datafusion.base import ObjectStorageProvider
@@ -107,6 +107,37 @@ def get_scheme(self) -> str:
107107
return "gs://"
108108

109109

110+
class AzureObjectStorageProvider(ObjectStorageProvider):
111+
"""Azure Object Storage Provider using DataFusion's MicrosoftAzure."""
112+
113+
@property
114+
def get_storage_type(self) -> StorageType:
115+
"""Return the storage type."""
116+
return StorageType.AZURE
117+
118+
def create_object_store(self, path: str, connection_config: ConnectionConfig | None = None):
119+
"""Create an Azure object store using DataFusion's MicrosoftAzure."""
120+
if connection_config is None:
121+
raise ValueError(f"connection_config must be provided for {self.get_storage_type.value}")
122+
123+
try:
124+
credentials = connection_config.credentials
125+
container = self.get_bucket(path)
126+
127+
azure_store = MicrosoftAzure(container_name=container, **credentials)
128+
self.log.info("Created Azure object store for container %s", container)
129+
130+
return azure_store
131+
132+
except BaseException as e:
133+
# A bad credential combination panics as pyo3_runtime.PanicException, not a plain Exception.
134+
raise ObjectStoreCreationException(f"Failed to create Azure object store: {e}")
135+
136+
def get_scheme(self) -> str:
137+
"""Return the scheme for Azure."""
138+
return "az://"
139+
140+
110141
class LocalObjectStorageProvider(ObjectStorageProvider):
111142
"""Local Object Storage Provider using DataFusion's LocalFileSystem."""
112143

@@ -126,10 +157,11 @@ def get_scheme(self) -> str:
126157

127158
def get_object_storage_provider(storage_type: StorageType) -> ObjectStorageProvider:
128159
"""Get an object storage provider based on the storage type."""
129-
# TODO: Add support for Azure, HTTP: https://datafusion.apache.org/python/autoapi/datafusion/object_store/index.html
160+
# TODO: Add support for HTTP: https://datafusion.apache.org/python/autoapi/datafusion/object_store/index.html
130161
providers: dict[StorageType, type] = {
131162
StorageType.S3: S3ObjectStorageProvider,
132163
StorageType.GCS: GCSObjectStorageProvider,
164+
StorageType.AZURE: AzureObjectStorageProvider,
133165
StorageType.LOCAL: LocalObjectStorageProvider,
134166
}
135167

0 commit comments

Comments
 (0)