8000 feat(rust): Auth config enums and Database option parsing by vikrantpuppala · Pull Request #321 · adbc-drivers/databricks · GitHub
[go: up one dir, main page]

Skip to content

feat(rust): Auth config enums and Database option parsing#321

Closed
vikrantpuppala wants to merge 3 commits intoadbc-drivers:mainfrom
vikrantpuppala:stack/pr-database-config
Closed

feat(rust): Auth config enums and Database option parsing#321
vikrantpuppala wants to merge 3 commits intoadbc-drivers:mainfrom
vikrantpuppala:stack/pr-database-config

Conversation

@vikrantpuppala
Copy link
Collaborator
@vikrantpuppala vikrantpuppala commented Mar 8, 2026

🥞 Stacked PR

Use this link to review incremental changes.


Summary

Replaces the ODBC-style numeric AuthMech/Auth_Flow config with a single string-based databricks.auth.type option:

  • AuthType enum (config.rs) — AccessToken, OAuthM2m, OAuthU2m parsed from string values (access_token, oauth_m2m, oauth_u2m)
  • AuthConfig — simplified to auth_type + credential fields (no more mechanism/flow)
  • Database::set_option — handles databricks.auth.type as a string; removed numeric mechanism/flow parsing
  • Database::new_connection — flat match auth_type instead of nested mechanism/flow matching
  • Validation — each auth type directly checks its required fields

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.rsAuthType enum and AuthConfig validation
  • src/auth/mod.rs — updated re-exports
  • src/database.rs — option parsing and new_connection() auth provider creation

This pull request was AI-assisted by Isaac.

@vikrantpuppala vikrantpuppala force-pushed the stack/pr-database-config branch from 094741a to f49c83e Compare March 9, 2026 05:18
@vikrantpuppala vikrantpuppala changed the title Add AuthMechanism and AuthFlow enums to database.rs\n\nTask ID: task-2.1-auth-enums [PECOBLR-2165] feat(rust/oauth): Auth config enums and Database option parsing Mar 9, 2026
@vikrantpuppala vikrantpuppala marked this pull request as ready for review March 9, 2026 05:42
@vikrantpuppala vikrantpuppala force-pushed the stack/pr-database-config branch 3 times, most recently from f23bf89 to 2e70c60 Compare March 12, 2026 07:34
@vikrantpuppala
Copy link
Collaborator Author
Range-diff: stack/pr-oauth-foundation (f23bf89 -> 2e70c60)
rust/src/auth/config.rs
@@ -24,58 +24,39 @@
 +use crate::error::DatabricksErrorHelper;
 +use driverbase::error::ErrorHelper;
 +
-+/// Authentication mechanism -- top-level selector.
-+/// Config values match the ODBC driver's AuthMech numeric codes.
++/// Authentication type -- single selector for the authentication method.
 +#[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,
++pub enum AuthType {
++    /// Personal access token.
++    AccessToken,
++    /// M2M: client credentials grant for service principals.
++    OAuthM2m,
++    /// U2M: browser-based authorization code + PKCE.
++    OAuthU2m,
 +}
 +
-+impl TryFrom<i64> for AuthMechanism {
++impl TryFrom<&str> for AuthType {
 +    type Error = crate::error::Error;
 +
-+    fn try_from(value: i64) -> std::result::Result<Self, Self::Error> {
++    fn try_from(value: &str) -> std::result::Result<Self, Self::Error> {
 +        match value {
-+            0 => Ok(AuthMechanism::Pat),
-+            11 => Ok(AuthMechanism::OAuth),
++            "access_token" => Ok(AuthType::AccessToken),
++            "oauth_m2m" => Ok(AuthType::OAuthM2m),
++            "oauth_u2m" => Ok(AuthType::OAuthU2m),
 +            _ => Err(DatabricksErrorHelper::invalid_argument().message(format!(
-+                "Invalid auth mechanism value: {}. Valid values are 0 (PAT) or 11 (OAuth)",
++                "Invalid auth type: '{}'. Valid values: 'access_token', 'oauth_m2m', 'oauth_u2m'",
 +                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
-+            ))),
++impl std::fmt::Display for AuthType {
++    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
++        match self {
++            AuthType::AccessToken => write!(f, "access_token"),
++            AuthType::OAuthM2m => write!(f, "oauth_m2m"),
++            AuthType::OAuthU2m => write!(f, "oauth_u2m"),
 +        }
 +    }
 +}
@@ -87,8 +68,7 @@
 +/// create the appropriate `AuthProvider`.
 +#[derive(Debug, Default, Clone)]
 +pub struct AuthConfig {
-+    pub mechanism: Option<AuthMechanism>,
-+    pub flow: Option<AuthFlow>,
++    pub auth_type: Option<AuthType>,
 +    pub client_id: Option<String>,
 +    pub client_secret: Option<String>,
 +    pub scopes: Option<String>,
@@ -97,60 +77,47 @@
 +}
 +
 +impl AuthConfig {
-+    /// Validates the auth configuration and returns the mechanism.
++    /// Validates the auth configuration and returns the auth type.
 +    ///
 +    /// This checks that:
-+    /// - A mechanism is specified
-+    /// - Required fields are present for the chosen mechanism/flow
-+    pub fn validate(&self, access_token: &Option<String>) -> crate::error::Result<AuthMechanism> {
-+        let mechanism = self.mechanism.ok_or_else(|| {
-+            DatabricksErrorHelper::invalid_argument()
-+                .message("databricks.auth.mechanism is required. Valid values: 0 (PAT), 11 (OAuth)")
++    /// - An auth type is specified
++    /// - Required fields are present for the chosen auth type
++    pub fn validate(&self, access_token: &Option<String>) -> crate::error::Result<AuthType> {
++        let auth_type = self.auth_type.ok_or_else(|| {
++            DatabricksErrorHelper::invalid_argument().message(
++                "databricks.auth.type is required. Valid values: 'access_token', 'oauth_m2m', 'oauth_u2m'",
++            )
 +        })?;
 +
-+        match mechanism {
-+            AuthMechanism::Pat => {
++        match auth_type {
++            AuthType::AccessToken => {
 +                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 type is 'access_token'",
++                    ));
 +                }
 +            }
-+            AuthMechanism::OAuth => {
-+                let flow = self.flow.ok_or_else(|| {
-+                    DatabricksErrorHelper::invalid_argument()
-+                        .message("databricks.auth.flow is required when auth mechanism is 11 (OAuth). Valid values: 0 (token passthrough), 1 (client credentials), 2 (browser)")
-+                })?;
-+
-+                match flow {
-+                    AuthFlow::TokenPassthrough => {
-+                        if access_token.is_none() {
-+                            return Err(DatabricksErrorHelper::invalid_argument()
-+                                .message("databricks.access_token is required when auth flow is 0 (token passthrough)"));
-+                        }
-+                    }
-+                    AuthFlow::ClientCredentials => {
-+                        if self.client_id.is_none() {
-+                            return Err(DatabricksErrorHelper::invalid_argument()
-+                                .message("databricks.auth.client_id is required when auth flow is 1 (client credentials)"));
-+                        }
-+                        if self.client_secret.is_none() {
-+                            return Err(DatabricksErrorHelper::invalid_argument()
-+                                .message("databricks.auth.client_secret is required when auth flow is 1 (client credentials)"));
-+                        }
-+                    }
-+                    AuthFlow::Browser => {
-+                        // Browser flow has no required fields - all parameters have defaults:
-+                        // - client_id defaults to "databricks-cli"
-+                        // - scopes defaults to "all-apis offline_access"
-+                        // - redirect_port defaults to 8020
-+                    }
++            AuthType::OAuthM2m => {
++                if self.client_id.is_none() {
++                    return Err(DatabricksErrorHelper::invalid_argument().message(
++                        "databricks.auth.client_id is required when auth type is 'oauth_m2m'",
++                    ));
++                }
++                if self.client_secret.is_none() {
++                    return Err(DatabricksErrorHelper::invalid_argument().message(
++                        "databricks.auth.client_secret is required when auth type is 'oauth_m2m'",
++                    ));
 +                }
 +            }
++            AuthType::OAuthU2m => {
++                // U2M flow has no required fields - all parameters have defaults:
++                // - client_id defaults to "databricks-cli"
++                // - scopes defaults to "all-apis offline_access"
++                // - redirect_port defaults to 8020
++            }
 +        }
 +
-+        Ok(mechanism)
++        Ok(auth_type)
 +    }
 +}
 +
@@ -159,61 +126,41 @@
 +    use super::*;
 +
 +    #[test]
-+    fn test_auth_mechanism_valid() {
-+        let mechanism = AuthMechanism::try_from(0).unwrap();
-+        assert_eq!(mechanism, AuthMechanism::Pat);
-+        assert_eq!(mechanism as u8, 0);
-+
-+        let mechanism = AuthMechanism::try_from(11).unwrap();
-+        assert_eq!(mechanism, AuthMechanism::OAuth);
-+        assert_eq!(mechanism as u8, 11);
-+    }
-+
-+    #[test]
-+    fn test_auth_mechanism_invalid() {
-+        assert!(AuthMechanism::try_from(1).is_err());
-+        assert!(AuthMechanism::try_from(10).is_err());
-+        assert!(AuthMechanism::try_from(12).is_err());
-+        assert!(AuthMechanism::try_from(-1).is_err());
-+        assert!(AuthMechanism::try_from(100).is_err());
++    fn test_auth_type_valid() {
++        assert_eq!(AuthType::try_from("access_token").unwrap(), AuthType::AccessToken);
++        assert_eq!(AuthType::try_from("oauth_m2m").unwrap(), AuthType::OAuthM2m);
++        assert_eq!(AuthType::try_from("oauth_u2m").unwrap(), AuthType::OAuthU2m);
 +    }
 +
 +    #[test]
-+    fn test_auth_flow_valid() {
-+        let flow = AuthFlow::try_from(0).unwrap();
-+        assert_eq!(flow, AuthFlow::TokenPassthrough);
-+        assert_eq!(flow as u8, 0);
-+
-+        let flow = AuthFlow::try_from(1).unwrap();
-+        assert_eq!(flow, AuthFlow::ClientCredentials);
-+        assert_eq!(flow as u8, 1);
-+
-+        let flow = AuthFlow::try_from(2).unwrap();
-+        assert_eq!(flow, AuthFlow::Browser);
-+        assert_eq!(flow as u8, 2);
++    fn test_auth_type_invalid() {
++        assert!(AuthType::try_from("pat").is_err());
++        assert!(AuthType::try_from("oauth").is_err());
++        assert!(AuthType::try_from("0").is_err());
++        assert!(AuthType::try_from("11").is_err());
++        assert!(AuthType::try_from("").is_err());
 +    }
 +
 +    #[test]
-+    fn test_auth_flow_invalid() {
-+        assert!(AuthFlow::try_from(3).is_err());
-+        assert!(AuthFlow::try_from(-1).is_err());
-+        assert!(AuthFlow::try_from(10).is_err());
-+        assert!(AuthFlow::try_from(100).is_err());
++    fn test_auth_type_display() {
++        assert_eq!(AuthType::AccessToken.to_string(), "access_token");
++        assert_eq!(AuthType::OAuthM2m.to_string(), "oauth_m2m");
++        assert_eq!(AuthType::OAuthU2m.to_string(), "oauth_u2m");
 +    }
 +
 +    #[test]
-+    fn test_validate_missing_mechanism() {
++    fn test_validate_missing_auth_type() {
 +        let config = AuthConfig::default();
 +        let result = config.validate(&None);
 +        assert!(result.is_err());
 +        let err_msg = format!("{:?}", result.unwrap_err());
-+        assert!(err_msg.contains("databricks.auth.mechanism is required"));
++        assert!(err_msg.contains("databricks.auth.type is required"));
 +    }
 +
 +    #[test]
-+    fn test_validate_pat_missing_token() {
++    fn test_validate_access_token_missing_token() {
 +        let config = AuthConfig {
-+            mechanism: Some(AuthMechanism::Pat),
++            auth_type: Some(AuthType::AccessToken),
 +            ..Default::default()
 +        };
 +        let result = config.validate(&None);
@@ -223,33 +170,20 @@
 +    }
 +
 +    #[test]
-+    fn test_validate_pat_with_token() {
++    fn test_validate_access_token_with_token() {
 +        let config = AuthConfig {
-+            mechanism: Some(AuthMechanism::Pat),
++            auth_type: Some(AuthType::AccessToken),
 +            ..Default::default()
 +        };
 +        let result = config.validate(&Some("token".to_string()));
 +        assert!(result.is_ok());
-+        assert_eq!(result.unwrap(), AuthMechanism::Pat);
++        assert_eq!(result.unwrap(), AuthType::AccessToken);
 +    }
 +
 +    #[test]
-+    fn test_validate_oauth_missing_flow() {
-+        let config = AuthConfig {
-+            mechanism: Some(AuthMechanism::OAuth),
-+            ..Default::default()
-+        };
-+        let result = config.validate(&None);
-+        assert!(result.is_err());
-+        let err_msg = format!("{:?}", result.unwrap_err());
-+        assert!(err_msg.contains("databricks.auth.flow is required"));
-+    }
-+
-+    #[test]
-+    fn test_validate_client_credentials_missing_client_id() {
++    fn test_validate_oauth_m2m_missing_client_id() {
 +        let config = AuthConfig {
-+            mechanism: Some(AuthMechanism::OAuth),
-+            flow: Some(AuthFlow::ClientCredentials),
++            auth_type: Some(AuthType::OAuthM2m),
 +            ..Default::default()
 +        };
 +        let result = config.validate(&None);
@@ -259,10 +193,9 @@
 +    }
 +
 +    #[test]
-+    fn test_validate_client_credentials_missing_secret() {
++    fn test_validate_oauth_m2m_missing_secret() {
 +        let config = AuthConfig {
-+            mechanism: Some(AuthMechanism::OAuth),
-+            flow: Some(AuthFlow::ClientCredentials),
++            auth_type: Some(AuthType::OAuthM2m),
 +            client_id: Some("id".to_string()),
 +            ..Default::default()
 +        };
@@ -273,37 +206,26 @@
 +    }
 +
 +    #[test]
-+    fn test_validate_client_credentials_valid() {
++    fn test_validate_oauth_m2m_valid() {
 +        let config = AuthConfig {
-+            mechanism: Some(AuthMechanism::OAuth),
-+            flow: Some(AuthFlow::ClientCredentials),
++            auth_type: Some(AuthType::OAuthM2m),
 +            client_id: Some("id".to_string()),
 +            client_secret: Some("secret".to_string()),
 +            ..Default::default()
 +        };
 +        let result = config.validate(&None);
 +        assert!(result.is_ok());
++        assert_eq!(result.unwrap(), AuthType::OAuthM2m);
 +    }
 +
 +    #[test]
-+    fn test_validate_token_passthrough_missing_token() {
++    fn test_validate_oauth_u2m_no_required_fields() {
 +        let config = AuthConfig {
-+            mechanism: Some(AuthMechanism::OAuth),
-+            flow: Some(AuthFlow::TokenPassthrough),
-+            ..Default::default()
-+        };
-+        let result = config.validate(&None);
-+        assert!(result.is_err());
-+    }
-+
-+    #[test]
-+    fn test_validate_browser_no_required_fields() {
-+        let config = AuthConfig {
-+            mechanism: Some(AuthMechanism::OAuth),
-+            flow: Some(AuthFlow::Browser),
++            auth_type: Some(AuthType::OAuthU2m),
 +            ..Default::default()
 +        };
 +        let result = config.validate(&None);
 +        assert!(result.is_ok());
++        assert_eq!(result.unwrap(), AuthType::OAuthU2m);
 +    }
 +}
\ No newline at end of file
rust/src/auth/mod.rs
@@ -8,7 +8,7 @@
  pub mod oauth;
  pub mod pat;
  
-+pub use config::{AuthConfig, AuthFlow, AuthMechanism};
++pub use config::{AuthConfig, AuthType};
  pub use pat::PersonalAccessToken;
  
  use crate::error::Result;
\ No newline at end of file
rust/src/database.rs
@@ -4,7 +4,7 @@
  
  //! 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::client::{
      DatabricksClient, DatabricksClientConfig, DatabricksHttpClient, HttpClientConfig, SeaClient,
@@ -20,18 +20,10 @@
                      }
                  }
  
-+                // OAuth configuration options
-+                "databricks.auth.mechanism" => {
-+                    if let Some(v) = Self::parse_int_option(&value) {
-+                        self.auth_config.mechanism = Some(AuthMechanism::try_from(v)?);
-+                        Ok(())
-+                    } else {
-+                        Err(DatabricksErrorHelper::set_invalid_option(&key, &value).to_adbc())
-+                    }
-+                }
-+                "databricks.auth.flow" => {
-+                    if let Some(v) = Self::parse_int_option(&value) {
-+                        self.auth_config.flow = Some(AuthFlow::try_from(v)?);
++                // Authentication configuration options
++                "databricks.auth.type" => {
++                    if let OptionValue::String(v) = value {
++                        self.auth_config.auth_type = Some(AuthType::try_from(v.as_str())?);
 +                        Ok(())
 +                    } else {
 +                        Err(DatabricksErrorHelper::set_invalid_option(&key, &value).to_adbc())
@@ -92,6 +84,15 @@
                          .message("option 'databricks.log_file' is not set")
                          .to_adbc()
                  }),
++                "databricks.auth.type" => self
++                    .auth_config
++                    .auth_type
++                    .ok_or_else(|| {
++                        DatabricksErrorHelper::invalid_state()
++                            .message("option 'databricks.auth.type' is not set")
++                            .to_adbc()
++                    })
++                    .map(|t| t.to_string()),
 +                "databricks.auth.client_id" => {
 +                    self.auth_config.client_id.clone().ok_or_else(|| {
 +                        DatabricksErrorHelper::invalid_state()
@@ -126,24 +127,6 @@
                  "databricks.cloudfetch.max_retries" => {
                      Ok(self.cloudfetch_config.max_retries as i64)
                  }
-+                "databricks.auth.mechanism" => self
-+                    .auth_config
-+                    .mechanism
-+                    .ok_or_else(|| {
-+                        DatabricksErrorHelper::invalid_state()
-+                            .message("option 'databricks.auth.mechanism' is not set")
-+                            .to_adbc()
-+                    })
-+                    .map(|m| m as i64),
-+                "databricks.auth.flow" => self
-+                    .auth_config
-+                    .flow
-+                    .ok_or_else(|| {
-+                        DatabricksErrorHelper::invalid_state()
-+                            .message("option 'databricks.auth.flow' is not set")
-+                            .to_adbc()
-+                    })
-+                    .map(|f| f as i64),
 +                "databricks.auth.redirect_port" => self
 +                    .auth_config
 +                    .redirect_port
@@ -193,108 +176,55 @@
              .is_err());
      }
 +
-+    // Note: AuthMechanism/AuthFlow enum tests are in auth::config::tests
++    // Note: AuthType enum tests are in auth::config::tests
 +
 +    #[test]
-+    fn test_database_set_auth_mechanism_option() {
++    fn test_database_set_auth_type_option() {
 +        let mut db = Database::new();
 +
-+        // Test valid PAT value (0)
 +        db.set_option(
-+            OptionDatabase::Other("databricks.auth.mechanism".into()),
-+            OptionValue::String("0".into()),
++            OptionDatabase::Other("databricks.auth.type".into()),
++            OptionValue::String("access_token".into()),
 +        )
 +        .unwrap();
-+        assert_eq!(db.auth_config.mechanism, Some(AuthMechanism::Pat));
++        assert_eq!(db.auth_config.auth_type, Some(AuthType::AccessToken));
 +
-+        // Test valid OAuth value (11)
 +        db.set_option(
-+            OptionDatabase::Other("databricks.auth.mechanism".into()),
-+            OptionValue::String("11".into()),
++            OptionDatabase::Other("databricks.auth.type".into()),
++            OptionValue::String("oauth_m2m".into()),
 +        )
 +        .unwrap();
-+        assert_eq!(db.auth_config.mechanism, Some(AuthMechanism::OAuth));
++        assert_eq!(db.auth_config.auth_type, Some(AuthType::OAuthM2m));
 +
-+        // Test with OptionValue::Int
 +        db.set_option(
-+            OptionDatabase::Other("databricks.auth.mechanism".into()),
-+            OptionValue::Int(0),
++            OptionDatabase::Other("databricks.auth.type".into()),
++            OptionValue::String("oauth_u2m".into()),
 +        )
 +        .unwrap();
-+        assert_eq!(db.auth_config.mechanism, Some(AuthMechanism::Pat));
-+    }
++        assert_eq!(db.auth_config.auth_type, Some(AuthType::OAuthU2m));
 +
-+    #[test]
-+    fn test_database_set_auth_mechanism_invalid() {
-+        let mut db = Database::new();
-+
-+        // Test invalid integer value
-+        let result = db.set_option(
-+            OptionDatabase::Other("databricks.auth.mechanism".into()),
-+            OptionValue::String("99".into()),
++        // Verify get_option_string round-trip
++        assert_eq!(
++            db.get_option_string(OptionDatabase::Other("databricks.auth.type".into()))
++                .unwrap(),
++            "oauth_u2m"
 +        );
-+        assert!(result.is_err());
-+
-+        // Test invalid non-integer string
-+        let result = db.set_option(
-+            OptionDatabase::Other("databricks.auth.mechanism".into()),
-+            OptionValue::String("invalid".into()),
-+        );
-+        assert!(result.is_err());
 +    }
 +
 +    #[test]
-+    fn test_database_set_auth_flow_option() {
-+        let mut db = Database::new();
-+
-+        // Test TokenPassthrough (0)
-+        db.set_option(
-+            OptionDatabase::Other("databricks.auth.flow".into()),
-+            OptionValue::String("0".into()),
-+        )
-+        .unwrap();
-+        assert_eq!(db.auth_config.flow, Some(AuthFlow::TokenPassthrough));
-+
-+        // Test ClientCredentials (1)
-+        db.set_option(
-+            OptionDatabase::Other("databricks.auth.flow".into()),
-+            OptionValue::String("1".into()),
-+        )
-+        .unwrap();
-+        assert_eq!(db.auth_config.flow, Some(AuthFlow::ClientCredentials));
-+
-+        // Test Browser (2)
-+        db.set_option(
-+            OptionDatabase::Other("databricks.auth.flow".into()),
-+            OptionValue::String("2".into()),
-+        )
-+        .unwrap();
-+        assert_eq!(db.auth_config.flow, Some(AuthFlow::Browser));
-+
-+        // Test with OptionValue::Int
-+        db.set_option(
-+            OptionDatabase::Other("databricks.auth.flow".into()),
-+            OptionValue::Int(1),
-+        )
-+        .unwrap();
-+        assert_eq!(db.auth_config.flow, Some(AuthFlow::ClientCredentials));
-+    }
-+
-+    #[test]
-+    fn test_database_set_auth_flow_invalid() {
++    fn test_database_set_auth_type_invalid() {
 +        let mut db = Database::new();
 +
-+        // Test invalid integer value
 +        let result = db.set_option(
-+            OptionDatabase::Other("databricks.auth.flow".into()),
-+            OptionValue::String("5".into()),
++            OptionDatabase::Other("databricks.auth.type".into()),
++            OptionValue::String("invalid".into()),
 +        );
 +        assert!(result.is_err());
 +
-+        // Test invalid non-integer string
++        // Non-string should fail
 +        let result = db.set_option(
-+            Op
8000
tionDatabase::Other("databricks.auth.flow".into()),
-+            OptionValue::String("browser".into()),
++            OptionDatabase::Other("databricks.auth.type".into()),
++            OptionValue::Int(0),
 +        );
 +        assert!(result.is_err());
 +    }
@@ -482,6 +412,9 @@
 +        let db = Database::new();
 +
 +        assert!(db
++            .get_option_string(OptionDatabase::Other("databricks.auth.type".into()))
++            .is_err());
++        assert!(db
 +            .get_option_string(OptionDatabase::Other("databricks.auth.client_id".into()))
 +            .is_err());
 +        assert!(db
@@ -496,94 +429,18 @@
 +            .get_option_string(OptionDatabase::Other(
 +                "databricks.auth.token_endpoint".into()
 +            ))
-+            .is_err());
-+        assert!(db
-+            .get_option_int(OptionDatabase::Other("databricks.auth.mechanism".into()))
-+            .is_err());
-+        assert!(db
-+            .get_option_int(OptionDatabase::Other("databricks.auth.flow".into()))
 +            .is_err());
 +        assert!(db
 +            .get_option_int(OptionDatabase::Other(
 +                "databricks.auth.redirect_port".into()
 +            ))
 +            .is_err());
-+    }
-+
-+    #[test]
-+    fn test_database_get_auth_mechanism_as_int() {
-+        let mut db = Database::new();
-+
-+        // Set PAT mechanism
-+        db.set_option(
-+            OptionDatabase::Other("databricks.auth.mechanism".into()),
-+            OptionValue::Int(0),
-+        )
-+        .unwrap();
-+        assert_eq!(
-+            db.get_option_int(OptionDatabase::Other("databricks.auth.mechanism".into()))
-+                .unwrap(),
-+            0
-+        );
-+
-+        // Set OAuth mechanism
-+        db.set_option(
-+            OptionDatabase::Other("databricks.auth.mechanism".into()),
-+            OptionValue::Int(11),
-+        )
-+        .unwrap();
-+        assert_eq!(
-+            db.get_option_int(OptionDatabase::Other("databricks.auth.mechanism".into()))
-+                .unwrap(),
-+            11
-+        );
-+    }
-+
-+    #[test]
-+    fn test_database_get_auth_flow_as_int() {
-+        let mut db = Database::new();
-+
-+        // Set TokenPassthrough flow
-+        db.set_option(
-+            OptionDatabase::Other("databricks.auth.flow".into()),
-+            OptionValue::Int(0),
-+        )
-+        .unwrap();
-+        assert_eq!(
-+            db.get_option_int(OptionDatabase::Other("databricks.auth.flow".into()))
-+                .unwrap(),
-+            0
-+        );
-+
-+        // Set ClientCredentials flow
-+        db.set_option(
-+            OptionDatabase::Other("databricks.auth.flow".into()),
-+            OptionValue::Int(1),
-+        )
-+        .unwrap();
-+        assert_eq!(
-+            db.get_option_int(OptionDatabase::Other("databricks.auth.flow".into()))
-+                .unwrap(),
-+            1
-+        );
-+
-+        // Set Browser flow
-+        db.set_option(
-+            OptionDatabase::Other("databricks.auth.flow".into()),
-+            OptionValue::Int(2),
-+        )
-+        .unwrap();
-+        assert_eq!(
-+            db.get_option_int(OptionDatabase::Other("databricks.auth.flow".into()))
-+                .unwrap(),
-+            2
-+        );
 +    }
 +
 +    // Config validation tests for new_connection()
 +
 +    #[test]
-+    fn test_new_connection_missing_mechanism() {
++    fn test_new_connection_missing_auth_type() {
 +        use adbc_core::Database as _;
 +
 +        let mut db = Database::new();
@@ -602,14 +459,14 @@
 +        assert!(result.is_err());
 +        let err_msg = format!("{:?}", result.unwrap_err());
 +        assert!(
-+            err_msg.contains("databricks.auth.mechanism is required"),
-+            "Expected error message about missing mechanism, got: {}",
++            err_msg.contains("databricks.auth.type is required"),
++            "Expected error message about missing auth type, got: {}",
 +            err_msg
 +        );
 +    }
 +
 +    #[test]
-+    fn test_new_connection_oauth_missing_flow() {
++    fn test_new_connection_oauth_m2m_missing_client_id() {
 +        use adbc_core::Database as _;
 +
 +        let mut db = Database::new();
@@ -624,47 +481,10 @@
 +        )
 +        .unwrap();
 +        db.set_option(
-+            OptionDatabase::Other("databricks.auth.mechanism".into()),
-+            OptionValue::Int(11), // OAuth
-+        )
-+        .unwrap();
-+
-+        let result = db.new_connection();
-+        assert!(result.is_err());
-+        let err_msg = format!("{:?}", result.unwrap_err());
-+        assert!(
-+            err_msg.contains("databricks.auth.flow is required"),
-+            "Expected error message about missing flow, got: {}",
-+            err_msg
-+        );
-+    }
-+
-+    #[test]
-+    fn test_new_connection_client_credentials_missing_client_id() {
-+        use adbc_core::Database as _;
-+
-+        let mut db = Database::new();
-+        db.set_option(
-+            OptionDatabase::Uri,
-+            OptionValue::String("https://example.databricks.com".into()),
++            OptionDatabase::Other("databricks.auth.type".into()),
++            OptionValue::String("oauth_m2m".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(1), // ClientCredentials
-+        )
-+        .unwrap();
-+        // Missing client_id
 +
 +        let result = db.new_connection();
 +        assert!(result.is_err());
@@ -677,7 +497,7 @@
 +    }
 +
 +    #[test]
-+    fn test_new_connection_client_credentials_missing_secret() {
++    fn test_new_connection_oauth_m2m_missing_secret() {
 +        use adbc_core::Database as _;
 +
 +        let mut db = Database::new();
@@ -692,13 +512,8 @@
 +        )
 +        .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(1), // ClientCredentials
++            OptionDatabase::Other("databricks.auth.type".into()),
++            OptionValue::String("oauth_m2m".into()),
 +        )
 +        .unwrap();
 +        db.set_option(
@@ -706,7 +521,6 @@
 +            OptionValue::String("test-client-id".into()),
 +        )
 +        .unwrap();
-+        // Missing client_secret
 +
 +        let result = db.new_connection();
 +        assert!(result.is_err());
@@ -719,7 +533,7 @@
 +    }
 +
 +    #[test]
-+    fn test_new_connection_pat_missing_token() {
++    fn test_new_connection_access_token_missing_token() {
 +        use adbc_core::Database as _;
 +
 +        let mut db = Database::new();
@@ -734,55 +548,17 @@
 +        )
 +        .unwrap();
 +        db.set_option(
-+            OptionDatabase::Other("databricks.auth.mechanism".into()),
-+            OptionValue::Int(0), // PAT
++            OptionDatabase::Other("databricks.auth.type".into()),
++            OptionValue::String("access_token".into()),
 +        )
 +        .unwrap();
-+        // Missing access_token
 +
 +        let result = db.new_connection();
 +        assert!(result.is_err());
 +        let err_msg = format!("{:?}", result.unwrap_err());
 +        assert!(
 +            err_msg.contains("databricks.access_token is required"),
-+            "Expected error message about missing access_token for PAT, got: {}",
-+            err_msg
-+        );
-+    }
-+
-+    #[test]
-+    fn test_new_connection_token_passthrough_missing_token() {
-+        use adbc_core::Database as _;
-+
-+        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(0), // TokenPassthrough
-+        )
-+        .unwrap();
-+        // Missing access_token
-+
-+        let result = db.new_connection();
-+        assert!(result.is_err());
-+        let err_msg = format!("{:?}", result.unwrap_err());
-+        assert!(
-+            err_msg.contains("databricks.access_token is required"),
-+            "Expected error message about missing access_token for TokenPassthrough, got: {}",
++            "Expected error message about missing access_token, got: {}",
 +            err_msg
 +        );
 +    }
rust/tests/integration.rs
@@ -0,0 +1,19 @@
+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());
+ }
+ 
+ #[test]
\ No newline at end of file

Reproduce locally: git range-diff 1651567..f23bf89 f0a904f..2e70c60 | Disable: git config gitstack.push-range-diff false

@vikrantpuppala vikrantpuppala force-pushed the stack/pr-database-config branch from 2e70c60 to 296931c 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-database-config branch 2 times, most recently from 36f6b59 to 1222489 Compare March 13, 2026 12:21
@vikrantpuppala vikrantpuppala changed the title [PECOBLR-2165] feat(rust/oauth): Auth config enums and Database option parsing feat(rust): Auth config enums and Database option parsing Mar 13, 2026
@vikrantpuppala vikrantpuppala force-pushed the stack/pr-database-config branch 3 times, most recently from fce60b0 to 164ada0 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-database-config branch from 164ada0 to 20b42e7 Compare March 13, 2026 16:33
@vikrantpuppala vikrantpuppala added the integration-test Trigger integration tests in internal repo label Mar 13, 2026
@vikrantpuppala vikrantpuppala force-pushed the stack/pr-database-config branch from 20b42e7 to 3cdbf7c Compare March 13, 2026 17:07
@vikrantpuppala vikrantpuppala added integration-test Trigger integration tests in internal repo and removed integration-test Trigger integration tests in internal repo labels Mar 13, 2026
@vikrantpuppala vikrantpuppala requested a review from gopalldb March 13, 2026 17:08
@vikrantpuppala
Copy link
Collaborator Author

Folded into PR #322 to reduce stack size.

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

Labels

integration-test Trigger integration tests in internal repo

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

0