omar-massfih commented on code in PR #75810:
URL: https://github.com/apache/airflow/pull/75810#discussion_r4402078896
##########
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:
Changed to `compute_engine.IDTokenCredentials` with
`use_metadata_identity_endpoint=True`, followed by refresh. Added coverage with
`GOOGLE_APPLICATION_CREDENTIALS` set.
##########
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:
Added the loaded flag. A successful `None` result is cached across WIF cache
hits and refreshes; loader errors remain retryable. Both behaviors have
regression tests.
##########
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:
Changed these errors to `AirflowOptionalProviderFeatureException` with
installation or upgrade guidance. Azure uses the existing `microsoft.azure`
extra.
##########
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:
Moved the prerequisites below the example and added WIF to both
authentication lists, preserving private-key precedence. The Snowflake
documentation build passes.
--
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]
