Refresh expired OAuth tokens in SnowflakeHook for long-running tasks #60027
+242
−96
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
Refactored OAuth token handling in
SnowflakeHookto 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_paramsand a dynamic_get_conn_paramsmethod (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_tokenand_get_conn_paramsnow 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_paramswas 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
expires_inso token expiry timestamps can be computed correctly, and adjusted tests to account for the newly added OAuth token request timeout._get_conn_paramsis now a method instead of a property, the relevant tests have been changed to reflect that.Notes
get_oauth_tokenand_get_conn_paramsby centralizing token validation and refresh logic in_get_valid_oauth_token.OAUTH_REQUEST_TIMEOUT) to OAuth token requests to prevent tasks from hanging indefinitely when the token endpoint does not respond.test_get_oauth_token_without_scopeto match the actual behavior ofget_oauth_token.Backwards Compatibility
This change does not modify the public
get_oauth_tokenAPI. 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