feat(rust): OAuth auth config, U2M and M2M providers

[vikrantpuppala]

🥞 Stacked PR

Use this link to review incremental changes.

• stack/pr-u2m-provider [Files changed]
  • stack/pr-integration-tests [Files changed]

---

Summary

Adds the AuthType configuration model and both OAuth authentication providers, wired into Database::new_connection():

Auth config (config.rs, database.rs)

• AuthType enum — AccessToken, OAuthM2m, OAuthU2m parsed from databricks.auth.type string option
• AuthConfig — collects all auth options (auth_type, client_id, client_secret, scopes, token_endpoint, redirect_port) with validation
• Database::set_option / new_connection() — flat match auth_type dispatches to the appropriate provider

OAuth providers

• CallbackServer (callback.rs) — lightweight TCP server on localhost for OAuth redirect callbacks with CSRF state validation
• AuthorizationCodeProvider (U2M) (u2m.rs) — authorization code + PKCE flow: OIDC discovery, browser launch, callback capture, token exchange, refresh with disk caching
• ClientCredentialsProvider (M2M) (m2m.rs) — client credentials grant with TokenStore-managed lifecycle

Configuration

databricks.auth.type = "access_token"    # Personal access token
databricks.auth.type = "oauth_m2m"       # Client credentials (service principal)
databricks.auth.type = "oauth_u2m"       # Authorization code + PKCE (browser)

Key files

• src/auth/config.rs — AuthType enum and AuthConfig validation
• src/auth/oauth/callback.rs — OAuth redirect callback server
• src/auth/oauth/u2m.rs — U2M provider
• src/auth/oauth/m2m.rs — M2M provider
• src/database.rs — option parsing and auth provider creation

This pull request was AI-assisted by Isaac.

[vikrantpuppala]

Range-diff: stack/pr-database-config (bba99fe -> 0919bed)

rust/src/auth/mod.rs

@@ -3,7 +3,7 @@
 +++ b/rust/src/auth/mod.rs
  pub mod pat;
  
- pub use config::{AuthConfig, AuthFlow, AuthMechanism};
+ pub use config::{AuthConfig, AuthType};
 +pub use oauth::{AuthorizationCodeProvider, ClientCredentialsProvider};
  pub use pat::PersonalAccessToken;
  

rust/src/database.rs

@@ -3,7 +3,7 @@
 +++ b/rust/src/database.rs
  //! Database implementation for the Databricks ADBC driver.
  
- use crate::auth::config::{AuthConfig, AuthFlow, AuthMechanism};
+ use crate::auth::config::{AuthConfig, AuthType};
 -use crate::auth::PersonalAccessToken;
 +use crate::auth::{AuthProvider, AuthorizationCodeProvider, PersonalAccessToken};
  use crate::client::{
@@ -31,11 +31,16 @@
  
          // Validate auth configuration
 -        self.auth_config
-+        let mechanism = self
++        let auth_type = self
 +            .auth_config
              .validate(&self.access_token)
              .map_err(|e| e.to_adbc())?;
  
+-        // Get access_token if needed (for PAT or TokenPassthrough)
++        // Get access_token if needed (for AccessToken type)
+         let access_token = self.access_token.as_ref();
+ 
+         // Create HTTP client (without auth provider - two-phase initialization)
          let http_client =
              Arc::new(DatabricksHttpClient::new(self.http_config.clone()).map_err(|e| e.to_adbc())?);
  
@@ -58,69 +63,52 @@
                  .to_adbc()
          })?;
  
-+        // Create auth provider based on mechanism and flow
-+        let auth_provider: Arc<dyn AuthProvider> = match mechanism {
-+            AuthMechanism::Pat => {
-+                // PAT mechanism - use PersonalAccessToken provider
++        // Create auth provider based on auth type
++        let auth_provider: Arc<dyn AuthProvider> = match auth_type {
++            AuthType::AccessToken => {
 +                Arc::new(PersonalAccessToken::new(
 +                    access_token
 +                        .expect("access_token should be validated above")
 +                        .clone(),
 +                ))
 +            }
-+            AuthMechanism::OAuth => {
-+                let flow = self
++            AuthType::OAuthM2m => {
++                // Client credentials flow not yet implemented
++                return Err(DatabricksErrorHelper::invalid_state()
++                    .message("Client credentials flow (M2M) not yet implemented")
++                    .to_adbc());
++            }
++            AuthType::OAuthU2m => {
++                // U2M flow - create AuthorizationCodeProvider
++                let client_id = self
 +                    .auth_config
-+                    .flow
-+                    .expect("flow should be validated above");
-+                match flow {
-+                    AuthFlow::TokenPassthrough => {
-+                        // Token passthrough - use PersonalAccessToken provider with the provided token
-+                        Arc::new(PersonalAccessToken::new(
-+                            access_token
-+                                .expect("access_token should be validated above")
-+                                .clone(),
-+                        ))
-+                    }
-+                    AuthFlow::ClientCredentials => {
-+                        // Client credentials flow not yet implemented
-+                        return Err(DatabricksErrorHelper::invalid_state()
-+                            .message("Client credentials flow (M2M) not yet implemented")
-+                            .to_adbc());
-+                    }
-+                    AuthFlow::Browser => {
-+                        // Browser flow (U2M) - create AuthorizationCodeProvider
-+                        let client_id = self
-+                            .auth_config
-+                            .client_id
-+                            .as_deref()
-+                            .unwrap_or("databricks-cli");
-+                        let scopes_str = self
-+                            .auth_config
-+                            .scopes
-+                            .as_deref()
-+                            .unwrap_or("all-apis offline_access");
-+                        let scopes: Vec<String> =
-+                            scopes_str.split_whitespace().map(String::from).collect();
-+                        let redirect_port = self.auth_config.redirect_port.unwrap_or(8020);
-+                        let callback_timeout = Duration::from_secs(120); // 2 minutes default
++                    .client_id
++                    .as_deref()
++                    .unwrap_or("databricks-cli");
++                let scopes_str = self
++                    .auth_config
++                    .scopes
++                    .as_deref()
++                    .unwrap_or("all-apis offline_access");
++                let scopes: Vec<String> =
++                    scopes_str.split_whitespace().map(String::from).collect();
++                let redirect_port = self.auth_config.redirect_port.unwrap_or(8020);
++                let callback_timeout = Duration::from_secs(120); // 2 minutes default
 +
-+                        // Create the provider (async operation - needs runtime)
-+                        let provider = runtime
-+                            .block_on(AuthorizationCodeProvider::new_with_full_config(
-+                                host,
-+                                client_id,
-+                                http_client.clone(),
-+                                scopes,
-+                                redirect_port,
-+                                callback_timeout,
-+                                self.auth_config.token_endpoint.clone(),
-+                            ))
-+                            .map_err(|e| e.to_adbc())?;
++                // Create the provider (async operation - needs runtime)
++                let provider = runtime
++                    .block_on(AuthorizationCodeProvider::new_with_full_config(
++                        host,
++                        client_id,
++                        http_client.clone(),
++                        scopes,
++                        redirect_port,
++                        callback_timeout,
++                        self.auth_config.token_endpoint.clone(),
++                    ))
++                    .map_err(|e| e.to_adbc())?;
 +
-+                        Arc::new(provider)
-+                    }
-+                }
++                Arc::new(provider)
 +            }
 +        };
 +
@@ -156,11 +144,10 @@
          );
      }
 +
-+    // Browser flow (U2M) configuration tests
++    // U2M (OAuth) configuration tests
 +
 +    #[test]
-+    fn test_database_browser_flow_config_with_defaults() {
-+        // Test that Browser flow accepts minimal configuration and uses defaults
++    fn test_database_oauth_u2m_config_with_defaults() {
 +        let mut db = Database::new();
 +        db.set_option(
 +            OptionDatabase::Uri,
@@ -173,48 +160,23 @@
 +        )
 +        .unwrap();
 +        db.set_option(
-+            OptionDatabase::Other("databricks.auth.mechanism".into()),
-+            OptionValue::Int(11), // OAuth
-+        )
-+        .unwrap();
-+        db.set_option(
-+            OptionDatabase::Other("databricks.auth.flow".into()),
-+            OptionValue::Int(2), // Browser
++            OptionDatabase::Other("databricks.auth.type".into()),
++            OptionValue::String("oauth_u2m".into()),
 +        )
 +        .unwrap();
 +
-+        // Verify config is valid - client_id, scopes, redirect_port should use defaults
-+        // We can't actually create a connection in unit tests (requires OIDC discovery),
-+        // but we can verify the config fields are accepted
-+        assert_eq!(db.auth_config.mechanism, Some(AuthMechanism::OAuth));
-+        assert_eq!(db.auth_config.flow, Some(AuthFlow::Browser));
++        assert_eq!(db.auth_config.auth_type, Some(AuthType::OAuthU2m));
 +        assert_eq!(db.auth_config.client_id, None); // Will use default "databricks-cli"
 +        assert_eq!(db.auth_config.scopes, None); // Will use default "all-apis offline_access"
 +        assert_eq!(db.auth_config.redirect_port, None); // Will use default 8020
 +    }
 +
 +    #[test]
-+    fn test_database_browser_flow_config_with_overrides() {
-+        // Test that Browser flow accepts custom configuration
++    fn test_database_oauth_u2m_config_with_overrides() {
 +        let mut db = Database::new();
 +        db.set_option(
-+            OptionDatabase::Uri,
-+            OptionValue::String("https://example.databricks.com".into()),
-+        )
-+        .unwrap();
-+        db.set_option(
-+            OptionDatabase::Other("databricks.warehouse_id".into()),
-+            OptionValue::String("test123".into()),
-+        )
-+        .unwrap();
-+        db.set_option(
-+            OptionDatabase::Other("databricks.auth.mechanism".into()),
-+            OptionValue::Int(11), // OAuth
-+        )
-+        .unwrap();
-+        db.set_option(
-+            OptionDatabase::Other("databricks.auth.flow".into()),
-+            OptionValue::Int(2), // Browser
++            OptionDatabase::Other("databricks.auth.type".into()),
++            OptionValue::String("oauth_u2m".into()),
 +        )
 +        .unwrap();
 +        db.set_option(
@@ -238,9 +200,7 @@
 +        )
 +        .unwrap();
 +
-+        // Verify config fields are set
-+        assert_eq!(db.auth_config.mechanism, Some(AuthMechanism::OAuth));
-+        assert_eq!(db.auth_config.flow, Some(AuthFlow::Browser));
++        assert_eq!(db.auth_config.auth_type, Some(AuthType::OAuthU2m));
 +        assert_eq!(
 +            db.auth_config.client_id,
 +            Some("custom-client-id".to_string())
@@ -257,36 +217,14 @@
 +    }
 +
 +    #[test]
-+    fn test_database_browser_flow_validation_passes() {
-+        // Test that Browser flow passes validation (no required fields beyond mechanism/flow)
-+        let mut db = Database::new();
-+        db.set_option(
-+            OptionDatabase::Uri,
-+            OptionValue::String("https://example.databricks.com".into()),
-+        )
-+        .unwrap();
-+        db.set_option(
-+            OptionDatabase::Other("databricks.warehouse_id".into()),
-+            OptionValue::String("test123".into()),
-+        )
-+        .unwrap();
-+        db.set_option(
-+            OptionDatabase::Other("databricks.auth.mechanism".into()),
-+            OptionValue::Int(11), // OAuth
-+        )
-+        .unwrap();
-+        db.set_option(
-+            OptionDatabase::Other("databricks.auth.flow".into()),
-+            OptionValue::Int(2), // Browser
-+        )
-+        .unwrap();
-+
-+        // Validation should pass - Browser flow doesn't require client_id, scopes, etc.
-+        // (they all have defaults)
-+        // We can't actually create a connection without mocking OIDC discovery,
-+        // but we can verify that validation logic doesn't reject minimal config
-+        // by checking the internal state
-+        assert_eq!(db.auth_config.mechanism, Some(AuthMechanism::OAuth));
-+        assert_eq!(db.auth_config.flow, Some(AuthFlow::Browser));
++    fn test_database_oauth_u2m_validation_passes() {
++        let config = AuthConfig {
++            auth_type: Some(AuthType::OAuthU2m),
++            ..Default::default()
++        };
++        // U2M doesn't require client_id, scopes, etc. (all have defaults)
++        let result = config.validate(&None);
++        assert!(result.is_ok());
++        assert_eq!(result.unwrap(), AuthType::OAuthU2m);
 +    }
  }
\ No newline at end of file

rust/tests/integration.rs

@@ -1,19 +1,10 @@
 diff --git a/rust/tests/integration.rs b/rust/tests/integration.rs
 --- a/rust/tests/integration.rs
 +++ b/rust/tests/integration.rs
- 
- #[test]
- fn test_auth_providers() {
--    use databricks_adbc::auth::{AuthProvider, OAuthCredentials, PersonalAccessToken};
-+    use databricks_adbc::auth::{AuthProvider, PersonalAccessToken};
- 
      // Test PAT
      let pat = PersonalAccessToken::new("test-token");
      assert_eq!(pat.get_auth_header().unwrap(), "Bearer test-token");
- 
--    // Test OAuth (not yet implemented)
--    let oauth = OAuthCredentials::new("client-id", "client-secret");
--    assert!(oauth.get_auth_header().is_err());
++
 +    // OAuth providers (ClientCredentialsProvider and AuthorizationCodeProvider)
 +    // require async initialization and HTTP client, so they are tested in unit tests
  }

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

_{Reproduce locally: git range-diff f23bf89..bba99fe 2e70c60..0919bed | Disable: git config gitstack.push-range-diff false}

[gopalldb]

should we also log this? Or is logging implicit in the error handling?

[vikrantpuppala]

Errors propagate to the caller via Result — the ADBC caller handles them. Adding logging here would duplicate the error path. We could add tracing::debug! for successful option sets, but that's noisy for normal operation. I'd leave it as-is.

[gopalldb]

but you are adding these providers, will this stitching be done in separate PR?

[vikrantpuppala]

Yes — M2M is fully wired up in PR #323 (the next PR in the stack). That PR adds ClientCredentialsProvider to the AuthType::OAuthM2m match arm in new_connection(), along with integration tests for both M2M and U2M.

[gopalldb]

shall we also implement Drop, which can abort the JoinHandle, if wait_for_code is not called due to some error?

[vikrantpuppala]

Good catch — added impl Drop for CallbackServer that aborts the server_handle if it's still running. Handles the case where an error occurs before wait_for_code is called.