kaxil commented on code in PR #73374:
URL: https://github.com/apache/airflow/pull/73374#discussion_r4139361214
##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -213,6 +317,50 @@ 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}
+ _AZURE_PUBLIC_SUFFIX = ".blob.core.windows.net"
+
+ @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`` 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. Returns ``None`` when neither is set, so the binding
falls back to
+ ``AZURE_STORAGE_ACCOUNT_NAME`` instead of targeting the literal string
``"None"``.
+ 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.
+
+ Only the public ``*.blob.core.windows.net`` cloud is supported unless
the worker sets
+ ``AZURE_STORAGE_ENDPOINT``/``AZURE_ENDPOINT``: DataFusion's Azure
binding takes no
+ endpoint override on this side, so a sovereign-cloud or emulator host
would otherwise be
+ silently misrouted to the public account of the same name -- but once
one of those
+ variables is set, the binding uses it verbatim instead of deriving the
URL from the
+ account name, so the host no longer needs to match.
+ """
+ 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 "." not in netloc:
+ # 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). Sovereign clouds and the Azurite
emulator are not "
+ "supported unless the AZURE_STORAGE_ENDPOINT environment
variable is set."
Review Comment:
The endpoint skip fixes the sovereign-cloud case but not the Azurite one
this message names: the account is still the first label of the netloc, so
`http://127.0.0.1:10000/devstoreaccount1` with `AZURE_STORAGE_ENDPOINT` set
resolves to `account="127"`. I ran it against a live Azurite container: with
`AZURE_ALLOW_HTTP=true` the request path is right but the header is `SharedKey
127:...` and Azurite returns 403; with the account forced to `devstoreaccount1`
the same query returns rows; without `AZURE_ALLOW_HTTP` the store sends nothing
at all. Two ways out: narrow the promise here and in the docs (sovereign clouds
work with the endpoint variable; Azurite needs the account in `login` with
`host` empty, plus `AZURE_ALLOW_HTTP=true`), or prefer `login` over the host
label when the endpoint variable is set and both are present. Either way
`test_resolve_wasb_account_allows_non_public_cloud_host_when_endpoint_set` only
pins the sovereign shape; a row with the Azurite URL would have caught
this.
##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -140,6 +141,23 @@ def _get_connection_config(self, conn_id: str) ->
ConnectionConfig:
extra_config=extra_config,
)
+ # Used by the wasb branch's env-precedence guard below.
Review Comment:
Shape note, not blocking: the triple check below (three `any()` calls plus
the nested `extend`) collapses to one comprehension once the three `any()`
results have names, and these four tuples sit 170 lines from
`_AZURE_PUBLIC_SUFFIX`, which they belong next to. If the guard becomes per
tier as suggested on line 284, a small mapping from tier to conflict tuple
replaces both.
##########
providers/common/sql/docs/operators.rst:
##########
@@ -362,6 +363,57 @@ 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 unless the worker sets
+``AZURE_STORAGE_ENDPOINT``/``AZURE_ENDPOINT``, since DataFusion's binding
+otherwise has no endpoint override. ``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.** DataFusion
+reads ``AZURE_*`` environment variables first, and an environment bearer
+token, access key, workload-identity token, or client secret wins over the
+connection's SAS token or shared key (a SAS token is the lowest-precedence
Review Comment:
`or shared key` overclaims. A connection shared key is overlaid onto the
builder and `build()` checks `access_key` before workload identity and client
secret, so only a bearer token (`AZURE_STORAGE_TOKEN`) can outrank it; the
round-3 wording (SAS token or client secret) was the accurate one. Once the
guard is per tier, this paragraph should describe the same rule as the code.
##########
providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py:
##########
@@ -204,6 +228,86 @@ def _get_gcp_extra_field(extra_dejson: dict[str, Any],
field_name: 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 _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)}
+ explicit_credential = False
+ 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}
+ )
+ explicit_credential = True
+ 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
+ explicit_credential = True
+ 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("?"))
+ explicit_credential = True
+ else:
+ access_key = conn.password or
_get_wasb_extra_field(extra_dejson, "account_key")
+ if access_key:
+ credentials["access_key"] = access_key
+ explicit_credential = True
+
+ if explicit_credential:
+ # object_store's from_env() precedence, high to low:
bearer token > access
+ # key > workload identity > client secret > SAS. Anything
above the
+ # connection's own tier would silently win; SAS is the
bottom tier, so an
+ # env-derived SAS key never can and isn't checked.
+ conflicting_env_vars = [
Review Comment:
Round 3 added the missing names, thanks, and this is now the right set for a
SAS connection. It is the wrong set for the other two shapes, because
`MicrosoftAzure.__init__` overlays the connection's values onto the
`from_env()` builder before `build()` runs its chain (bearer > access key >
workload identity > client secret > SAS). A connection shared key replaces any
env key via `with_access_key`, and `build()` picks `access_key` ahead of
workload identity and client secret, so only `AZURE_STORAGE_TOKEN` can outrank
it. On an AKS pod with the workload-identity webhook, which injects
`AZURE_CLIENT_ID`, `AZURE_TENANT_ID` and `AZURE_FEDERATED_TOKEN_FILE` into
every labelled pod, this raises for a plain `shared_access_key` connection that
the store would have honored, and neither remedy in the message applies: the
variables are not the user's to unset, and dropping the key switches the task
to the pod identity. I checked on datafusion 50.0.0 (the floor) and 51.0.0:
with that triple se
t and a shared-key connection, construction fails with `InvalidAccessKey` for
a bad key and succeeds for a good one, so the connection key is what gets used.
A service-principal connection is similar in the other direction: the env
client-secret triple is overlaid and never wins, while
`AZURE_FEDERATED_TOKEN_FILE` does, since the connection's `client_id` and
`tenant_id` complete the workload-identity triple with the pod token.
In round 3 I called the extra single-var names at worst conservative. This
case is different because it lands on the most common credential on the most
common Azure Kubernetes deployment, and the error tells the user to do two
things that do not work. Could the conflict set be chosen per tier:
`{AZURE_STORAGE_TOKEN}` for a key connection; that plus the three access-key
spellings and `AZURE_FEDERATED_TOKEN_FILE` for a service principal; the current
full list for SAS? The comment above then just states that rule, plus the
reason the guard exists at all (the binding always calls `from_env()` and there
is no way to skip it, so this can only be caught on the Python side), which the
round-3 rewrite dropped. It also says `from_env()` precedence where the chain
is `build()`'s.
##########
providers/common/sql/tests/unit/common/sql/datafusion/test_engine.py:
##########
@@ -394,6 +395,352 @@ def
test_get_credentials_gcs_rejects_legacy_extra_prefixed_unsupported_field(sel
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_var",
+ [
+ "AZURE_FEDERATED_TOKEN_FILE",
+ "AZURE_STORAGE_ACCOUNT_KEY",
+ "AZURE_STORAGE_ACCESS_KEY",
+ "AZURE_STORAGE_MASTER_KEY",
+ "AZURE_STORAGE_TOKEN",
+ ],
+ )
+ def
test_get_credentials_azure_rejects_when_env_would_outrank_explicit_credential(
+ self, env_var, monkeypatch
+ ):
+ """DataFusion's binding always reads these AZURE_* vars via from_env()
before overlaying
+ an explicit credential, and object_store checks the env-derived slots
before the
+ connection's SAS slot -- so any of these would silently win."""
+ monkeypatch.setenv(env_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()
+
+ with pytest.raises(ValueError, match=env_var):
+ engine._get_credentials(mock_conn)
+
+ def test_get_credentials_azure_ignores_sas_key_env_var(self, monkeypatch):
+ """AZURE_STORAGE_SAS_KEY sits below every explicit credential in
object_store's
+ precedence order, so it can never outrank a connection's credential
and must not raise."""
+ monkeypatch.setenv("AZURE_STORAGE_SAS_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 = {"sas_token": "?sv=2020-08-04&sp=rl&sig=abc"}
+ engine = DataFusionEngine()
+
+ credentials, extra_config = engine._get_credentials(mock_conn)
+
+ assert "sas_query_pairs" in credentials
+ assert extra_config == {}
+
+ def
test_get_credentials_azure_rejects_when_env_client_secret_triple_would_outrank_sas(self,
monkeypatch):
+ """A full client_id+client_secret+tenant_id triple in env sits above
SAS in
+ object_store's precedence order, so it would silently win over a
connection's SAS
+ credential."""
+ monkeypatch.setenv("AZURE_CLIENT_ID", "some-client-id")
+ monkeypatch.setenv("AZURE_CLIENT_SECRET", "some-client-secret")
+ monkeypatch.setenv("AZURE_TENANT_ID", "some-tenant-id")
+ 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()
+
+ with pytest.raises(ValueError, match="AZURE_CLIENT_ID,
AZURE_CLIENT_SECRET, AZURE_TENANT_ID"):
Review Comment:
The reply on the guard thread said both spellings were pinned, but only
`AZURE_CLIENT_*` is: deleting `AZURE_STORAGE_CLIENT_ID`,
`AZURE_STORAGE_CLIENT_SECRET`, `AZURE_STORAGE_TENANT_ID`,
`AZURE_STORAGE_AUTHORITY_ID` and `AZURE_AUTHORITY_ID` from the tuples leaves
all 71 tests in this file green. Parametrizing this test over the spellings
would close that. While here, the two no-raise tests above and below assert
only `"sas_query_pairs" in credentials`; the full-dict equality the round-2
thread asked for would pin the account as well.
--
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]
