mirror of
https://github.com/apache/airflow.git
synced 2026-11-04 07:12:48 +00:00
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>
This commit is contained in:
1 parent
40832901c5
commit
5487c080be
11 files changed
+793
-6
No files matched your search
@@ -70,7 +70,7 @@ With Object Storage or a Database Table
|
||||
|
||||
Use ``data_sources`` with
|
||||
:class:`~airflow.providers.common.sql.config.DataSourceConfig` to include
|
||||
object-storage sources (S3 or GCS Parquet, CSV, Iceberg, etc.) in the comparison.
|
||||
object-storage sources (S3, GCS, Azure Blob Storage, Parquet, CSV, Iceberg, etc.) in the comparison.
|
||||
These can be freely combined with ``db_conn_ids``. Whether an entry is
|
||||
introspected via ``DbApiHook`` or DataFusion depends on what its ``conn_id``
|
||||
resolves to, not on its ``uri``/``format`` fields: a ``DataSourceConfig``
|
||||
|
||||
@@ -65,7 +65,7 @@ With Object Storage
|
||||
-------------------
|
||||
|
||||
Use ``datasource_config`` to generate queries for data stored in object storage
|
||||
(e.g., S3, GCS, local filesystem) via `DataFusion <https://datafusion.apache.org/>`_.
|
||||
(e.g., S3, GCS, Azure Blob Storage, local filesystem) via `DataFusion <https://datafusion.apache.org/>`_.
|
||||
The operator uses :class:`~airflow.providers.common.sql.config.DataSourceConfig`
|
||||
to register the object storage source as a table so the LLM can include it in
|
||||
the schema context.
|
||||
|
||||
@@ -26,7 +26,7 @@ Files with DataFusion: ``DataFusionToolset``
|
||||
Curated toolset wrapping
|
||||
:class:`~airflow.providers.common.sql.datafusion.engine.DataFusionEngine`
|
||||
with three tools (``list_tables``, ``get_schema``, and ``query``) for
|
||||
querying files on object stores (S3, GCS, local filesystem, Iceberg) via Apache DataFusion.
|
||||
querying files on object stores (S3, GCS, Azure Blob Storage, local filesystem, Iceberg) via Apache DataFusion.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
|
||||
@@ -297,10 +297,11 @@ Supported Storage Systems
|
||||
-------------------------
|
||||
- S3
|
||||
- GCS
|
||||
- Azure Blob Storage
|
||||
- Local File System
|
||||
|
||||
.. note::
|
||||
Azure, HTTP, Delta are not yet supported but will be added in the future.
|
||||
HTTP, Delta are not yet supported but will be added in the future.
|
||||
|
||||
|
||||
|
||||
@@ -362,6 +363,63 @@ 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. The account name
|
||||
comes from ``host`` (its first DNS label) when set, falling back to
|
||||
``login`` only when ``host`` is empty; 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
|
||||
hostname. A ``host:port`` address (the Azurite emulator's shape, e.g.
|
||||
``azurite:10000``) is never a real DNS hostname and always raises instead,
|
||||
even with those variables set: put the account name in ``login`` with
|
||||
``host`` left empty for Azurite, alongside those same variables and
|
||||
``AZURE_ALLOW_HTTP=true``. ``client_secret_auth_config`` (the authority
|
||||
override ``WasbHook`` honors) is not read here.
|
||||
|
||||
The connection supplies one of the following credentials, checked in this
|
||||
order (matching ``WasbHook.get_conn``):
|
||||
|
||||
1. Azure AD service principal -- ``tenant_id`` extra, with ``login`` as the
|
||||
client ID and ``password`` as the client secret (both required together)
|
||||
2. Shared key -- the ``shared_access_key`` extra
|
||||
3. SAS token -- ``sas_token`` extra, as a query string
|
||||
4. Shared key -- ``password``, or the ``account_key`` extra
|
||||
5. None of the above -- ambient auth (see below)
|
||||
|
||||
**A worker environment variable can override the connection, but only from a
|
||||
higher-priority tier.** DataFusion reads ``AZURE_*`` environment variables
|
||||
first, and resolves credentials in this priority order:
|
||||
|
||||
1. Bearer token
|
||||
2. Access key -- the shared-key connection's tier
|
||||
3. Workload identity
|
||||
4. Client secret -- the service-principal connection's tier
|
||||
5. SAS -- the SAS connection's tier
|
||||
|
||||
Only a complete credential from an earlier tier can override the
|
||||
connection -- a single bearer or access-key variable, or every variable a
|
||||
multi-field tier needs. This raises when found, naming the variables,
|
||||
instead of silently using the wrong identity.
|
||||
|
||||
With no explicit credential, authentication falls back to ``AZURE_*``
|
||||
environment variables, managed identity, or workload identity. Unlike
|
||||
``WasbHook`` (which tries ``az login`` automatically), DataFusion's
|
||||
underlying ``object_store`` binding only tries the Azure CLI if
|
||||
``AZURE_USE_AZURE_CLI=true`` is set; otherwise it defaults straight to
|
||||
IMDS managed identity.
|
||||
|
||||
``connection_string``, ``managed_identity_client_id``, ``workload_identity_tenant_id``,
|
||||
and a URL-form ``sas_token`` are not supported.
|
||||
|
||||
.. exampleinclude:: /../../sql/src/airflow/providers/common/sql/example_dags/example_analytics.py
|
||||
:language: python
|
||||
:dedent: 4
|
||||
:start-after: [START howto_analytics_operator_with_azure]
|
||||
:end-before: [END howto_analytics_operator_with_azure]
|
||||
|
||||
Local File System Storage
|
||||
-------------------------
|
||||
.. exampleinclude:: /../../sql/src/airflow/providers/common/sql/example_dags/example_analytics.py
|
||||
|
||||
@@ -48,6 +48,7 @@ class StorageType(str, Enum):
|
||||
|
||||
S3 = "s3"
|
||||
GCS = "gcs"
|
||||
AZURE = "azure"
|
||||
LOCAL = "local"
|
||||
|
||||
|
||||
@@ -117,6 +118,8 @@ class DataSourceConfig:
|
||||
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
|
||||
raise ValueError(f"Unsupported storage type for URI: {self.uri}")
|
||||
@@ -18,6 +18,7 @@ from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from urllib.parse import parse_qsl, urlsplit
|
||||
|
||||
from datafusion import SessionContext
|
||||
|
||||
@@ -161,6 +162,12 @@ class DataFusionEngine(LoggingMixin):
|
||||
return extra_dejson[field_name]
|
||||
return extra_dejson.get(f"extra__google_cloud_platform__{field_name}")
|
||||
|
||||
def _get_wasb_extra_field(extra_dejson: dict[str, Any], field_name: str) -> Any:
|
||||
# Older Airflow connection UIs wrote custom extra fields as
|
||||
# extra__wasb__<field_name> instead of the bare key; WasbHook still reads that
|
||||
# legacy spelling as a fallback, so this must too.
|
||||
return extra_dejson.get(field_name, extra_dejson.get(f"extra__wasb__{field_name}"))
|
||||
|
||||
match conn.conn_type:
|
||||
case "aws":
|
||||
try:
|
||||
@@ -204,6 +211,66 @@ class DataFusionEngine(LoggingMixin):
|
||||
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 _get_wasb_extra_field(extra_dejson, 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)}
|
||||
credential_tier: str | None = None
|
||||
if tenant_id := _get_wasb_extra_field(extra_dejson, "tenant_id"):
|
||||
if not conn.login or not conn.password:
|
||||
# Falling through here would silently switch identity (ambient auth, or
|
||||
# the client secret sent as a shared key) instead of failing clearly.
|
||||
missing = "login (client_id)" if not conn.login else "password (client_secret)"
|
||||
raise ValueError(
|
||||
f"Connection extra 'tenant_id' is set for DataFusion Azure Blob Storage "
|
||||
f"service-principal auth, but {missing} is not."
|
||||
)
|
||||
credentials.update(
|
||||
{"client_id": conn.login, "client_secret": conn.password, "tenant_id": tenant_id}
|
||||
)
|
||||
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.
|
||||
credentials["access_key"] = shared_access_key
|
||||
credential_tier = "access_key"
|
||||
elif sas_token := _get_wasb_extra_field(extra_dejson, "sas_token"):
|
||||
if sas_token.startswith("http"):
|
||||
raise ValueError(
|
||||
"A URL-form `sas_token` is not supported for DataFusion Azure Blob Storage "
|
||||
"access; provide the SAS token as a query string instead."
|
||||
)
|
||||
credentials["sas_query_pairs"] = parse_qsl(sas_token.lstrip("?"))
|
||||
credential_tier = "sas"
|
||||
else:
|
||||
access_key = conn.password or _get_wasb_extra_field(extra_dejson, "account_key")
|
||||
if access_key:
|
||||
credentials["access_key"] = access_key
|
||||
credential_tier = "access_key"
|
||||
|
||||
if credential_tier is not None:
|
||||
conflicting_env_vars = self._find_conflicting_azure_env_vars(credential_tier)
|
||||
if conflicting_env_vars:
|
||||
raise ValueError(
|
||||
f"Worker environment variable(s) {', '.join(conflicting_env_vars)} would "
|
||||
"silently take precedence over this connection's explicit credential in "
|
||||
"DataFusion's Azure Blob Storage binding. Unset them on the worker, or "
|
||||
"remove the explicit credential from this connection to rely on the "
|
||||
"environment instead."
|
||||
)
|
||||
credentials = self._remove_none_values(credentials)
|
||||
|
||||
case _:
|
||||
raise ValueError(f"Unknown connection type {conn.conn_type}")
|
||||
return credentials, extra_config
|
||||
@@ -213,6 +280,124 @@ class DataFusionEngine(LoggingMixin):
|
||||
"""Filter out None values from the dictionary."""
|
||||
return {k: v for k, v in params.items() if v is not None}
|
||||
|
||||
_AZURE_PUBLIC_SUFFIX = ".blob.core.windows.net"
|
||||
|
||||
# object_store's build() precedence, high to low: bearer token > access key > workload
|
||||
# identity (client_id+tenant_id+federated_token_file) > client secret
|
||||
# (client_id+client_secret+tenant_id) > SAS. Used by _find_conflicting_azure_env_vars.
|
||||
_AZURE_ENV_BEARER_VARS = ("AZURE_STORAGE_TOKEN",)
|
||||
_AZURE_ENV_ACCESS_KEY_VARS = (
|
||||
"AZURE_STORAGE_ACCOUNT_KEY",
|
||||
"AZURE_STORAGE_ACCESS_KEY",
|
||||
"AZURE_STORAGE_MASTER_KEY",
|
||||
)
|
||||
_AZURE_ENV_CLIENT_ID_VARS = ("AZURE_STORAGE_CLIENT_ID", "AZURE_CLIENT_ID")
|
||||
_AZURE_ENV_CLIENT_SECRET_VARS = ("AZURE_STORAGE_CLIENT_SECRET", "AZURE_CLIENT_SECRET")
|
||||
_AZURE_ENV_TENANT_ID_VARS = (
|
||||
"AZURE_STORAGE_TENANT_ID",
|
||||
"AZURE_STORAGE_AUTHORITY_ID",
|
||||
"AZURE_TENANT_ID",
|
||||
"AZURE_AUTHORITY_ID",
|
||||
)
|
||||
_AZURE_ENV_FEDERATED_TOKEN_FILE_VAR = "AZURE_FEDERATED_TOKEN_FILE"
|
||||
|
||||
@classmethod
|
||||
def _find_conflicting_azure_env_vars(cls, credential_tier: str) -> list[str]:
|
||||
"""
|
||||
Return worker env vars that would silently outrank the connection's own credential.
|
||||
|
||||
The binding always calls `from_env()` with no way to skip it, so this can only be
|
||||
caught here, not avoided. The connection's fields overwrite the same fields from
|
||||
`from_env()`, so only a tier above the connection's own, or an unoccupied tier
|
||||
(workload identity or client secret, for SAS) with *all* its fields present, can
|
||||
actually take over.
|
||||
"""
|
||||
|
||||
def env_set(*var_groups: tuple[str, ...]) -> list[str]:
|
||||
return [var for group in var_groups for var in group if os.environ.get(var)]
|
||||
|
||||
conflicting = env_set(cls._AZURE_ENV_BEARER_VARS)
|
||||
if credential_tier == "access_key":
|
||||
return conflicting
|
||||
conflicting += env_set(cls._AZURE_ENV_ACCESS_KEY_VARS)
|
||||
if credential_tier == "client_secret":
|
||||
# client_id and tenant_id are already the connection's own; only the federated
|
||||
# token file is left for env to complete the workload-identity triple with.
|
||||
if os.environ.get(cls._AZURE_ENV_FEDERATED_TOKEN_FILE_VAR):
|
||||
conflicting.append(cls._AZURE_ENV_FEDERATED_TOKEN_FILE_VAR)
|
||||
return conflicting
|
||||
# SAS occupies none of these fields, so each whole triple must come from env.
|
||||
if (
|
||||
any(os.environ.get(var) for var in cls._AZURE_ENV_CLIENT_ID_VARS)
|
||||
and any(os.environ.get(var) for var in cls._AZURE_ENV_TENANT_ID_VARS)
|
||||
and os.environ.get(cls._AZURE_ENV_FEDERATED_TOKEN_FILE_VAR)
|
||||
):
|
||||
conflicting += env_set(cls._AZURE_ENV_CLIENT_ID_VARS, cls._AZURE_ENV_TENANT_ID_VARS)
|
||||
conflicting.append(cls._AZURE_ENV_FEDERATED_TOKEN_FILE_VAR)
|
||||
if (
|
||||
any(os.environ.get(var) for var in cls._AZURE_ENV_CLIENT_ID_VARS)
|
||||
and any(os.environ.get(var) for var in cls._AZURE_ENV_CLIENT_SECRET_VARS)
|
||||
and any(os.environ.get(var) for var in cls._AZURE_ENV_TENANT_ID_VARS)
|
||||
):
|
||||
conflicting += env_set(
|
||||
cls._AZURE_ENV_CLIENT_ID_VARS,
|
||||
cls._AZURE_ENV_CLIENT_SECRET_VARS,
|
||||
cls._AZURE_ENV_TENANT_ID_VARS,
|
||||
)
|
||||
return list(dict.fromkeys(conflicting))
|
||||
|
||||
@classmethod
|
||||
def _resolve_wasb_account(cls, host: str | None, login: str | None) -> str | None:
|
||||
"""
|
||||
Return the storage account name the way WasbHook resolves it.
|
||||
|
||||
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
|
||||
``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
|
||||
the microsoft-azure provider's Azure SDK dependency for one stdlib string operation.
|
||||
|
||||
Only the public ``*.blob.core.windows.net`` cloud is supported, unless
|
||||
``AZURE_STORAGE_ENDPOINT``/``AZURE_ENDPOINT`` is set for a real sovereign-cloud hostname
|
||||
(DataFusion's binding otherwise has no endpoint override). A ``host:port`` netloc (the
|
||||
Azurite emulator's shape) is never a real hostname, so it always raises instead --
|
||||
checked before the dotless-host fallback below, which would otherwise mask it.
|
||||
"""
|
||||
if not host and not login:
|
||||
return None
|
||||
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 host and ":" in netloc:
|
||||
raise ValueError(
|
||||
f"Connection host {host!r} looks like an emulator address (host:port), which "
|
||||
"DataFusion's Azure Blob Storage binding cannot resolve an account name from -- "
|
||||
"even with AZURE_STORAGE_ENDPOINT set. Put the account name in `login` with "
|
||||
"`host` empty instead, alongside AZURE_STORAGE_ENDPOINT and AZURE_ALLOW_HTTP=true."
|
||||
)
|
||||
if "." not in netloc:
|
||||
if not login:
|
||||
raise ValueError(
|
||||
f"Connection host {host!r} is not a full URL or DNS name, and no `login` was "
|
||||
"given to resolve it as an Active Directory ID instead."
|
||||
)
|
||||
# Only an Active Directory ID was given, not a full URL or DNS name.
|
||||
netloc = f"{login}.blob.core.windows.net"
|
||||
if not netloc.endswith(cls._AZURE_PUBLIC_SUFFIX) and not (
|
||||
os.environ.get("AZURE_STORAGE_ENDPOINT") or os.environ.get("AZURE_ENDPOINT")
|
||||
):
|
||||
raise ValueError(
|
||||
f"Connection host {host!r} does not resolve to the public {cls._AZURE_PUBLIC_SUFFIX} "
|
||||
"cloud, which is the only one DataFusion's Azure Blob Storage binding can target (it "
|
||||
"has no endpoint override). A sovereign cloud is supported once the "
|
||||
"AZURE_STORAGE_ENDPOINT environment variable is set."
|
||||
)
|
||||
# Azure storage account names are capped at 24 characters.
|
||||
return netloc.split(".", 1)[0][:24]
|
||||
|
||||
def get_schema(self, table_name: str):
|
||||
"""Get the schema of a table."""
|
||||
schema = str(self.session_context.table(table_name).schema())
|
||||
|
||||
+34
-2
@@ -20,7 +20,7 @@ import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from datafusion.object_store import AmazonS3, GoogleCloud, LocalFileSystem
|
||||
from datafusion.object_store import AmazonS3, GoogleCloud, LocalFileSystem, MicrosoftAzure
|
||||
|
||||
from airflow.providers.common.sql.config import ConnectionConfig, StorageType
|
||||
from airflow.providers.common.sql.datafusion.base import ObjectStorageProvider
|
||||
@@ -107,6 +107,37 @@ class GCSObjectStorageProvider(ObjectStorageProvider):
|
||||
return "gs://"
|
||||
|
||||
|
||||
class AzureObjectStorageProvider(ObjectStorageProvider):
|
||||
"""Azure Object Storage Provider using DataFusion's MicrosoftAzure."""
|
||||
|
||||
@property
|
||||
def get_storage_type(self) -> StorageType:
|
||||
"""Return the storage type."""
|
||||
return StorageType.AZURE
|
||||
|
||||
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:
|
||||
raise ValueError(f"connection_config must be provided for {self.get_storage_type.value}")
|
||||
|
||||
try:
|
||||
credentials = connection_config.credentials
|
||||
container = self.get_bucket(path)
|
||||
|
||||
azure_store = MicrosoftAzure(container_name=container, **credentials)
|
||||
self.log.info("Created Azure object store for container %s", container)
|
||||
|
||||
return azure_store
|
||||
|
||||
except BaseException as e:
|
||||
# A bad credential combination panics as pyo3_runtime.PanicException, not a plain Exception.
|
||||
raise ObjectStoreCreationException(f"Failed to create Azure object store: {e}")
|
||||
|
||||
def get_scheme(self) -> str:
|
||||
"""Return the scheme for Azure."""
|
||||
return "az://"
|
||||
|
||||
|
||||
class LocalObjectStorageProvider(ObjectStorageProvider):
|
||||
"""Local Object Storage Provider using DataFusion's LocalFileSystem."""
|
||||
|
||||
@@ -126,10 +157,11 @@ class LocalObjectStorageProvider(ObjectStorageProvider):
|
||||
|
||||
def get_object_storage_provider(storage_type: StorageType) -> ObjectStorageProvider:
|
||||
"""Get an object storage provider based on the storage type."""
|
||||
# TODO: Add support for Azure, HTTP: https://datafusion.apache.org/python/autoapi/datafusion/object_store/index.html
|
||||
# TODO: Add support for HTTP: https://datafusion.apache.org/python/autoapi/datafusion/object_store/index.html
|
||||
providers: dict[StorageType, type] = {
|
||||
StorageType.S3: S3ObjectStorageProvider,
|
||||
StorageType.GCS: GCSObjectStorageProvider,
|
||||
StorageType.AZURE: AzureObjectStorageProvider,
|
||||
StorageType.LOCAL: LocalObjectStorageProvider,
|
||||
}
|
||||
|
||||
|
||||
+13
@@ -34,6 +34,10 @@ datasource_config_gcs = DataSourceConfig(
|
||||
conn_id="google_cloud_default", table_name="users_data", uri="gs://bucket/path/", format="parquet"
|
||||
)
|
||||
|
||||
datasource_config_azure = DataSourceConfig(
|
||||
conn_id="wasb_default", table_name="users_data", uri="az://container/path/", format="parquet"
|
||||
)
|
||||
|
||||
datasource_config_iceberg = DataSourceConfig(
|
||||
conn_id="iceberg_default",
|
||||
table_name="users_data",
|
||||
@@ -85,6 +89,15 @@ with DAG(
|
||||
analytics_with_s3 >> analytics_with_gcs
|
||||
# [END howto_analytics_operator_with_gcs]
|
||||
|
||||
# [START howto_analytics_operator_with_azure]
|
||||
analytics_with_azure = AnalyticsOperator(
|
||||
task_id="analytics_with_azure",
|
||||
datasource_configs=[datasource_config_azure],
|
||||
queries=["SELECT * FROM users_data", "SELECT count(*) FROM users_data"],
|
||||
)
|
||||
analytics_with_s3 >> analytics_with_azure
|
||||
# [END howto_analytics_operator_with_azure]
|
||||
|
||||
# [START howto_analytics_operator_with_local]
|
||||
analytics_with_local = AnalyticsOperator(
|
||||
task_id="analytics_with_local",
|
||||
|
||||
@@ -85,6 +85,7 @@ class TestDataFusionEngine:
|
||||
("s3", "csv", "s3"),
|
||||
("s3", "avro", "s3"),
|
||||
("gcs", "parquet", "gs"),
|
||||
("azure", "parquet", "az"),
|
||||
],
|
||||
)
|
||||
@patch("airflow.providers.common.sql.datafusion.engine.get_object_storage_provider", autospec=True)
|
||||
@@ -394,6 +395,452 @@ class TestDataFusionEngine:
|
||||
with pytest.raises(ValueError, match="'impersonation_chain' is not supported"):
|
||||
engine._get_credentials(mock_conn)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("password", "extra_dejson", "expected_access_key"),
|
||||
[
|
||||
("mykey", {}, "mykey"),
|
||||
(None, {"shared_access_key": "extra-key"}, "extra-key"),
|
||||
(None, {"account_key": "extra-key"}, "extra-key"),
|
||||
],
|
||||
ids=["password", "shared_access_key_extra", "account_key_extra"],
|
||||
)
|
||||
def test_get_credentials_azure_with_shared_key(self, password, extra_dejson, expected_access_key):
|
||||
mock_conn = MagicMock()
|
||||
mock_conn.conn_type = "wasb"
|
||||
mock_conn.host = None
|
||||
mock_conn.login = "myaccount"
|
||||
mock_conn.password = password
|
||||
mock_conn.extra_dejson = extra_dejson
|
||||
engine = DataFusionEngine()
|
||||
|
||||
credentials, extra_config = engine._get_credentials(mock_conn)
|
||||
|
||||
assert credentials == {"account": "myaccount", "access_key": expected_access_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 == {}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("login", "password", "missing"),
|
||||
[
|
||||
(None, "client-secret", "login"),
|
||||
("client-id", None, "password"),
|
||||
(None, None, "login"),
|
||||
],
|
||||
)
|
||||
def test_get_credentials_azure_partial_service_principal_raises(self, login, password, missing):
|
||||
"""A partial service-principal config must raise, not silently authenticate with a
|
||||
different identity (ambient auth, or the client secret sent as a shared key) --
|
||||
DataFusion's binding also 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 = login
|
||||
mock_conn.password = password
|
||||
mock_conn.extra_dejson = {"tenant_id": "tenant-id"}
|
||||
engine = DataFusionEngine()
|
||||
|
||||
with pytest.raises(ValueError, match=f"{missing}.*is not"):
|
||||
engine._get_credentials(mock_conn)
|
||||
|
||||
def test_get_credentials_azure_fully_empty_connection_omits_account(self):
|
||||
"""Neither host nor login set (the shape of the ``wasb_default`` connection ``airflow
|
||||
db`` creates) must drop `account` entirely, not send the literal string 'None' --
|
||||
the binding then falls back to AZURE_STORAGE_ACCOUNT_NAME."""
|
||||
mock_conn = MagicMock()
|
||||
mock_conn.conn_type = "wasb"
|
||||
mock_conn.host = None
|
||||
mock_conn.login = None
|
||||
mock_conn.password = None
|
||||
mock_conn.extra_dejson = {}
|
||||
engine = DataFusionEngine()
|
||||
|
||||
credentials, extra_config = engine._get_credentials(mock_conn)
|
||||
|
||||
assert credentials == {}
|
||||
assert extra_config == {}
|
||||
|
||||
def test_get_credentials_azure_shared_access_key_takes_priority_over_sas_token(self):
|
||||
"""Matches WasbHook.get_conn, which checks the `shared_access_key` extra before
|
||||
`sas_token`."""
|
||||
mock_conn = MagicMock()
|
||||
mock_conn.conn_type = "wasb"
|
||||
mock_conn.host = None
|
||||
mock_conn.login = "myaccount"
|
||||
mock_conn.password = None
|
||||
mock_conn.extra_dejson = {
|
||||
"sas_token": "?sv=2020-08-04&sp=rl&sig=abc",
|
||||
"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_sas_token(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 = {"sas_token": "?sv=2020-08-04&sp=rl&sig=abc"}
|
||||
engine = DataFusionEngine()
|
||||
|
||||
credentials, extra_config = engine._get_credentials(mock_conn)
|
||||
|
||||
assert credentials == {
|
||||
"account": "myaccount",
|
||||
"sas_query_pairs": [("sv", "2020-08-04"), ("sp", "rl"), ("sig", "abc")],
|
||||
}
|
||||
assert extra_config == {}
|
||||
|
||||
def test_get_credentials_azure_without_credentials_uses_ambient_auth(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 = {}
|
||||
engine = DataFusionEngine()
|
||||
|
||||
credentials, extra_config = engine._get_credentials(mock_conn)
|
||||
|
||||
assert credentials == {"account": "myaccount"}
|
||||
assert extra_config == {}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"unsupported_field",
|
||||
["connection_string", "managed_identity_client_id", "workload_identity_tenant_id"],
|
||||
)
|
||||
def test_get_credentials_azure_rejects_unsupported_identity_fields(self, unsupported_field):
|
||||
mock_conn = MagicMock()
|
||||
mock_conn.conn_type = "wasb"
|
||||
mock_conn.host = None
|
||||
mock_conn.extra_dejson = {unsupported_field: "some-value"}
|
||||
engine = DataFusionEngine()
|
||||
|
||||
with pytest.raises(ValueError, match=f"{unsupported_field!r} is not supported"):
|
||||
engine._get_credentials(mock_conn)
|
||||
|
||||
def test_get_credentials_azure_reads_legacy_extra_prefixed_sas_token(self):
|
||||
"""Older Airflow connection UIs wrote custom extra fields as
|
||||
extra__wasb__<field>; WasbHook still reads that spelling as a fallback."""
|
||||
mock_conn = MagicMock()
|
||||
mock_conn.conn_type = "wasb"
|
||||
mock_conn.host = None
|
||||
mock_conn.login = "myaccount"
|
||||
mock_conn.password = None
|
||||
mock_conn.extra_dejson = {"extra__wasb__sas_token": "?sv=2020-08-04&sp=rl&sig=abc"}
|
||||
engine = DataFusionEngine()
|
||||
|
||||
credentials, extra_config = engine._get_credentials(mock_conn)
|
||||
|
||||
assert credentials == {
|
||||
"account": "myaccount",
|
||||
"sas_query_pairs": [("sv", "2020-08-04"), ("sp", "rl"), ("sig", "abc")],
|
||||
}
|
||||
assert extra_config == {}
|
||||
|
||||
def test_get_credentials_azure_rejects_legacy_extra_prefixed_unsupported_field(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 = {"extra__wasb__connection_string": "some-conn-string"}
|
||||
engine = DataFusionEngine()
|
||||
|
||||
with pytest.raises(ValueError, match="'connection_string' is not supported"):
|
||||
engine._get_credentials(mock_conn)
|
||||
|
||||
def test_get_credentials_azure_rejects_url_form_sas_token(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 = {"sas_token": "https://myaccount.blob.core.windows.net/?sv=2020-08-04"}
|
||||
engine = DataFusionEngine()
|
||||
|
||||
with pytest.raises(ValueError, match="URL-form `sas_token` is not supported"):
|
||||
engine._get_credentials(mock_conn)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("env_vars", "should_raise"),
|
||||
[
|
||||
(["AZURE_STORAGE_ACCOUNT_KEY"], True),
|
||||
(["AZURE_STORAGE_ACCESS_KEY"], True),
|
||||
(["AZURE_STORAGE_MASTER_KEY"], True),
|
||||
(["AZURE_STORAGE_TOKEN"], True),
|
||||
(["AZURE_STORAGE_SAS_KEY"], False),
|
||||
(["AZURE_CLIENT_ID", "AZURE_TENANT_ID", "AZURE_FEDERATED_TOKEN_FILE"], True),
|
||||
(["AZURE_FEDERATED_TOKEN_FILE"], False),
|
||||
(["AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET", "AZURE_TENANT_ID"], True),
|
||||
(["AZURE_STORAGE_CLIENT_ID", "AZURE_STORAGE_CLIENT_SECRET", "AZURE_STORAGE_TENANT_ID"], True),
|
||||
(["AZURE_STORAGE_CLIENT_ID", "AZURE_STORAGE_CLIENT_SECRET", "AZURE_STORAGE_AUTHORITY_ID"], True),
|
||||
(["AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET", "AZURE_AUTHORITY_ID"], True),
|
||||
(["AZURE_CLIENT_ID"], False),
|
||||
],
|
||||
ids=[
|
||||
"account-key",
|
||||
"access-key",
|
||||
"master-key",
|
||||
"bearer",
|
||||
"sas-key-checked-after-query-pairs",
|
||||
"workload-identity-triple",
|
||||
"workload-identity-partial",
|
||||
"client-secret-triple-bare",
|
||||
"client-secret-triple-storage-prefixed",
|
||||
"client-secret-triple-storage-authority",
|
||||
"client-secret-triple-bare-authority",
|
||||
"client-secret-partial",
|
||||
],
|
||||
)
|
||||
def test_get_credentials_azure_sas_env_precedence(self, env_vars, should_raise, monkeypatch):
|
||||
"""A SAS connection sits at the bottom of object_store's precedence order and occupies
|
||||
none of the higher tiers' fields, so a bearer token or any access-key spelling always
|
||||
outranks it, and a complete workload-identity or client-secret triple (any recognized
|
||||
spelling) does too -- but a partial triple, or AZURE_STORAGE_SAS_KEY (checked only after
|
||||
the connection's own sas_query_pairs within SAS resolution itself), must not raise."""
|
||||
for var in env_vars:
|
||||
monkeypatch.setenv(var, "some-value")
|
||||
mock_conn = MagicMock()
|
||||
mock_conn.conn_type = "wasb"
|
||||
mock_conn.host = None
|
||||
mock_conn.login = "myaccount"
|
||||
mock_conn.password = None
|
||||
mock_conn.extra_dejson = {"sas_token": "?sv=2020-08-04&sp=rl&sig=abc"}
|
||||
engine = DataFusionEngine()
|
||||
|
||||
if should_raise:
|
||||
with pytest.raises(ValueError, match=", ".join(env_vars)):
|
||||
engine._get_credentials(mock_conn)
|
||||
else:
|
||||
credentials, extra_config = engine._get_credentials(mock_conn)
|
||||
assert credentials == {
|
||||
"account": "myaccount",
|
||||
"sas_query_pairs": [("sv", "2020-08-04"), ("sp", "rl"), ("sig", "abc")],
|
||||
}
|
||||
assert extra_config == {}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("env_vars", "should_raise"),
|
||||
[
|
||||
(["AZURE_STORAGE_TOKEN"], True),
|
||||
(["AZURE_STORAGE_ACCOUNT_KEY"], False),
|
||||
(["AZURE_STORAGE_ACCESS_KEY"], False),
|
||||
(["AZURE_STORAGE_MASTER_KEY"], False),
|
||||
(["AZURE_FEDERATED_TOKEN_FILE"], False),
|
||||
(["AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET", "AZURE_TENANT_ID"], False),
|
||||
],
|
||||
ids=[
|
||||
"bearer",
|
||||
"account-key",
|
||||
"access-key",
|
||||
"master-key",
|
||||
"federated-token-file",
|
||||
"client-secret-triple",
|
||||
],
|
||||
)
|
||||
def test_get_credentials_azure_access_key_env_precedence(self, env_vars, should_raise, monkeypatch):
|
||||
"""A shared-key connection overwrites the same field object_store's build() would
|
||||
otherwise read from env, and build() picks access_key ahead of every lower tier -- so
|
||||
only a bearer token, the one tier above access_key, can actually outrank it, even on an
|
||||
AKS pod where the workload-identity webhook injects AZURE_FEDERATED_TOKEN_FILE into
|
||||
every labelled pod."""
|
||||
for var in env_vars:
|
||||
monkeypatch.setenv(var, "some-value")
|
||||
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()
|
||||
|
||||
if should_raise:
|
||||
with pytest.raises(ValueError, match=", ".join(env_vars)):
|
||||
engine._get_credentials(mock_conn)
|
||||
else:
|
||||
credentials, extra_config = engine._get_credentials(mock_conn)
|
||||
assert credentials == {"account": "myaccount", "access_key": "mykey"}
|
||||
assert extra_config == {}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("env_vars", "should_raise"),
|
||||
[
|
||||
(["AZURE_STORAGE_TOKEN"], True),
|
||||
(["AZURE_STORAGE_ACCOUNT_KEY"], True),
|
||||
(["AZURE_STORAGE_ACCESS_KEY"], True),
|
||||
(["AZURE_STORAGE_MASTER_KEY"], True),
|
||||
(["AZURE_FEDERATED_TOKEN_FILE"], True),
|
||||
(["AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET", "AZURE_TENANT_ID"], False),
|
||||
],
|
||||
ids=[
|
||||
"bearer",
|
||||
"account-key",
|
||||
"access-key",
|
||||
"master-key",
|
||||
"federated-token-file",
|
||||
"client-secret-triple",
|
||||
],
|
||||
)
|
||||
def test_get_credentials_azure_service_principal_env_precedence(
|
||||
self, env_vars, should_raise, monkeypatch
|
||||
):
|
||||
"""A bearer token or any access-key spelling sits above client secret in object_store's
|
||||
precedence order, and the connection's own client_id and tenant_id already complete two
|
||||
of workload identity's three fields, so a federated token file alone is enough for env to
|
||||
complete that tier -- but a client-secret triple in env is moot, since the connection's
|
||||
own client_id/client_secret/tenant_id overwrite the same fields."""
|
||||
for var in env_vars:
|
||||
monkeypatch.setenv(var, "some-value")
|
||||
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()
|
||||
|
||||
if should_raise:
|
||||
with pytest.raises(ValueError, match=", ".join(env_vars)):
|
||||
engine._get_credentials(mock_conn)
|
||||
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",
|
||||
}
|
||||
assert extra_config == {}
|
||||
|
||||
def test_get_credentials_azure_env_precedence_guard_ignores_pure_ambient_auth(self, monkeypatch):
|
||||
"""The guard only fires for an explicit connection credential -- a connection with none
|
||||
at all is meant to rely on the environment, so the same env vars must not raise here."""
|
||||
monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_KEY", "some-value")
|
||||
mock_conn = MagicMock()
|
||||
mock_conn.conn_type = "wasb"
|
||||
mock_conn.host = None
|
||||
mock_conn.login = "myaccount"
|
||||
mock_conn.password = None
|
||||
mock_conn.extra_dejson = {}
|
||||
engine = DataFusionEngine()
|
||||
|
||||
credentials, extra_config = engine._get_credentials(mock_conn)
|
||||
|
||||
assert credentials == {"account": "myaccount"}
|
||||
assert extra_config == {}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("host", "login", "expected"),
|
||||
[
|
||||
("myaccount.blob.core.windows.net", None, "myaccount"),
|
||||
("https://myaccount.blob.core.windows.net/container/path", None, "myaccount"),
|
||||
(None, "myaccount", "myaccount"),
|
||||
("justanadid", "myaccount", "myaccount"),
|
||||
("a" * 40 + ".blob.core.windows.net", "x", "a" * 24),
|
||||
(None, None, None),
|
||||
],
|
||||
)
|
||||
def test_resolve_wasb_account(self, host, login, expected):
|
||||
assert DataFusionEngine._resolve_wasb_account(host, login) == expected
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"host",
|
||||
[
|
||||
"myaccount.blob.core.chinacloudapi.cn",
|
||||
"myaccount.blob.core.usgovcloudapi.net",
|
||||
],
|
||||
)
|
||||
def test_resolve_wasb_account_rejects_non_public_cloud_hosts(self, host):
|
||||
"""The binding has no endpoint override, so a sovereign-cloud host would otherwise be
|
||||
silently misrouted to the public *.blob.core.windows.net account of the same name."""
|
||||
with pytest.raises(ValueError, match="does not resolve to the public"):
|
||||
DataFusionEngine._resolve_wasb_account(host, None)
|
||||
|
||||
@pytest.mark.parametrize("env_var", ["AZURE_STORAGE_ENDPOINT", "AZURE_ENDPOINT"])
|
||||
def test_resolve_wasb_account_allows_non_public_cloud_host_when_endpoint_set(self, env_var, monkeypatch):
|
||||
"""Once the worker sets an endpoint override, object_store uses it verbatim instead of
|
||||
deriving the URL from the account name, so a real (dotted, portless) sovereign-cloud
|
||||
hostname no longer needs to match *.blob.core.windows.net."""
|
||||
monkeypatch.setenv(env_var, "https://myaccount.blob.core.chinacloudapi.cn")
|
||||
|
||||
account = DataFusionEngine._resolve_wasb_account("myaccount.blob.core.chinacloudapi.cn", None)
|
||||
|
||||
assert account == "myaccount"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("host", "login"),
|
||||
[
|
||||
("http://127.0.0.1:10000/devstoreaccount1", None),
|
||||
("azurite:10000", None),
|
||||
("azurite:10000", "devstoreaccount1"),
|
||||
],
|
||||
ids=["ip-port-url", "dotless-host-port", "dotless-host-port-with-login"],
|
||||
)
|
||||
def test_resolve_wasb_account_rejects_emulator_style_hosts(self, host, login, monkeypatch):
|
||||
"""A host:port netloc (Azurite's shape) is never a real DNS hostname -- resolving the
|
||||
account from it would otherwise silently produce the wrong value ("127", or "None" once
|
||||
the dotless fallback for an Active Directory ID kicks in), so this must always raise
|
||||
regardless of login or of the endpoint override, which does not help this shape either."""
|
||||
monkeypatch.setenv("AZURE_STORAGE_ENDPOINT", host.split("/")[0])
|
||||
|
||||
with pytest.raises(ValueError, match="looks like an emulator address"):
|
||||
DataFusionEngine._resolve_wasb_account(host, login)
|
||||
|
||||
def test_resolve_wasb_account_rejects_dotless_host_without_login(self):
|
||||
"""The Active Directory ID fallback assumes login holds the real account name; without
|
||||
one there's nothing to build a netloc from, and silently using the literal string "None"
|
||||
is exactly the bug this function's own empty-connection case is meant to avoid."""
|
||||
with pytest.raises(ValueError, match="no `login` was given"):
|
||||
DataFusionEngine._resolve_wasb_account("azurite", None)
|
||||
|
||||
def test_resolve_wasb_account_azurite_workaround_uses_login_with_empty_host(self):
|
||||
"""The documented Azurite workaround -- the account name in login, host left empty --
|
||||
already resolves correctly without needing the non-public-suffix check at all."""
|
||||
assert DataFusionEngine._resolve_wasb_account(None, "devstoreaccount1") == "devstoreaccount1"
|
||||
|
||||
def test_get_credentials_unknown_type(self):
|
||||
mock_conn = MagicMock()
|
||||
mock_conn.conn_type = "dummy"
|
||||
|
||||
@@ -25,6 +25,7 @@ import pytest
|
||||
from airflow.providers.common.sql.config import ConnectionConfig, StorageType
|
||||
from airflow.providers.common.sql.datafusion.exceptions import ObjectStoreCreationException
|
||||
from airflow.providers.common.sql.datafusion.object_storage_provider import (
|
||||
AzureObjectStorageProvider,
|
||||
GCSObjectStorageProvider,
|
||||
LocalObjectStorageProvider,
|
||||
S3ObjectStorageProvider,
|
||||
@@ -141,6 +142,52 @@ class TestObjectStorageProvider:
|
||||
with pytest.raises(ValueError, match="connection_config must be provided for gcs"):
|
||||
provider.create_object_store("gs://demo-data/path")
|
||||
|
||||
@patch("airflow.providers.common.sql.datafusion.object_storage_provider.MicrosoftAzure")
|
||||
def test_azure_provider_success(self, mock_azure):
|
||||
provider = AzureObjectStorageProvider()
|
||||
connection_config = ConnectionConfig(
|
||||
conn_id="wasb_default",
|
||||
credentials={"account": "myaccount", "access_key": "fake_key"},
|
||||
)
|
||||
|
||||
store = provider.create_object_store("az://demo-container/path", 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_storage_type == StorageType.AZURE
|
||||
assert provider.get_scheme() == "az://"
|
||||
|
||||
def test_azure_provider_failure(self):
|
||||
provider = AzureObjectStorageProvider()
|
||||
connection_config = ConnectionConfig(conn_id="wasb_default")
|
||||
|
||||
with patch(
|
||||
"airflow.providers.common.sql.datafusion.object_storage_provider.MicrosoftAzure",
|
||||
side_effect=Exception("Error"),
|
||||
):
|
||||
with pytest.raises(ObjectStoreCreationException, match="Failed to create Azure object store"):
|
||||
provider.create_object_store("az://demo-container/path", connection_config)
|
||||
|
||||
def test_azure_provider_requires_connection_config(self):
|
||||
provider = AzureObjectStorageProvider()
|
||||
|
||||
with pytest.raises(ValueError, match="connection_config must be provided for azure"):
|
||||
provider.create_object_store("az://demo-container/path")
|
||||
|
||||
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."""
|
||||
provider = AzureObjectStorageProvider()
|
||||
connection_config = ConnectionConfig(
|
||||
conn_id="wasb_default",
|
||||
credentials={"client_id": "only-this-one-set"},
|
||||
)
|
||||
|
||||
with pytest.raises(ObjectStoreCreationException, match="Failed to create Azure object store"):
|
||||
provider.create_object_store("az://demo-container/path", connection_config)
|
||||
|
||||
@patch("airflow.providers.common.sql.datafusion.object_storage_provider.LocalFileSystem")
|
||||
def test_local_provider(self, mock_local):
|
||||
provider = LocalObjectStorageProvider()
|
||||
@@ -152,6 +199,7 @@ class TestObjectStorageProvider:
|
||||
def test_get_object_storage_provider(self):
|
||||
assert isinstance(get_object_storage_provider(StorageType.S3), S3ObjectStorageProvider)
|
||||
assert isinstance(get_object_storage_provider(StorageType.GCS), GCSObjectStorageProvider)
|
||||
assert isinstance(get_object_storage_provider(StorageType.AZURE), AzureObjectStorageProvider)
|
||||
assert isinstance(get_object_storage_provider(StorageType.LOCAL), LocalObjectStorageProvider)
|
||||
|
||||
with pytest.raises(ValueError, match="Unsupported storage type"):
|
||||
|
||||
@@ -34,6 +34,7 @@ class TestDataSourceConfig:
|
||||
[
|
||||
("s3://bucket/path", StorageType.S3),
|
||||
("gs://bucket/path", StorageType.GCS),
|
||||
("az://container/path", StorageType.AZURE),
|
||||
("file:///path/to/file", StorageType.LOCAL),
|
||||
],
|
||||
)
|
||||
|
||||
Reference in new issue
Block a user