feat(rust): OAuth token primitives, OIDC discovery, cache and token store by vikrantpuppala · Pull Request #320 · adbc-drivers/databricks

vikrantpuppala · 2026-03-08T06:08:10Z

🥞 Stacked PR

Use this link to review incremental changes.

stack/oauth-u2m-m2m-design [Files changed]
- stack/pr-oauth-foundation [Files changed]
  - stack/pr-database-config [Files changed]
    - stack/pr-u2m-provider [Files changed]
      - stack/pr-integration-tests [Files changed]
        
        stack/pr-final-validation [Files changed]

Summary

Core OAuth token infrastructure used by both U2M and M2M flows:

OAuthToken (token.rs) — token struct with expiry tracking (40s buffer), stale detection (min(TTL*0.5, 20min)), and serde support
TokenStore (token_store.rs) — thread-safe token lifecycle state machine (Empty → Fresh → Stale → Expired) with RwLock + AtomicBool coordination, background refresh via tokio::task::spawn_blocking, panic-safe RefreshGuard
OIDC discovery (oidc.rs) — fetches authorization and token endpoints from /.well-known/oauth-authorization-server
TokenCache (cache.rs) — file-based U2M token persistence at ~/.config/databricks-adbc/oauth/ with SHA-256 hashed filenames and 0o600/0o700 permissions
DatabricksHttpClient (http.rs) — extended with OnceLock-based auth provider for two-phase initialization and execute_without_auth() for token endpoint calls
Cargo deps — oauth2, sha2, dirs, serde, open crates
License config — added MPL-2.0 and CDLA-Permissive-2.0 to about.toml for transitive deps

Key files

src/auth/oauth/token.rs — OAuthToken struct
src/auth/oauth/token_store.rs — token lifecycle state machine
src/auth/oauth/oidc.rs — OIDC endpoint discovery
src/auth/oauth/cache.rs — file-based token cache
src/client/http.rs — HTTP client auth provider integration

This pull request was AI-assisted by Isaac.

vikrantpuppala · 2026-03-12T07:10:30Z

Range-diff: stack/oauth-u2m-m2m-design (fe4a411 -> 1651567)

rust/docs/designs/oauth-u2m-m2m-design.md

@@ -19,41 +19,4 @@
 +
  # OAuth Authentication Design: U2M and M2M Flows
  
- ## Overview
-     participant DB as Database
-     participant M2M as ClientCredsProvider
-     participant OIDC as OIDC Discovery
--    participant TE as Token Endpoint
-+    participant TokenEP as Token Endpoint
- 
-     App->>DB: new_connection()
-     DB->>M2M: new(host, client_id, client_secret, scopes)
-     M2M->>OIDC: GET {host}/oidc/.well-known/oauth-authorization-server
-     OIDC-->>M2M: OidcEndpoints
- 
--    Note over App,TE: First get_auth_header() call
-+    Note over App,TokenEP: First get_auth_header() call
-     App->>M2M: get_auth_header()
--    M2M->>TE: client.exchange_client_credentials()
--    TE-->>M2M: access_token (no refresh_token)
-+    M2M->>TokenEP: client.exchange_client_credentials()
-+    TokenEP-->>M2M: access_token (no refresh_token)
-     M2M-->>App: "Bearer {access_token}"
- 
--    Note over App,TE: Subsequent calls (token FRESH)
-+    Note over App,TokenEP: Subsequent calls (token FRESH)
-     App->>M2M: get_auth_header()
-     M2M-->>App: "Bearer {cached_token}"
- 
--    Note over App,TE: Token becomes STALE
-+    Note over App,TokenEP: Token becomes STALE
-     App->>M2M: get_auth_header()
-     M2M->>M2M: Spawn background refresh
-     M2M-->>App: "Bearer {current_token}"
--    M2M->>TE: client.exchange_client_credentials() (background)
--    TE-->>M2M: new access_token
-+    M2M->>TokenEP: client.exchange_client_credentials() (background)
-+    TokenEP-->>M2M: new access_token
- ```
- 
- ### Token Exchange (M2M)
\ No newline at end of file
+ ## Overview
\ No newline at end of file

rust/src/auth/oauth/token_store.rs

@@ -28,7 +28,7 @@
 +//!   (computed as `min(initial_TTL * 0.5, 20 minutes)` when the token was
 +//!   first acquired). The current token is still valid but eligible for
 +//!   background refresh. The token is returned immediately to the caller,
-+//!   and a background refresh is spawned via `std::thread::spawn`.
++//!   and a background refresh is spawned via `tokio::task::spawn_blocking`.
 +//!
 +//! - **EXPIRED**: Token's remaining TTL is less than 40 seconds. The caller
 +//!   is blocked until a refresh completes. This ensures no requests are made
@@ -103,7 +103,7 @@
 +    /// - **Empty (no token)**: Blocks and calls `refresh_fn` to fetch the initial token.
 +    /// - **FRESH**: Returns the token immediately without calling `refresh_fn`.
 +    /// - **STALE**: Returns the current token immediately and spawns a background
-+    ///   refresh via `std::thread::spawn`. Only one background refresh runs at a time.
++    ///   refresh via `tokio::task::spawn_blocking`. Only one background refresh runs at a time.
 +    /// - **EXPIRED**: Blocks the caller and calls `refresh_fn` to obtain a fresh token
 +    ///   before returning.
 +    ///
@@ -224,9 +224,9 @@
 +        }
 +    }
 +
-+    /// Spawns a background refresh thread.
++    /// Spawns a background refresh task on the tokio blocking thread pool.
 +    ///
-+    /// This is called when the token is STALE. The background thread will
++    /// This is called when the token is STALE. The background task will
 +    /// attempt to refresh the token without blocking the caller.
 +    ///
 +    /// If a refresh is already in progress, this method does nothing (no-op).
@@ -240,12 +240,12 @@
 +            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
 +            .is_ok()
 +        {
-+            // Clone Arc pointers to move into the thread
++            // Clone Arc pointers to move into the task
 +            let token = self.token.clone();
 +            let refreshing = self.refreshing.clone();
 +
-+            // Spawn background thread to perform the refresh
-+            std::thread::spawn(move || {
++            // Spawn on tokio's blocking thread pool (reuses threads, cheaper than std::thread::spawn)
++            tokio::task::spawn_blocking(move || {
 +                let result = refresh_fn();
 +
 +                // Store the new token if successful (ignore errors in background)
@@ -259,7 +259,7 @@
 +                refreshing.store(false, Ordering::SeqCst);
 +            });
 +        }
-+        // If another thread is already refreshing, do nothing
++        // If another task is already refreshing, do nothing
 +    }
 +
 +    /// Waits for an in-progress refresh to complete, then returns the new token.
@@ -428,8 +428,8 @@
 +        }
 +    }
 +
-+    #[test]
-+    fn test_store_stale_returns_current_token() {
++    #[tokio::test(flavor = "multi_thread")]
++    async fn test_store_stale_returns_current_token() {
 +        let store = TokenStore::new();
 +        let refresh_count = Arc::new(AtomicUsize::new(0));
 +        let refresh_count_clone = refresh_count.clone();

_{Reproduce locally: git range-diff cf7d895..fe4a411 4dd8f51..1651567 | Disable: git config gitstack.push-range-diff false}

gopalldb · 2026-03-13T11:02:20Z

rust/src/auth/oauth/token_store.rs

+    fn wait_for_refresh(&self) -> Result<OAuthToken> {
+        // Spin-wait until the refresh completes
+        while self.refreshing.load(Ordering::SeqCst) {
+            std::thread::yield_now();


yield_now will consume cpu, can we add sleep instead?

check parking_lot::Condvar

Agreed — replaced yield_now() with sleep(1ms) to avoid busy-waiting. A Condvar would be cleaner but more invasive for the rare case where we actually hit the wait path (only when multiple threads race on an expired token simultaneously).

gopalldb · 2026-03-13T11:04:59Z

rust/src/auth/oauth/token_store.rs

+            }
+
+            // Release the refresh lock
+            self.refreshing.store(false, Ordering::SeqCst);


we should also set refresh to false for panic cases

Good catch — added a RefreshGuard RAII type that resets refreshing to false on Drop. Used in both blocking_refresh and spawn_background_refresh, so the flag is always cleared even if refresh_fn panics.

gopalldb · 2026-03-13T11:08:07Z

rust/src/client/http.rs

                .and_then(|b| b.as_bytes())
                .map(|b| b.to_vec())
        } else {
            None


you don't need request body ?

No request body needed — OIDC discovery is a standard GET request to the well-known endpoint. The server returns JSON with the authorization and token endpoints.

gopalldb · 2026-03-13T11:11:32Z

rust/src/auth/oauth/cache.rs

+
+        // Create cache directory if it doesn't exist
+        if let Some(parent) = cache_path.parent() {
+            if let Err(e) = fs::create_dir_all(parent) {


this usually uses 755 permission, through which others can also get access on the file

Good point — tightened directory permissions to 0o700 (owner only) after create_dir_all. The files themselves are already 0o600.

gopalldb · 2026-03-13T11:13:47Z

rust/src/auth/oauth/oidc.rs

+    /// ```
+    pub async fn discover(host: &str, http_client: &Arc<DatabricksHttpClient>) -> Result<Self> {
+        // Construct the well-known endpoint URL
+        let discovery_url = format!("{}/oidc/.well-known/oauth-authorization-server", host);


if host has trailing '/', the discovery url will have 2 '/'s

Good catch — added host.trim_end_matches('/') before building the discovery URL.

## 🥞 Stacked PR Use this [link](https://github.com/adbc-drivers/databricks/pull/319/files) to review incremental changes. - [**stack/oauth-u2m-m2m-design**](#319) [[Files changed](https://github.com/adbc-drivers/databricks/pull/319/files)] - [stack/pr-oauth-foundation](#320) [[Files changed](https://github.com/adbc-drivers/databricks/pull/320/files/250ff3d91c3001f671f08084f68e949e556bc5d2..bd474c189621aa70c1f14e97c32d64605275e07d)] - [stack/pr-database-config](#321) [[Files changed](https://github.com/adbc-drivers/databricks/pull/321/files/bd474c189621aa70c1f14e97c32d64605275e07d..296931cd396d82dccb1b548a51f6b9d31be3683e)] - [stack/pr-u2m-provider](#322) [[Files changed](https://github.com/adbc-drivers/databricks/pull/322/files/296931cd396d82dccb1b548a51f6b9d31be3683e..c96689981e79c04f43e8251f2cbd5690371dfca5)] - [stack/pr-integration-tests](#323) [[Files changed](https://github.com/adbc-drivers/databricks/pull/323/files/c96689981e79c04f43e8251f2cbd5690371dfca5..83d639337ca30688abb7bdba85aa16426d76eb31)] - [stack/pr-final-validation](#324) [[Files changed](https://github.com/adbc-drivers/databricks/pull/324/files/83d639337ca30688abb7bdba85aa16426d76eb31..e2cd82bf1e9510169735774784591074f30351d3)] --------- ## Summary - Design document for adding OAuth 2.0 authentication to the Rust ADBC driver covering both U2M (Authorization Code + PKCE) and M2M (Client Credentials) flows - Sprint plan breaking the implementation into 3 tasks: foundation + HTTP client changes, M2M provider, U2M provider - Uses the `oauth2` crate for protocol-level operations, unified `DatabricksHttpClient` with two-phase `OnceLock` init, and ODBC-aligned numeric config values (`AuthMech`/`Auth_Flow`) ## Key decisions and alternatives considered - **`oauth2` crate adoption** over hand-rolling OAuth protocol (eliminates ~200 lines of boilerplate, handles PKCE/token exchange/refresh) - **Unified HTTP client** (`DatabricksHttpClient` with `OnceLock`) over separate `reqwest::Client` for token calls (shared retry logic, connection pooling) - **ODBC-aligned numeric config** (`mechanism=0/11`, `flow=0/1/2`) over string-based or auto-detection (explicit, predictable, matches ODBC driver) - **Separate U2M/M2M providers** over single OAuthProvider (different flows, refresh strategies, caching needs) - **Separate token cache** (`~/.config/databricks-adbc/oauth/`) over sharing Python SDK cache (fragile cross-SDK compatibility) ## Areas needing specific review focus - Two-phase HTTP client initialization pattern (OnceLock for auth provider) — is this the right approach for breaking the circular dependency? - Token refresh state machine (FRESH/STALE/EXPIRED) — are the thresholds (40s expiry buffer, min(TTL*0.5, 20min) stale) appropriate? - Config option naming (`databricks.auth.mechanism`, `databricks.auth.flow`) — alignment with ODBC driver - Sprint plan task breakdown — is the scope realistic for 2 weeks? --- *Replaces #318 (closed — converted to stacked branch)* 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>

…-1.1-oauth-token-struct

…ache

…D: task-1.4-token-store

gopalldb

Nice overall

…task-1.5-oauth-module-setup

vikrantpuppala force-pushed the stack/pr-oauth-foundation branch from 62e94c8 to e0a0fdc Compare March 9, 2026 05:18

vikrantpuppala changed the title ~~Create OAuthToken struct with expiry and stale logic\n\nTask ID: task-1.1-oauth-token-struct~~ [PECOBLR-2158] feat(rust/oauth): Token primitives, OIDC discovery, cache and token store Mar 9, 2026

vikrantpuppala marked this pull request as ready for review March 9, 2026 05:42

vikrantpuppala force-pushed the stack/pr-oauth-foundation branch 2 times, most recently from fe4a411 to 1651567 Compare March 12, 2026 07:10

vikrantpuppala force-pushed the stack/pr-oauth-foundation branch 2 times, most recently from f0a904f to bd474c1 Compare March 12, 2026 13:49

gopalldb reviewed Mar 13, 2026

View reviewed changes

vikrantpuppala added 4 commits March 13, 2026 11:20

Create OAuthToken struct with expiry and stale logic\n\nTask ID: task…

4f5028d

…-1.1-oauth-token-struct

Implement OIDC endpoint discovery\n\nTask ID: task-1.2-oidc-discovery

e3837e5

Implement file-based token cache for U2M\n\nTask ID: task-1.3-token-c…

06636dd

…ache

Implement TokenStore with FRESH/STALE/EXPIRED state machine\n\nTask I…

4fa4768

…D: task-1.4-token-store

vikrantpuppala force-pushed the stack/pr-oauth-foundation branch 2 times, most recently from d601309 to 0968101 Compare March 13, 2026 12:21

vikrantpuppala changed the title ~~[PECOBLR-2158] feat(rust/oauth): Token primitives, OIDC discovery, cache and token store~~ feat(rust): OAuth token primitives, OIDC discovery, cache and token store Mar 13, 2026

vikrantpuppala force-pushed the stack/pr-oauth-foundation branch 2 times, most recently from 15ef8c4 to 7fd8031 Compare March 13, 2026 12:37

gopalldb approved these changes Mar 13, 2026

View reviewed changes

Set up oauth directory module and add Cargo dependencies\n\nTask ID: …

78b9ec8

…task-1.5-oauth-module-setup

vikrantpuppala force-pushed the stack/pr-oauth-foundation branch from 7fd8031 to 78b9ec8 Compare March 13, 2026 15:59

vikrantpuppala added the e2e-test Trigger E2E tests on this PR label Mar 13, 2026

vikrantpuppala enabled auto-merge March 13, 2026 16:31

vikrantpuppala removed the e2e-test Trigger E2E tests on this PR label Mar 13, 2026

vikrantpuppala disabled auto-merge March 13, 2026 16:32

vikrantpuppala merged commit 0eb5ee7 into adbc-drivers:main Mar 13, 2026
25 of 27 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(rust): OAuth token primitives, OIDC discovery, cache and token store#320

feat(rust): OAuth token primitives, OIDC discovery, cache and token store#320
vikrantpuppala merged 5 commits intoadbc-drivers:mainfrom
vikrantpuppala:stack/pr-oauth-foundation

vikrantpuppala commented Mar 8, 2026 •

edited

Loading

Uh oh!

vikrantpuppala commented Mar 12, 2026

Uh oh!

gopalldb Mar 13, 2026

Uh oh!

vikrantpuppala Mar 13, 2026

Uh oh!

gopalldb Mar 13, 2026

Uh oh!

vikrantpuppala Mar 13, 2026

Uh oh!

gopalldb Mar 13, 2026

Uh oh!

vikrantpuppala Mar 13, 2026

Uh oh!

gopalldb Mar 13, 2026

Uh oh!

vikrantpuppala Mar 13, 2026

Uh oh!

gopalldb Mar 13, 2026

Uh oh!

vikrantpuppala Mar 13, 2026

Uh oh!

gopalldb left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

vikrantpuppala commented Mar 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🥞 Stacked PR

Summary

Key files

Uh oh!

vikrantpuppala commented Mar 12, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gopalldb left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

vikrantpuppala commented Mar 8, 2026 •

edited

Loading