kaxil commented on code in PR #75810:
URL: https://github.com/apache/airflow/pull/75810#discussion_r4467410122
##########
providers/snowflake/tests/unit/snowflake/hooks/test_sql_api.py:
##########
@@ -539,6 +539,23 @@ def
test_get_headers_uses_key_pair_for_workload_identity_connection_with_private
result = hook.get_headers()
assert result["X-Snowflake-Authorization-Token-Type"] == "KEYPAIR_JWT"
+
@mock.patch("airflow.providers.snowflake.utils._rest_auth._fetch_gcp_identity_token",
autospec=True)
+ @mock.patch(f"{HOOK_PATH}.get_private_key", autospec=True,
return_value=None)
+ @mock.patch(f"{HOOK_PATH}._get_conn_params", autospec=True)
+ def test_get_headers_uses_workload_identity_without_private_key(
+ self, mock_conn_param, mock_private_key, mock_fetch_token
+ ):
+ mock_conn_param.return_value = {
+ **CONN_PARAMS,
+ "region": "us-east-2.aws",
+ "workload_identity_provider": "GCP",
Review Comment:
With the new `WORKLOAD_IDENTITY` check this test never reaches WIF:
`CONN_PARAMS` sets `authenticator: "snowflake"`, so `get_headers()` raises the
"none is configured" `ValueError`.
`test_request_sends_workload_identity_headers` in `test_cortex_agent.py` fails
the same way, since it sets no authenticator at all. Adding `"authenticator":
"WORKLOAD_IDENTITY"` to both gets them passing again.
##########
providers/snowflake/docs/connections/snowflake.rst:
##########
@@ -189,3 +189,28 @@ and is granted a role with access to the target objects.
"workload_identity_provider": "GCP"
}
}'
+
+``SnowflakeSqlApiHook``, ``SnowflakeCortexAgentHook`` and the
``pydanticai_snowflake`` connection type call
+Snowflake's REST APIs, which accept the workload's identity token directly
instead of a connector session.
+The token is requested from the cloud and reused until shortly before it
expires. A connection that also
+sets a private key uses key-pair authentication for these REST hooks. Each
provider needs its cloud's client
+library installed:
+
+* ``AWS``: ``boto3`` with ``botocore>=1.41.0`` for ``sts:GetWebIdentityToken``
support.
+ The workload's IAM role must be allowed to
+ call it, and AWS outbound identity federation must be enabled for the
account. The call is not available on
Review Comment:
`SnowflakeHook` doesn't need outbound identity federation on AWS: connector
4.8.1 signs a `GetCallerIdentity` request by default and only calls
`get_web_identity_token` when `workload_identity_aws_use_outbound_token` or
`SNOWFLAKE_ENABLE_AWS_WIF_OUTBOUND_TOKEN` is set, and
`_add_workload_identity_config` passes neither. So on an account without it,
`SnowflakeHook` works and these hooks fail, which the "same connection works
for" line up top doesn't hint at. Worth a sentence saying only the REST hooks
need it?
##########
providers/snowflake/tests/unit/snowflake/utils/test__rest_auth.py:
##########
@@ -40,6 +48,16 @@
"password": "my_pat_token_value",
}
CONN_PARAMS_KEYPAIR = {"account": "airflow", "authenticator": "snowflake",
"user": "user"}
Review Comment:
With `authenticator: "snowflake"` here,
`test_get_token_workload_identity_with_private_key_uses_key_pair`
(`{**CONN_PARAMS_KEYPAIR, "workload_identity_provider": "AWS"}`) no longer
qualifies for WIF, so it would pass even if WIF were picked ahead of the
private key. Same for
`test_get_headers_uses_key_pair_for_workload_identity_connection_with_private_key`
in `test_sql_api.py`. Setting `authenticator` to `WORKLOAD_IDENTITY` in those
two would keep them testing the precedence.
##########
providers/snowflake/src/airflow/providers/snowflake/utils/_rest_auth.py:
##########
@@ -163,6 +242,104 @@ def build_auth_headers(self) -> dict[str, str]:
}
+def _validate_wif_host(conn_config: dict[str, Any]) -> None:
+ """Restrict ambient workload credentials to trusted Snowflake
destinations."""
+ host =
get_cortex_base_url(conn_config).removeprefix("https://";).lower().removesuffix(".")
+ suffixes = ["snowflakecomputing.com", "snowflakecomputing.cn",
"snowflakecomputing.mil"]
+ suffixes.extend(
+ entry.strip().lower().strip(".")
+ for entry in os.environ.get("SNOWFLAKE_WIF_ALLOWED_HOST_SUFFIXES",
"").split(",")
+ if entry.strip()
+ )
+ dns_name =
r"[a-z0-9_](?:[a-z0-9_-]{0,61}[a-z0-9_])?(?:\.[a-z0-9_](?:[a-z0-9_-]{0,61}[a-z0-9_])?)*"
+ if (
+ len(host) <= 253
+ and re.fullmatch(dns_name, host)
+ and any(
+ re.fullmatch(dns_name, suffix) and (host == suffix or
host.endswith(f".{suffix}"))
+ for suffix in suffixes
+ )
+ ):
+ return
+ raise ValueError("Workload Identity Federation requires a recognized
Snowflake host.")
+
+
+def _fetch_aws_identity_token() -> str:
+ """Return the JWT from regional STS for the workload's AWS identity."""
+ try:
+ import boto3
+ from botocore.session import get_session
+ from botocore.utils import InstanceMetadataRegionFetcher
+ except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(
+ "Workload Identity Federation with the AWS provider requires
boto3. Install boto3."
+ ) from e
+ region = (
+ os.environ.get("AWS_REGION")
+ or os.environ.get("AWS_DEFAULT_REGION")
+ or InstanceMetadataRegionFetcher().retrieve_region()
+ )
+ if not region:
+ raise ValueError(
+ "AWS workload identity requires a region: set AWS_REGION or
AWS_DEFAULT_REGION, "
+ "or enable EC2 instance metadata region discovery."
+ )
+ session = get_session()
+ session.set_config_variable("sts_regional_endpoints", "regional")
+ sts = boto3.Session(botocore_session=session).client("sts",
region_name=region)
+ if not hasattr(sts, "get_web_identity_token"):
+ raise AirflowOptionalProviderFeatureException(
+ "Workload Identity Federation with the AWS provider requires a
boto3 version that "
+ "supports sts:GetWebIdentityToken (botocore>=1.41.0). Upgrade
boto3."
+ )
+ response = sts.get_web_identity_token(Audience=[WIF_AUDIENCE],
SigningAlgorithm="ES384")
+ return response["WebIdentityToken"]
+
+
+def _fetch_azure_identity_token() -> str:
+ """
+ Return an Entra access token for Snowflake's application, issued to the
workload's identity.
+
+ Uses the AKS workload identity when the pod has one (the Azure Workload
Identity webhook sets
+ ``AZURE_CLIENT_ID``, ``AZURE_TENANT_ID`` and
``AZURE_FEDERATED_TOKEN_FILE``), and the managed
+ identity otherwise. A user-assigned managed identity is selected with
+ ``MANAGED_IDENTITY_CLIENT_ID``.
+ """
+ try:
+ from azure.identity import ManagedIdentityCredential,
WorkloadIdentityCredential
+ except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(
+ "Workload Identity Federation with the AZURE provider requires
azure-identity. "
+ "Install apache-airflow-providers-snowflake[microsoft.azure]."
+ ) from e
+ credential: ManagedIdentityCredential | WorkloadIdentityCredential
+ if all(
+ os.environ.get(name) for name in ("AZURE_CLIENT_ID",
"AZURE_TENANT_ID", "AZURE_FEDERATED_TOKEN_FILE")
+ ) and os.path.exists(os.environ["AZURE_FEDERATED_TOKEN_FILE"]):
+ credential = WorkloadIdentityCredential()
+ else:
+ credential =
ManagedIdentityCredential(client_id=os.environ.get("MANAGED_IDENTITY_CLIENT_ID"))
Review Comment:
This one's on me, I missed it when I suggested the file check:
`ManagedIdentityCredential()` picks `WorkloadIdentityCredential` by itself
whenever `AZURE_AUTHORITY_HOST`, `AZURE_TENANT_ID` and
`AZURE_FEDERATED_TOKEN_FILE` are set (azure-identity 1.27.0), and it doesn't
look at the file. The AKS webhook sets all three, so a missing mount fails here
anyway instead of falling back to IMDS like the connector does. The only
opt-out I found is the private `_exclude_workload_identity_credential=True`
kwarg, and the test mocks both classes so it can't see this. I'd probably just
drop the existence check, unless you want to take on the private kwarg?
##########
providers/snowflake/src/airflow/providers/snowflake/utils/_rest_auth.py:
##########
@@ -117,18 +141,18 @@ def get_token(self) -> SnowflakeRestToken:
)
return SnowflakeRestToken(token=pat,
token_type="PROGRAMMATIC_ACCESS_TOKEN")
- if self._private_key is None:
+ if not self._private_key_loaded:
self._private_key = self._private_key_loader()
+ self._private_key_loaded = True
if self._private_key is None:
- if conn_config.get("workload_identity_provider"):
- raise ValueError(
- "Workload identity federation is not supported for
Snowflake REST APIs; use "
- "OAuth, PAT, or key-pair authentication instead. A
connection that also "
- "sets a private key uses key-pair authentication."
- )
+ if (
+ conn_config.get("authenticator") or ""
+ ).upper() == "WORKLOAD_IDENTITY" and
conn_config.get("workload_identity_provider"):
+ return self._get_workload_identity_token(conn_config)
raise ValueError(
"Snowflake REST API authentication requires an OAuth
access token, a "
- "Programmatic Access Token (PAT), or a private key for
key-pair JWT auth; "
+ "Programmatic Access Token (PAT), a private key for
key-pair JWT auth, or a "
+ "workload_identity_provider for Workload Identity
Federation; "
Review Comment:
A connection that sets `workload_identity_provider` but not `authenticator`
now ends up here, and the message tells them to set the extra they already
have. The connector raises "workload_identity_provider was set but
authenticator was not set to WORKLOAD_IDENTITY" for that case. Could this do
the same?
--
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]
