Skip to content

Conversation

@SameerMesiah97
Copy link
Contributor

@SameerMesiah97 SameerMesiah97 commented Jan 1, 2026

Description

Refactored OAuth token handling in SnowflakeHook to ensure expired access tokens are refreshed for long-running tasks.

Connection resolution has been split into static and dynamic layers by introducing a cached _get_static_conn_params and a dynamic _get_conn_params method (previously a cached property) that resolves OAuth tokens at access time. OAuth token request logic is centralized in a new internal helper, _get_valid_oauth_token, which validates cached tokens and refreshes them on expiry.

Both get_oauth_token and _get_conn_params now delegate token resolution to _get_valid_oauth_token, ensuring consistent behavior and preventing reuse of expired tokens.

Rationale

OAuth access tokens are time-bound credentials and should not be cached alongside static connection configuration. Previously, caching connection parameters caused expired tokens to be reused within long-running tasks, even when valid refresh credentials were available.

By separating static connection parameters from token resolution and revalidating tokens at access time, the hook now handles both short- and long-running workloads correctly without sacrificing the benefits of caching immutable configuration.

_get_conn_params was converted from a property to a method to make its dynamic behavior and potential OAuth token refresh explicit, avoiding the misleading implication of cheap, side-effect-free attribute access.`

Tests

  • Added a unit test covering the refresh-token flow that simulates token expiry in a long-running task and verifies that expired tokens are refreshed while static connection parameters remain unchanged.
  • Updated existing mocks to include expires_in so token expiry timestamps can be computed correctly, and adjusted tests to account for the newly added OAuth token request timeout.
  • As _get_conn_params is now a method instead of a property, the relevant tests have been changed to reflect that.

Notes

  • Removed a cyclical call relationship between get_oauth_token and _get_conn_params by centralizing token validation and refresh logic in _get_valid_oauth_token.
  • Added a 30-second timeout (refer to OAUTH_REQUEST_TIMEOUT) to OAuth token requests to prevent tasks from hanging indefinitely when the token endpoint does not respond.
  • Fixed docstring for test_get_oauth_token_without_scope to match the actual behavior of get_oauth_token.

Backwards Compatibility

This change does not modify the public get_oauth_token API. The method signature is unchanged, and default user-facing behavior for the refresh-token flow and grant-type handling remains the same as prior to this refactor.

Closes: #60023

@SameerMesiah97 SameerMesiah97 requested a review from potiuk as a code owner January 1, 2026 17:38
@boring-cyborg boring-cyborg bot added area:providers provider:snowflake Issues related to Snowflake provider labels Jan 1, 2026
@SameerMesiah97 SameerMesiah97 force-pushed the 60023-SFHook-Token-Refactor branch 3 times, most recently from 046859a to a6b6ff6 Compare January 1, 2026 20:48
@SameerMesiah97
Copy link
Contributor Author

Requesting review for this.

@potiuk potiuk merged commit 2c5487c into apache:main Jan 8, 2026
87 checks passed
@potiuk
Copy link
Member

potiuk commented Jan 8, 2026

Nice

chirodip98 pushed a commit to chirodip98/airflow-contrib that referenced this pull request Jan 9, 2026
…tests. (apache#60027)

Co-authored-by: Sameer Mesiah <smesiah971@gmail.com>
stegololz pushed a commit to stegololz/airflow that referenced this pull request Jan 9, 2026
…tests. (apache#60027)

Co-authored-by: Sameer Mesiah <smesiah971@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:providers provider:snowflake Issues related to Snowflake provider

Projects

None yet

Development

Successfully merging this pull request may close these issues.

SnowflakeHook reuses expired OAuth tokens for long-running tasks

2 participants