fix(rocky-snowflake): set auth-token-type header per auth mode#283
Merged
hugocorreia90 merged 1 commit intomainfrom Apr 29, 2026
Merged
fix(rocky-snowflake): set auth-token-type header per auth mode#283hugocorreia90 merged 1 commit intomainfrom
hugocorreia90 merged 1 commit intomainfrom
Conversation
…th mode The connector hard-coded `X-Snowflake-Authorization-Token-Type: KEYPAIR_JWT` on every request, regardless of which auth mode was active. For OAuth and password / session-token users this caused Snowflake to reject the request as an invalid JWT, the connector's 401 retry path would drop the cache and re-mint, and the same wrong-typed token would 401 again — burning retries for what was, fundamentally, a header bug. Per Snowflake's SQL API v2 spec (https://docs.snowflake.com/en/developer-guide/sql-api/authenticating): Auth mode | X-Snowflake-Authorization-Token-Type ----------------------|-------------------------------------- Key-pair JWT (RS256) | KEYPAIR_JWT OAuth | OAUTH Password / session | (omit; server sniffs from token) This adds a public `auth::TokenType` enum with a `header_value()` mapping, exposes `Auth::token_type()`, and routes submit/poll/cancel through a single `apply_auth_headers` helper so the same per-mode header logic applies everywhere, not just submit. Three new wiremock tests assert the correct header rides each auth mode.
5 tasks
hugocorreia90
added a commit
that referenced
this pull request
Apr 29, 2026
Engine 1.18.0 ships the rocky preview workflow end-to-end (#279, #280, #281, #282), the [budget].max_bytes_scanned threshold (#288), the audit-sweep closeout (#283, #285–#287, #290–#293), and the rocky-server auth + CORS gate (#291). Dagster 1.15.0 picks up the regenerated Pydantic models for the rocky preview surface and ships the P1 cluster (#289) + FR-014 follow-on (#284). VS Code 1.10.0 regenerates TypeScript bindings for rocky preview and RunCostSummary.total_bytes_scanned. See per-artifact CHANGELOG entries for the full breakdown.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Summary
Snowflake's SQL REST API v2 connector hard-coded
X-Snowflake-Authorization-Token-Type: KEYPAIR_JWTon every request, even when the resolved auth mode was OAuth or password / session-token. For non-keypair users the warehouse rejected the request as an invalid JWT, the 401 retry path dropped the cache and re-minted the same wrong-typed token, and the loop kept failing — what looked like an auth/refresh problem was a header bug.Bug
engine/crates/rocky-snowflake/src/connector.rs:344(pre-fix) attachedKEYPAIR_JWTunconditionally on submit. Poll and cancel paths sent no token-type header at all. Tests attests/wiremock_tests.rs::test_bearer_token_sentencoded the broken assertion against an OAuth-configured connector — it passed only because the test mock never validated the token shape.Fix
Per Snowflake's SQL API v2 auth docs:
X-Snowflake-Authorization-Token-TypeKEYPAIR_JWTOAUTHChanges:
pub enum auth::TokenType { KeypairJwt, OAuth, SessionToken }with aheader_value() -> Option<&'static str>mapping.Auth::token_type()returning the resolved mode.SnowflakeConnector::apply_auth_headers()helper used bysubmit_and_wait_once, the polling loop, andcancel_statement— header logic now lives in one place and applies to every auth-bearing call site, not just submit.prime_cache_withtest helper onAuthis nowpubunderfeature = "test-support"so integration tests can skip the JWT/login roundtrip.Test plan
cargo test -p rocky-snowflake(107 unit tests pass)cargo test -p rocky-snowflake --features test-support --test wiremock_tests(22/22 pass, including 3 new cases:test_oauth_auth_headers,test_keypair_auth_headers,test_password_auth_omits_token_type_header, plustest_token_type_header_mapping)cargo clippy -p rocky-snowflake --all-targets --features test-support -- -D warningscleancargo fmt --checkcleanKnown follow-up (out of scope)
Snowflake's SQL API docs document the
Authorization: Bearer <token>format for OAuth and key-pair, but session tokens minted via/session/v1/login-requestare typically passed asAuthorization: Snowflake Token="<token>". The connector still usesBearerfor all three modes — that is unchanged by this PR. Password mode was already broken before this change; tracking a follow-up.