kaxil commented on code in PR #75810:
URL: https://github.com/apache/airflow/pull/75810#discussion_r4394293564
##########
providers/snowflake/src/airflow/providers/snowflake/utils/_rest_auth.py:
##########
@@ -154,6 +171,60 @@ def get_token(self) -> SnowflakeRestToken:
token = self._jwt_generator.get_token()
return SnowflakeRestToken(token=token, token_type="KEYPAIR_JWT")
+ def _get_workload_identity_token(self, conn_config: dict[str, Any]) ->
SnowflakeRestToken:
+ """
+ Return a Workload Identity Federation token, fetching a new one when
the cached one expires.
+
+ The caller must hold ``self._lock``.
+ """
+ if self._wif_token is not None and time.time() <
self._wif_token_expires_at:
+ return self._wif_token
+
+ provider_name = str(conn_config["workload_identity_provider"]).upper()
Review Comment:
Should this check the host first? The connector refuses WORKLOAD_IDENTITY
unless the host ends in `snowflakecomputing.com`, `.cn` or `.mil` (or a
`SNOWFLAKE_WIF_ALLOWED_HOST_SUFFIXES` entry), and it does that before fetching
anything from the cloud (`_verify_host_allowed_for_workload_identity` in 4.8;
5.0 has it in the core).
The Cortex hooks take their URL straight from the `host` extra via
`get_cortex_base_url`, so with this branch a connection pointing `host`
somewhere else gets sent the worker's own identity token. A test that a
non-Snowflake host fails before any `_fetch_*` call would pin it.
##########
providers/snowflake/src/airflow/providers/snowflake/utils/_rest_auth.py:
##########
@@ -163,6 +234,64 @@ def build_auth_headers(self) -> dict[str, str]:
}
+def _fetch_aws_identity_token() -> str:
+ """Return the JWT from ``sts:GetWebIdentityToken`` for the workload's AWS
identity."""
+ try:
+ import boto3
+ except ImportError as e:
+ raise ValueError(
+ "Workload Identity Federation with the AWS provider requires the
boto3 package."
+ ) from e
+ sts = boto3.client("sts")
Review Comment:
`boto3.client("sts")` with no `region_name` only reads `AWS_DEFAULT_REGION`
or the config file, not `AWS_REGION` or the EC2 metadata region. When it finds
none it resolves to `aws-global` and calls `https://sts.amazonaws.com/`, where
`GetWebIdentityToken` isn't available. I checked with botocore 1.43.116 and
only `AWS_REGION=eu-west-1` set: the client region came back as `aws-global`
and the request went to the global endpoint.
So on a plain EC2 host, or anywhere only `AWS_REGION` is set,
`SnowflakeHook` authenticates (the connector's `get_aws_region` reads
`AWS_REGION`, then `AWS_DEFAULT_REGION`, then instance metadata) and the REST
hooks get an STS error. Could this resolve the region the same way, pass
`region_name=`, and raise a clear error when nothing is found? The test at line
370 pins `client("sts")` with no region, so it would change too.
##########
providers/snowflake/src/airflow/providers/snowflake/utils/_rest_auth.py:
##########
@@ -163,6 +234,64 @@ def build_auth_headers(self) -> dict[str, str]:
}
+def _fetch_aws_identity_token() -> str:
+ """Return the JWT from ``sts:GetWebIdentityToken`` for the workload's AWS
identity."""
+ try:
+ import boto3
+ except ImportError as e:
+ raise ValueError(
+ "Workload Identity Federation with the AWS provider requires the
boto3 package."
+ ) from e
+ sts = boto3.client("sts")
+ if not hasattr(sts, "get_web_identity_token"):
+ raise ValueError(
+ "Workload Identity Federation with the AWS provider requires a
boto3 version that "
+ "supports sts:GetWebIdentityToken."
+ )
+ 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 ValueError(
+ "Workload Identity Federation with the AZURE provider requires the
azure-identity package."
+ ) from e
+ credential: ManagedIdentityCredential | WorkloadIdentityCredential
+ if all(
+ os.environ.get(name) for name in ("AZURE_CLIENT_ID",
"AZURE_TENANT_ID", "AZURE_FEDERATED_TOKEN_FILE")
+ ):
+ credential = WorkloadIdentityCredential()
+ else:
+ credential =
ManagedIdentityCredential(client_id=os.environ.get("MANAGED_IDENTITY_CLIENT_ID"))
+ try:
+ return credential.get_token(f"{WIF_AZURE_RESOURCE}/.default").token
+ finally:
+ credential.close()
+
+
+def _fetch_gcp_identity_token() -> str:
+ """Return a Google ID token for Snowflake's audience, issued to the
workload's service account."""
+ try:
+ from google.auth.transport.requests import Request
+ from google.oauth2 import id_token
+ except ImportError as e:
+ raise ValueError(
+ "Workload Identity Federation with the GCP provider requires the
google-auth package."
+ ) from e
+ return id_token.fetch_id_token(Request(), WIF_AUDIENCE)
Review Comment:
`id_token.fetch_id_token` uses a service-account key file from
`GOOGLE_APPLICATION_CREDENTIALS` before it tries the metadata server, while the
connector always asks the metadata server. On a worker that sets
`GOOGLE_APPLICATION_CREDENTIALS` for other GCP hooks, the REST hooks would
present that key's service account and `SnowflakeHook` the workload's, and the
Snowflake WIF user trusts only one of them. Would
`compute_engine.IDTokenCredentials(Request(), WIF_AUDIENCE,
use_metadata_identity_endpoint=True)` work here? Otherwise a docs note on the
precedence would help.
##########
providers/snowflake/src/airflow/providers/snowflake/utils/_rest_auth.py:
##########
@@ -121,14 +141,11 @@ def get_token(self) -> SnowflakeRestToken:
self._private_key = self._private_key_loader()
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."
- )
+ return self._get_workload_identity_token(conn_config)
Review Comment:
For a WIF connection `self._private_key` stays `None`, so the loader runs
again on every `get_token()`, including calls that return the cached WIF token.
Both loaders (`hook.get_private_key` and
`SnowflakeSqlApiHook._load_private_key`) call `get_connection`, so every SQL
API status poll now does a connection lookup that `_get_conn_params` otherwise
avoids by caching. A separate "loaded" flag would keep the docstring's "called
at most once" true.
##########
providers/snowflake/src/airflow/providers/snowflake/utils/_rest_auth.py:
##########
@@ -163,6 +234,64 @@ def build_auth_headers(self) -> dict[str, str]:
}
+def _fetch_aws_identity_token() -> str:
+ """Return the JWT from ``sts:GetWebIdentityToken`` for the workload's AWS
identity."""
+ try:
+ import boto3
+ except ImportError as e:
+ raise ValueError(
Review Comment:
Elsewhere in this provider a missing optional dependency raises
`AirflowOptionalProviderFeatureException` (`hooks/snowflake.py:475`,
`hooks/cortex_model.py:40`). A `ValueError` here reads like a connection
misconfiguration. Same for the azure-identity and google-auth branches below.
On your extras question: `azure-identity` already comes in through the
`microsoft.azure` extra.
##########
providers/snowflake/tests/unit/snowflake/utils/test__rest_auth.py:
##########
@@ -227,6 +346,96 @@ def
test_build_auth_headers_contains_only_auth_headers(self, mock_conn_params):
assert headers["X-Snowflake-Authorization-Token-Type"] ==
"PROGRAMMATIC_ACCESS_TOKEN"
+class TestFetchIdentityTokens:
+ @staticmethod
+ def _fake_google() -> tuple[mock.MagicMock, dict[str, mock.MagicMock]]:
+ google = mock.MagicMock()
+ modules = {
+ "google": google,
+ "google.auth": google.auth,
+ "google.auth.transport": google.auth.transport,
+ "google.auth.transport.requests": google.auth.transport.requests,
+ "google.oauth2": google.oauth2,
+ "google.oauth2.id_token": google.oauth2.id_token,
+ }
+ return google, modules
+
+ def test_aws_requests_web_identity_token_for_snowflake(self):
+ boto3 = mock.MagicMock()
Review Comment:
boto3, azure-identity and google-auth are installed in breeze, so these
could exercise the real libraries instead of bare `MagicMock` modules:
`botocore.stub.Stubber` on a real STS client checks the operation name and
params against the service model, and `mock.patch(..., autospec=True)` on
`ManagedIdentityCredential`, `WorkloadIdentityCredential` and `fetch_id_token`
checks the signatures. As written, a misspelled method or keyword would still
pass, because the assertion repeats whatever the code calls.
##########
providers/snowflake/docs/connections/snowflake.rst:
##########
@@ -175,6 +175,20 @@ workload runs on (here ``GCP``). No password, key-pair or
token is stored; the w
is the credential. The Snowflake side needs a ``TYPE = SERVICE`` user that
trusts the workload's identity
and is granted a role with access to the target objects.
+``SnowflakeSqlApiHook``, ``SnowflakeCortexAgentHook`` and the
``pydanticai_snowflake`` connection type call
Review Comment:
This paragraph now sits between the intro that says "(here ``GCP``)" and the
code block it refers to. Moving it below the example, or under its own
subheading, keeps the two together.
`operators/cortex_agent.rst` and `connections/pydantic_ai_snowflake.rst`
also still list only OAuth, PAT and key-pair, with key-pair as "the default
when neither of the above is set". Both could get a WIF bullet that links here.
--
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]
