[omar-massfih]

The problem

SnowflakeSqlApiOperator can't be used with Workload Identity Federation (WIF). If a Snowflake connection is set up for WIF (authenticator=WORKLOAD_IDENTITY), the SQL API hook fails with:

Workload identity federation is not supported for Snowflake REST APIs

SnowflakeHook already supports WIF (#68107), and the connection docs say the same connection works for the SQL API hook, so the two disagreed. The same limit applied to the Cortex Agent hook and the pydanticai_snowflake connection type, because they share the same authentication code.

Snowflake's REST API supports WIF: you send the cloud identity token as Authorization: Bearer WIF.<provider>.<token>, together with the header X-Snowflake-Authorization-Token-Type: WORKLOAD_IDENTITY_FEDERATION (Snowflake docs).

closes: #75808

What this changes

A WIF connection now works for these hooks. The fix is one new branch in SnowflakeRestTokenProvider (utils/_rest_auth.py), which all the REST hooks use.

• How the token is obtained. The provider asks the cloud for an identity token using that cloud's own library:

  Provider	Library	How
  AWS	boto3	sts:GetWebIdentityToken
  AZURE	azure-identity	AKS workload identity if the pod has one, otherwise the managed identity (MANAGED_IDENTITY_CLIENT_ID picks a user-assigned one)
  GCP	google-auth	ID token for the workload's service account
  OIDC	none	the connection's token or token_file_path

• No new dependencies. Each library is imported only when a connection uses that cloud. If it's missing, the error names the package to install.

• Existing behaviour is unchanged. The order is still OAuth, then PAT, then key-pair. A connection that has a private key keeps using key-pair, even if WIF is also configured. WIF is only used when nothing else applies.

• Tokens are cached. The SQL API operator checks its status repeatedly, so the token is reused until 30 seconds before it expires, instead of calling the cloud every time.

• Docs. The connection page now says what each provider needs installed.

Worth a look

• The cloud libraries are not dependencies of the provider. Should boto3, azure-identity and google-auth be added as optional extras?
• workload_identity_entra_resource and workload_identity_impersonation_path aren't passed through by SnowflakeHook today, so they aren't supported here either.

How it was tested

Unit tests cover each cloud's token helper (with fake cloud libraries), OIDC from a token or a file, caching and refresh near expiry, missing libraries, and that key-pair still wins. They also check the headers from SnowflakeSqlApiHook.get_headers() and a Cortex Agent request. 171 tests pass, and ruff and mypy are clean on the changed code.

Live test on Azure: an Azure Container Apps job, running this branch with a managed identity and no stored secret, ran a query through SnowflakeSqlApiHook against a real Snowflake account. The token came from azure-identity, the SQL API accepted Bearer WIF.AZURE.<token>, and select current_user(), current_role() finished successfully.

Not yet tested live: AKS workload identity, AWS, GCP and OIDC (unit tests only).

Two notes on the setup: I ran the tests in a scratch environment because the repo needs a newer uv than mine, and I skipped the mypy-devel-common pre-commit hook, which can't run on that uv; CI will cover it.

---

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)

Generated-by: Claude Code (Claude Sonnet 5.5) following the guidelines