pankajastro commented on code in PR #73374:
URL: https://github.com/apache/airflow/pull/73374#discussion_r4147144036
##########
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:
Parametrized over both spellings plus the authority-id variant — confirmed
the gap first by deleting the `AZURE_STORAGE_*` entries and rerunning:
everything stayed green before the fix. Also switched the two no-raise tests to
full-dict equality.
---
Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting
##########
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:
Went with a tier-dispatch method rather than a literal mapping — the triple
checks need AND logic across spelling groups, which a flat tuple can't express
— and moved the constants next to `_AZURE_PUBLIC_SUFFIX` as you suggested.
---
Drafted-by: Claude Code (Sonnet 5); reviewed by @pankajastro before posting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
