Skip to content

fix(rust): Test suite fixes from final validation#324

Closed
vikrantpuppala wants to merge 10 commits intoadbc-drivers:mainfrom
vikrantpuppala:stack/pr-final-validation
Closed

fix(rust): Test suite fixes from final validation#324
vikrantpuppala wants to merge 10 commits intoadbc-drivers:mainfrom
vikrantpuppala:stack/pr-final-validation

Conversation

@vikrantpuppala
Copy link
Collaborator

@vikrantpuppala vikrantpuppala commented Mar 8, 2026

🥞 Stacked PR

Use this link to review incremental changes.


Summary

Minor fixes discovered while running the full test suite:

  • Fixed callback server HTML response formatting (content-length calculation)

Key files

  • src/auth/oauth/callback.rs — HTML response fix

This pull request was AI-assisted by Isaac.

scopes: Vec<String>,
}

#[allow(dead_code)] // Used in Phase 3 (U2M)
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

why this? i think this is used now? we should fix all such instances across

<p>You can close this tab and return to your application.</p>
</div>
</body>
</html>"#;
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

let's make this page simpler, no need for fancy fonts colours, check marks, etc.

<p>An error occurred during authentication. You can close this tab and try again.</p>
</div>
</body>
</html>"#;
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

use simple pages

// Use block_in_place to avoid blocking the runtime if we're in one,
// and get the handle to block_on the async operation
tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(async {
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

shouldn't this happen async? check all instances and see if we're handling this correctly

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

should we just remove this?

Comment on lines +35 to +88
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum AuthMechanism {
/// Personal access token (no OAuth). Config value: 0
Pat = 0,
/// OAuth 2.0 -- requires AuthFlow to select the specific flow. Config value: 11
OAuth = 11,
}

impl TryFrom<i64> for AuthMechanism {
type Error = crate::error::Error;

fn try_from(value: i64) -> std::result::Result<Self, Self::Error> {
match value {
0 => Ok(AuthMechanism::Pat),
11 => Ok(AuthMechanism::OAuth),
_ => Err(DatabricksErrorHelper::invalid_argument().message(format!(
"Invalid auth mechanism value: {}. Valid values are 0 (PAT) or 11 (OAuth)",
value
))),
}
}
}

/// OAuth authentication flow -- selects the specific OAuth grant type.
/// Config values match the ODBC driver's Auth_Flow numeric codes.
/// Only applicable when AuthMechanism is OAuth.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum AuthFlow {
/// Use a pre-obtained OAuth access token directly. Config value: 0
TokenPassthrough = 0,
/// M2M: client credentials grant for service principals. Config value: 1
ClientCredentials = 1,
/// U2M: browser-based authorization code + PKCE. Config value: 2
Browser = 2,
}

impl TryFrom<i64> for AuthFlow {
type Error = crate::error::Error;

fn try_from(value: i64) -> std::result::Result<Self, Self::Error> {
match value {
0 => Ok(AuthFlow::TokenPassthrough),
1 => Ok(AuthFlow::ClientCredentials),
2 => Ok(AuthFlow::Browser),
_ => Err(DatabricksErrorHelper::invalid_argument().message(format!(
"Invalid auth flow value: {}. Valid values are 0 (token passthrough), 1 (client credentials), or 2 (browser)",
value
))),
}
}
}

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

is this the right place to define this? seems like database is becoming bloated?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

overall would be great if we can abstract out the auth related stuff in this file so that it is less bloated

rust/CLAUDE.md Outdated
- Test names: `test_<function>_<scenario>`
- E2E tests that require real Databricks connection should be marked with `#[ignore]`

#### Running E2E OAuth Tests
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

don't think we need this section

@vikrantpuppala vikrantpuppala force-pushed the stack/pr-final-validation branch from 0e69f4a to 13f8c28 Compare March 9, 2026 05:18
@vikrantpuppala vikrantpuppala changed the title Run full test suite and fix any failures\n\nTask ID: task-5.1-final-validation [PECOBLR-2089] fix(rust/oauth): Test suite fixes from final validation Mar 9, 2026
@vikrantpuppala vikrantpuppala force-pushed the stack/pr-final-validation branch 2 times, most recently from f87df69 to b8141fb Compare March 12, 2026 07:10
@vikrantpuppala
Copy link
Collaborator Author

Range-diff: stack/pr-integration-tests (f87df69 -> b8141fb)
rust/src/auth/oauth/callback.rs
@@ -51,19 +51,4 @@
 -                "GET /callback?error=access_denied&error_description=User%20denied%20consent HTTP/1.1\r\nHost: localhost:{}\r\n\r\n",
 +                "GET /?error=access_denied&error_description=User%20denied%20consent HTTP/1.1\r\nHost: localhost:{}\r\n\r\n",
                  port
-             );
- 
-             uri
-         );
-         // Verify the port is a valid non-zero number
--        let port: u16 = uri.strip_prefix("http://localhost:").unwrap().strip_suffix("/callback").unwrap_or_else(|| uri.strip_prefix("http://localhost:").unwrap()).parse().unwrap();
-+        let port: u16 = uri
-+            .strip_prefix("http://localhost:")
-+            .unwrap()
-+            .strip_suffix("/callback")
-+            .unwrap_or_else(|| uri.strip_prefix("http://localhost:").unwrap())
-+            .parse()
-+            .unwrap();
-         assert!(port > 0, "Expected a non-zero port, got: {}", port);
-     }
- }
\ No newline at end of file
+             );
\ No newline at end of file
rust/src/auth/config.rs
@@ -1,16 +0,0 @@
-diff --git a/rust/src/auth/config.rs b/rust/src/auth/config.rs
---- a/rust/src/auth/config.rs
-+++ b/rust/src/auth/config.rs
-         match mechanism {
-             AuthMechanism::Pat => {
-                 if access_token.is_none() {
--                    return Err(DatabricksErrorHelper::invalid_argument()
--                        .message(
--                            "databricks.access_token is required when auth mechanism is 0 (PAT)",
--                        ));
-+                    return Err(DatabricksErrorHelper::invalid_argument().message(
-+                        "databricks.access_token is required when auth mechanism is 0 (PAT)",
-+                    ));
-                 }
-             }
-             AuthMechanism::OAuth => {
\ No newline at end of file
rust/src/auth/oauth/cache.rs
@@ -1,17 +0,0 @@
-diff --git a/rust/src/auth/oauth/cache.rs b/rust/src/auth/oauth/cache.rs
---- a/rust/src/auth/oauth/cache.rs
-+++ b/rust/src/auth/oauth/cache.rs
-     scopes: Vec<String>,
- }
- 
--
- impl CacheKey {
-     /// Creates a new cache key from the given parameters.
-     fn new(host: &str, client_id: &str, scopes: &[String]) -> Self {
- /// Cache I/O errors are logged as warnings and never block authentication.
- pub struct TokenCache;
- 
--
- impl TokenCache {
-     /// Returns the cache directory path.
-     /// Location: `~/.config/databricks-adbc/oauth/`
\ No newline at end of file
rust/src/auth/oauth/m2m.rs
@@ -1,19 +0,0 @@
-diff --git a/rust/src/auth/oauth/m2m.rs b/rust/src/auth/oauth/m2m.rs
---- a/rust/src/auth/oauth/m2m.rs
-+++ b/rust/src/auth/oauth/m2m.rs
-         token_endpoint_override: Option<String>,
-     ) -> Result<Self> {
-         // Discover OIDC endpoints or use override
--        let (token_endpoint, auth_endpoint) = if let Some(token_endpoint) = token_endpoint_override {
-+        let (token_endpoint, auth_endpoint) = if let Some(token_endpoint) = token_endpoint_override
-+        {
-             // Use override for token endpoint, still discover auth endpoint
-             // (auth endpoint is not typically overridden for M2M since it's not used)
-             let endpoints = OidcEndpoints::discover(host, &http_client).await?;
-             scopes,
-         })
-     }
--
- }
- 
- impl AuthProvider for ClientCredentialsProvider {
\ No newline at end of file
rust/src/auth/oauth/token_store.rs
@@ -1,10 +0,0 @@
-diff --git a/rust/src/auth/oauth/token_store.rs b/rust/src/auth/oauth/token_store.rs
---- a/rust/src/auth/oauth/token_store.rs
-+++ b/rust/src/auth/oauth/token_store.rs
-     refreshing: Arc<AtomicBool>,
- }
- 
--
- impl TokenStore {
-     /// Creates a new empty token store.
-     pub fn new() -> Self {
\ No newline at end of file
rust/src/database.rs
@@ -1,141 +0,0 @@
-diff --git a/rust/src/database.rs b/rust/src/database.rs
---- a/rust/src/database.rs
-+++ b/rust/src/database.rs
-                             .to_adbc()
-                     })
-                 }
--                "databricks.auth.scopes" => {
--                    self.auth_config.scopes.clone().ok_or_else(|| {
--                        DatabricksErrorHelper::invalid_state()
--                            .message("option 'databricks.auth.scopes' is not set")
--                            .to_adbc()
--                    })
--                }
-+                "databricks.auth.scopes" => self.auth_config.scopes.clone().ok_or_else(|| {
-+                    DatabricksErrorHelper::invalid_state()
-+                        .message("option 'databricks.auth.scopes' is not set")
-+                        .to_adbc()
-+                }),
-                 "databricks.auth.token_endpoint" => {
-                     self.auth_config.token_endpoint.clone().ok_or_else(|| {
-                         DatabricksErrorHelper::invalid_state()
-         })?;
- 
-         // Validate auth configuration
--        let mechanism = self.auth_config
-+        let mechanism = self
-+            .auth_config
-             .validate(&self.access_token)
-             .map_err(|e| e.to_adbc())?;
- 
-                 ))
-             }
-             AuthMechanism::OAuth => {
--                let flow = self.auth_config.flow.expect("flow should be validated above");
-+                let flow = self
-+                    .auth_config
-+                    .flow
-+                    .expect("flow should be validated above");
-                 match flow {
-                     AuthFlow::TokenPassthrough => {
-                         // Token passthrough - use PersonalAccessToken provider with the provided token
-                     }
-                     AuthFlow::ClientCredentials => {
-                         // Client credentials flow (M2M) - create ClientCredentialsProvider
--                        let client_id = self.auth_config.client_id.as_ref()
-+                        let client_id = self
-+                            .auth_config
-+                            .client_id
-+                            .as_ref()
-                             .expect("client_id should be validated above");
--                        let client_secret = self.auth_config.client_secret.as_ref()
-+                        let client_secret = self
-+                            .auth_config
-+                            .client_secret
-+                            .as_ref()
-                             .expect("client_secret should be validated above");
- 
-                         // Default scope for M2M is "all-apis" (no offline_access since M2M has no refresh token)
-                         let scopes_str = self.auth_config.scopes.as_deref().unwrap_or("all-apis");
--                        let scopes: Vec<String> = scopes_str.split_whitespace().map(String::from).collect();
-+                        let scopes: Vec<String> =
-+                            scopes_str.split_whitespace().map(String::from).collect();
- 
-                         // Create the provider (async operation - needs runtime)
-                         let provider = runtime
--                            .block_on(crate::auth::ClientCredentialsProvider::new_with_full_config(
--                                host,
--                                client_id,
--                                client_secret,
--                                http_client.clone(),
--                                scopes,
--                                self.auth_config.token_endpoint.clone(),
--                            ))
-+                            .block_on(
-+                                crate::auth::ClientCredentialsProvider::new_with_full_config(
-+                                    host,
-+                                    client_id,
-+                                    client_secret,
-+                                    http_client.clone(),
-+                                    scopes,
-+                                    self.auth_config.token_endpoint.clone(),
-+                                ),
-+                            )
-                             .map_err(|e| e.to_adbc())?;
- 
-                         Arc::new(provider)
-         )
-         .unwrap();
- 
--        assert_eq!(db.auth_config.client_secret, Some("test-secret".to_string()));
-+        assert_eq!(
-+            db.auth_config.client_secret,
-+            Some("test-secret".to_string())
-+        );
-         assert_eq!(
-             db.get_option_string(OptionDatabase::Other(
-                 "databricks.auth.client_secret".into()
-         )
-         .unwrap();
- 
--        assert_eq!(db.auth_config.scopes, Some("all-apis offline_access".to_string()));
-+        assert_eq!(
-+            db.auth_config.scopes,
-+            Some("all-apis offline_access".to_string())
-+        );
-         assert_eq!(
-             db.get_option_string(OptionDatabase::Other("databricks.auth.scopes".into()))
-                 .unwrap(),
-         .unwrap();
- 
-         // Verify config is stored correctly
--        assert_eq!(db.auth_mechanism, Some(AuthMechanism::OAuth));
--        assert_eq!(db.auth_flow, Some(AuthFlow::ClientCredentials));
--        assert_eq!(db.auth_client_id, Some("test-client-id".to_string()));
--        assert_eq!(db.auth_client_secret, Some("test-secret".to_string()));
-+        assert_eq!(db.auth_config.mechanism, Some(AuthMechanism::OAuth));
-+        assert_eq!(db.auth_config.flow, Some(AuthFlow::ClientCredentials));
-+        assert_eq!(db.auth_config.client_id, Some("test-client-id".to_string()));
-+        assert_eq!(
-+            db.auth_config.client_secret,
-+            Some("test-secret".to_string())
-+        );
-     }
- 
-     #[test]
- 
-         // Verify custom scopes are stored
-         assert_eq!(
--            db.auth_scopes,
-+            db.auth_config.scopes,
-             Some("custom-scope-1 custom-scope-2".to_string())
-         );
-     }
- 
-         // Verify token_endpoint override is stored
-         assert_eq!(
--            db.auth_token_endpoint,
-+            db.auth_config.token_endpoint,
-             Some("https://custom.endpoint/token".to_string())
-         );
-     }
\ No newline at end of file
rust/tests/oauth_e2e.rs
@@ -1,23 +0,0 @@
-diff --git a/rust/tests/oauth_e2e.rs b/rust/tests/oauth_e2e.rs
---- a/rust/tests/oauth_e2e.rs
-+++ b/rust/tests/oauth_e2e.rs
- 
-     // Create driver and database
-     let mut driver = Driver::new();
--    let mut database = driver
--        .new_database()
--        .expect("Failed to create database");
-+    let mut database = driver.new_database().expect("Failed to create database");
- 
-     // Configure database with M2M OAuth
-     database
- 
-     // Create driver and database
-     let mut driver = Driver::new();
--    let mut database = driver
--        .new_database()
--        .expect("Failed to create database");
-+    let mut database = driver.new_database().expect("Failed to create database");
- 
-     // Configure database with U2M OAuth
-     database
\ No newline at end of file

Reproduce locally: git range-diff 4e111ed..f87df69 f3e647c..b8141fb | Disable: git config gitstack.push-range-diff false

@vikrantpuppala vikrantpuppala force-pushed the stack/pr-final-validation branch 2 times, most recently from 7dfee16 to e2cd82b Compare March 12, 2026 13:49
vikrantpuppala added a commit that referenced this pull request Mar 13, 2026
## 🥞 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>
@vikrantpuppala vikrantpuppala force-pushed the stack/pr-final-validation branch 2 times, most recently from 23d69f3 to 01cfd62 Compare March 13, 2026 12:21
@vikrantpuppala vikrantpuppala changed the title [PECOBLR-2089] fix(rust/oauth): Test suite fixes from final validation fix(rust): Test suite fixes from final validation Mar 13, 2026
@vikrantpuppala vikrantpuppala force-pushed the stack/pr-final-validation branch 3 times, most recently from 85d9339 to 2d6ccb0 Compare March 13, 2026 15:59
vikrantpuppala added a commit that referenced this pull request Mar 13, 2026
…tore (#320)

## 🥞 Stacked PR
Use this
[link](https://github.com/adbc-drivers/databricks/pull/320/files) to
review incremental changes.
-
[**stack/pr-oauth-foundation**](#320)
[[Files
changed](https://github.com/adbc-drivers/databricks/pull/320/files)]
-
[stack/pr-database-config](#321)
[[Files
changed](https://github.com/adbc-drivers/databricks/pull/321/files/78b9ec88459f895c76bd1aea99fcb47e5eb94893..164ada04d14660306c7e44dd3d52a7943050aa27)]
-
[stack/pr-u2m-provider](#322)
[[Files
changed](https://github.com/adbc-drivers/databricks/pull/322/files/164ada04d14660306c7e44dd3d52a7943050aa27..abc00ced51d89f1a652f78209f692775eba05e73)]
-
[stack/pr-integration-tests](#323)
[[Files
changed](https://github.com/adbc-drivers/databricks/pull/323/files/abc00ced51d89f1a652f78209f692775eba05e73..75b18d6c594eeba89a30450152d6d6f672239614)]
-
[stack/pr-final-validation](#324)
[[Files
changed](https://github.com/adbc-drivers/databricks/pull/324/files/75b18d6c594eeba89a30450152d6d6f672239614..2d6ccb09e121015aa6a0da6e992529a686bb0f04)]

---------
## Summary

Adds the core OAuth token infrastructure used by both U2M and M2M flows:

- **`OAuthToken`** — token struct with expiry tracking, stale detection
(40s buffer / 50% TTL), and serde support
- **OIDC discovery** — fetches `authorization_endpoint` and
`token_endpoint` from `/.well-known/oauth-authorization-server`
- **`TokenCache`** — file-based persistence at
`~/.config/databricks-adbc/oauth/` with SHA-256 hashed filenames and
`0o600` permissions
- **`TokenStore`** — thread-safe token lifecycle (Empty → Fresh → Stale
→ Expired) with coordinated refresh via `RwLock` + `AtomicBool`
- **Cargo dependencies** — `oauth2`, `sha2`, `dirs`, `serde`, `open`
crates
- **`DatabricksHttpClient`** — extended with `OnceLock`-based auth
provider and `inner()` accessor for the `oauth2` crate

### Key files
- `src/auth/oauth/token.rs` — `OAuthToken` struct
- `src/auth/oauth/oidc.rs` — OIDC endpoint discovery
- `src/auth/oauth/cache.rs` — file-based token cache
- `src/auth/oauth/token_store.rs` — token lifecycle state machine
- `src/client/http.rs` — HTTP client auth provider integration
@vikrantpuppala vikrantpuppala force-pushed the stack/pr-final-validation branch from 2d6ccb0 to 04d7b09 Compare March 13, 2026 16:33
@vikrantpuppala vikrantpuppala force-pushed the stack/pr-final-validation branch from 04d7b09 to 720434f Compare March 13, 2026 17:07
@vikrantpuppala
Copy link
Collaborator Author

Folded into PR #323 to reduce stack size.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant