Skip to content

Support PAT and key-pair auth in the Snowflake Cortex Agent hook - #73932

Merged
Lee-W merged 2 commits into
apache:mainfrom
astronomer:snowflake-cortex-auth-model
Oct 7, 2026
Merged

Lee-W merged 2 commits into
apache:mainfrom
astronomer:snowflake-cortex-auth-model

Conversation

@Lee-W

@Lee-W Lee-W commented Sep 30, 2026 •

Copy link
Copy Markdown
Member

SnowflakeCortexAgentHook only accepted an OAuth access token, so ones whose
Snowflake connections use key-pair or programmatic access tokens could not
call Cortex Agents without a second, OAuth-only connection. The SQL API hook
already minted all 3 header kinds; pulling that logic into a shared token
provider lets both hooks authenticate the same way and keeps a key-pair JWT
cached across requests until its renewal window, instead of re-signing on
every call. A renewal delta at or beyond the token lifetime now re-signs every
call rather than handing out an expired JWT.

It also adds PydanticAISnowflakeHook (the pydanticai_snowflake connection
type), which points a Common AI model at an existing Snowflake connection, so
Cortex chat completions reuse the same credentials.

A connection with authenticator: programmatic_access_token and the PAT in
Password now also works for SnowflakeHook, so one connection serves the SQL
connector, the SQL API hook and both Cortex hooks. Without a host extra, the
Cortex base URL now appends the connection's region the way
SnowflakeHook.account_identifier does, so locator-style accounts outside AWS
us-west-2 reach the right host; a connection that uses an org-account
identifier but still sets region gets a region-suffixed host, as the SQL API
hook already does.


Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

Generated-by: [Claude] following the guidelines


  • Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
  • For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

@Lee-W
Lee-W force-pushed the snowflake-cortex-auth-model branch 2 times, most recently from a6b4496 to 955a2c5 Compare October 1, 2026 09:59
Comment thread providers/snowflake/docs/index.rst
@Lee-W
Lee-W marked this pull request as ready for review October 1, 2026 12:36
Comment thread providers/snowflake/src/airflow/providers/snowflake/utils/rest_auth.py Outdated
Comment thread providers/snowflake/docs/connections/pydantic_ai_snowflake.rst
Comment thread providers/snowflake/docs/connections/pydantic_ai_snowflake.rst
Comment thread providers/common/ai/docs/model_providers.rst
Comment thread providers/snowflake/docs/operators/snowflake_cortex_agent.rst Outdated
@Lee-W
Lee-W force-pushed the snowflake-cortex-auth-model branch from 955a2c5 to cc80a25 Compare October 2, 2026 06:48
Lee-W added 2 commits October 6, 2026 06:13
SnowflakeCortexAgentHook only accepted an OAuth access token, so teams whose
Snowflake connections use key-pair or programmatic access tokens could not
call Cortex Agents without a second, OAuth-only connection. The SQL API hook
already minted all three header kinds; pulling that logic into a shared token
provider lets both hooks authenticate the same way and keeps a key-pair JWT
cached across requests until its renewal window, instead of re-signing on
every call. A renewal delta at or beyond the token lifetime now re-signs every
call rather than handing out an expired JWT.
Cortex chat models could only be reached through the snowflake: model prefix
with SNOWFLAKE_ACCOUNT and SNOWFLAKE_TOKEN set on the worker: one credential
per worker, no secrets backend, and no key-pair auth since the JWT expires
within an hour. Setting a password on a pydanticai connection failed outright
because the base hook passes it as api_key, which SnowflakeProvider rejects.
Building the provider from an existing Snowflake connection, with the token
refreshed per request, lets long agent runs outlive a single JWT and reuses
the auth teams already configured for their warehouse.

The hook sits behind an optional common.ai extra so the provider keeps
working without Common AI.
@Lee-W
Lee-W force-pushed the snowflake-cortex-auth-model branch from cc80a25 to 5037697 Compare October 6, 2026 03:15

@kaxil kaxil left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, LGTM

@Lee-W
Lee-W merged commit cdcbf12 into apache:main Oct 7, 2026
99 checks passed
@Lee-W
Lee-W deleted the snowflake-cortex-auth-model branch October 7, 2026 05:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants