Run full test suite and fix any failures\n\nTask ID: task-5.1-final-validation

Author: vikrantpuppala

rust/src/auth/config.rs

@@ -107,10 +107,9 @@ impl AuthConfig {
         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 => {

rust/src/auth/oauth/cache.rs

@@ -35,7 +35,6 @@ struct CacheKey {
     scopes: Vec<String>,
 }
 
-
 impl CacheKey {
     /// Creates a new cache key from the given parameters.
     fn new(host: &str, client_id: &str, scopes: &[String]) -> Self {
@@ -71,7 +70,6 @@ impl CacheKey {
 /// 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/`

rust/src/auth/oauth/callback.rs

@@ -152,9 +152,10 @@ impl CallbackServer {
     /// Returns the redirect URI for this callback server.
     ///
     /// This URI should be used when constructing the OAuth authorization URL.
-    /// The format is `http://localhost:{port}/callback`.
+    /// The format is `http://localhost:{port}` (matching the redirect URI
+    /// registered for the `databricks-cli` OAuth application).
     pub fn redirect_uri(&self) -> String {
-        format!("http://localhost:{}/callback", self.port)
+        format!("http://localhost:{}", self.port)
     }
 
     /// Waits for the OAuth callback with the authorization code.
@@ -276,7 +277,7 @@ impl CallbackServer {
 
         let request = String::from_utf8_lossy(&buffer[..n]);
 
-        // Parse the request line (e.g., "GET /callback?code=...&state=... HTTP/1.1")
+        // Parse the request line (e.g., "GET /?code=...&state=... HTTP/1.1")
         let first_line = request.lines().next().unwrap_or("");
         let parts: Vec<&str> = first_line.split_whitespace().collect();
 
@@ -424,7 +425,7 @@ mod tests {
                 .expect("Failed to connect to callback server");
 
             let request = format!(
-                "GET /callback?code=test-auth-code-123&state={} HTTP/1.1\r\nHost: localhost:{}\r\n\r\n",
+                "GET /?code=test-auth-code-123&state={} HTTP/1.1\r\nHost: localhost:{}\r\n\r\n",
                 expected_state, port
             );
 
@@ -465,7 +466,7 @@ mod tests {
                 .expect("Failed to connect");
 
             let request = format!(
-                "GET /callback?code=auth-code&state={} HTTP/1.1\r\nHost: localhost:{}\r\n\r\n",
+                "GET /?code=auth-code&state={} HTTP/1.1\r\nHost: localhost:{}\r\n\r\n",
                 wrong_state, port
             );
 
@@ -529,7 +530,7 @@ mod tests {
 
             // Send callback without 'code' parameter
             let request = format!(
-                "GET /callback?state=test-state HTTP/1.1\r\nHost: localhost:{}\r\n\r\n",
+                "GET /?state=test-state HTTP/1.1\r\nHost: localhost:{}\r\n\r\n",
                 port
             );
 
@@ -560,7 +561,7 @@ mod tests {
 
             // Simulate authorization server error response
             let request = format!(
-                "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
             );
 
@@ -619,7 +620,13 @@ mod tests {
             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);
     }
 }

rust/src/auth/oauth/m2m.rs

@@ -188,7 +188,8 @@ impl ClientCredentialsProvider {
         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?;
@@ -209,7 +210,6 @@ impl ClientCredentialsProvider {
             scopes,
         })
     }
-
 }
 
 impl AuthProvider for ClientCredentialsProvider {

rust/src/auth/oauth/token_store.rs

@@ -83,7 +83,6 @@ pub(crate) struct TokenStore {
     refreshing: Arc<AtomicBool>,
 }
 
-
 impl TokenStore {
     /// Creates a new empty token store.
     pub fn new() -> Self {

rust/src/database.rs

@@ -406,13 +406,11 @@ impl Optionable for Database {
                             .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()
@@ -511,7 +509,8 @@ impl adbc_core::Database for Database {
         })?;
 
         // Validate auth configuration
-        let mechanism = self.auth_config
+        let mechanism = self
+            .auth_config
             .validate(&self.access_token)
             .map_err(|e| e.to_adbc())?;
 
@@ -540,7 +539,10 @@ impl adbc_core::Database for Database {
                 ))
             }
             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
@@ -552,25 +554,34 @@ impl adbc_core::Database for Database {
                     }
                     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)
@@ -932,7 +943,10 @@ mod tests {
         )
         .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()
@@ -951,7 +965,10 @@ mod tests {
         )
         .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(),
@@ -1564,10 +1581,13 @@ mod tests {
         .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]
@@ -1612,7 +1632,7 @@ mod tests {
 
         // Verify custom scopes are stored
         assert_eq!(
-            db.auth_scopes,
+            db.auth_config.scopes,
             Some("custom-scope-1 custom-scope-2".to_string())
         );
     }
@@ -1659,7 +1679,7 @@ mod tests {
 
         // Verify token_endpoint override is stored
         assert_eq!(
-            db.auth_token_endpoint,
+            db.auth_config.token_endpoint,
             Some("https://custom.endpoint/token".to_string())
         );
     }

rust/tests/oauth_e2e.rs

@@ -129,9 +129,7 @@ fn test_m2m_end_to_end() {
 
     // 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
@@ -286,9 +284,7 @@ fn test_u2m_end_to_end() {
 
     // 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