omar-massfih opened a new pull request, #75810: URL: https://github.com/apache/airflow/pull/75810
`SnowflakeSqlApiHook` (used by `SnowflakeSqlApiOperator`), `SnowflakeCortexAgentHook` and the `pydanticai_snowflake` connection type raised `Workload identity federation is not supported for Snowflake REST APIs` for connections that use `authenticator=WORKLOAD_IDENTITY`, although `SnowflakeHook` has supported them since #68107. The connection docs even listed these hooks as working with a WIF connection. Snowflake's REST API documentation lists Workload Identity Federation as a supported method ([Authenticating Snowflake REST APIs with Snowflake](https://docs.snowflake.com/en/developer-guide/snowflake-rest-api/authentication)): ``` Authorization: Bearer WIF.{provider}.{token} X-Snowflake-Authorization-Token-Type: WORKLOAD_IDENTITY_FEDERATION ``` closes: #75808 ## What changed All REST callers share `SnowflakeRestTokenProvider`, so the change is one new branch in `utils/_rest_auth.py`: - A connection with `workload_identity_provider` and no private key now gets a WIF token instead of the error. The branch order is otherwise unchanged: OAuth, PAT, private key (key-pair still wins when a private key is also set), then WIF. - The platform token comes from `snowflake.connector.wif_util.create_attestation`, so the AWS, Azure, GCP and OIDC flows match the connector. OIDC reads `token` or `token_file_path`. - For AWS the REST API needs the JWT issued by `sts:GetWebIdentityToken`, not the signed STS request the connector uses for its login protocol, so `aws_use_outbound_token=True` is passed. That needs `snowflake-connector-python>=4.7.1`; with an older connector the user gets an explicit error naming the version. - The token is cached until 30 seconds before its `exp` claim, since the SQL API operator polls `get_headers()` repeatedly. A credential without a readable `exp` is not cached. - Docs: the connection page now describes the REST behavior and the AWS requirements. ## Things for reviewers to look at - `snowflake.connector.wif_util` is not a documented public API of the connector. It is imported lazily and only on this branch. Is relying on it acceptable, or would you rather have the provider implement the attestation itself? - I did not raise the provider's `snowflake-connector-python` floor (the AWS path needs 4.7.1, the floor is 3.17.0). It is handled with a runtime error instead. Say if you would rather bump it. - `workload_identity_entra_resource` and `workload_identity_impersonation_path` are not forwarded by `SnowflakeHook` today, so they are not supported here either. ## Testing - Unit tests for each provider, OIDC token and token file, caching and refresh near expiry, the old-connector error, and key-pair precedence. Also `get_headers()` on `SnowflakeSqlApiHook` and a Cortex Agent request. The old "WIF raises" test was updated. - `pytest` on `test__rest_auth.py`, `test_sql_api.py`, `test_cortex_agent.py` and `test_cortex_model.py`: 163 passed (the `test_cortex_model.py` file skipped because `common.ai` was not installed in my scratch environment). Ran with `--noconftest` in a scratch venv with connector 4.8.1, because the repo requires `uv>=0.13.4` and mine is older. - `ruff check`/`ruff format` and `mypy` on the changed module are clean. The `mypy-devel-common` pre-commit hook could not run for the same `uv` reason, so I skipped that hook only; CI will cover it. - **Not tested against a live Snowflake account or a real cloud workload.** Snowflake's page says "REST APIs" and does not state explicitly that the SQL API endpoint accepts WIF tokens. Opening as a draft until this is verified. --- ##### Was generative AI tooling used to co-author this PR? - [X] Yes (please specify the tool below) Generated-by: Claude Code (Claude Sonnet 5.5) following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions) -- 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]
