Wire U2M provider into Database::new_connection()\n\nTask ID: task-3.3-database-integration

vikrantpuppala · vikrantpuppala · commit ec7600d3990f · 2026-03-13T11:21:41.000Z
diff --git a/rust/src/auth/oauth/cache.rs b/rust/src/auth/oauth/cache.rs
@@ -27,15 +27,14 @@ use std::path::PathBuf;
 
 /// Cache key used for generating the cache filename.
 /// This struct is serialized to JSON and then hashed with SHA-256 to produce a unique filename.
-#[allow(dead_code)] // Used in Phase 3 (U2M)
+
 #[derive(Debug, Serialize, Deserialize)]
 struct CacheKey {
     host: String,
     client_id: String,
     scopes: Vec<String>,
 }
 
-#[allow(dead_code)] // Used in Phase 3 (U2M)
 impl CacheKey {
     /// Creates a new cache key from the given parameters.
     fn new(host: &str, client_id: &str, scopes: &[String]) -> Self {
@@ -69,10 +68,8 @@ impl CacheKey {
 /// set to `0o600` (owner read/write only) for security.
 ///
 /// Cache I/O errors are logged as warnings and never block authentication.
-#[allow(dead_code)] // Used in Phase 3 (U2M)
 pub(crate) struct TokenCache;
 
-#[allow(dead_code)] // Used in Phase 3 (U2M)
 impl TokenCache {
     /// Returns the cache directory path.
     /// Location: `~/.config/databricks-adbc/oauth/`
diff --git a/rust/src/auth/oauth/callback.rs b/rust/src/auth/oauth/callback.rs
@@ -33,80 +33,20 @@ use tokio::sync::oneshot;
 /// HTML response sent to the browser after successful callback.
 const SUCCESS_HTML: &str = r#"<!DOCTYPE html>
 <html>
-<head>
-    <title>Authentication Successful</title>
-    <style>
-        body {
-            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
-            display: flex;
-            justify-content: center;
-            align-items: center;
-            height: 100vh;
-            margin: 0;
-            background-color: #f5f5f5;
-        }
-        .message {
-            text-align: center;
-            padding: 2rem;
-            background: white;
-            border-radius: 8px;
-            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
-        }
-        h1 {
-            color: #2e7d32;
-            margin: 0 0 1rem 0;
-        }
-        p {
-            color: #666;
-            margin: 0;
-        }
-    </style>
-</head>
+<head><title>Authentication Successful</title></head>
 <body>
-    <div class="message">
-        <h1>✓ Authentication Successful</h1>
-        <p>You can close this tab and return to your application.</p>
-    </div>
+<h1>Authentication Successful</h1>
+<p>You can close this tab and return to your application.</p>
 </body>
 </html>"#;
 
 /// HTML response sent to the browser when an error occurs.
 const ERROR_HTML: &str = r#"<!DOCTYPE html>
 <html>
-<head>
-    <title>Authentication Error</title>
-    <style>
-        body {
-            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
-            display: flex;
-            justify-content: center;
-            align-items: center;
-            height: 100vh;
-            margin: 0;
-            background-color: #f5f5f5;
-        }
-        .message {
-            text-align: center;
-            padding: 2rem;
-            background: white;
-            border-radius: 8px;
-            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
-        }
-        h1 {
-            color: #c62828;
-            margin: 0 0 1rem 0;
-        }
-        p {
-            color: #666;
-            margin: 0;
-        }
-    </style>
-</head>
+<head><title>Authentication Error</title></head>
 <body>
-    <div class="message">
-        <h1>✗ Authentication Error</h1>
-        <p>An error occurred during authentication. You can close this tab and try again.</p>
-    </div>
+<h1>Authentication Error</h1>
+<p>An error occurred during authentication. You can close this tab and try again.</p>
 </body>
 </html>"#;
 
@@ -668,9 +608,24 @@ mod tests {
 
     #[tokio::test]
     async fn test_redirect_uri_format() {
-        let server = CallbackServer::new(8020)
+        // Use port 0 to let the OS assign an available port, avoiding conflicts
+        let server = CallbackServer::new(0)
             .await
             .expect("Failed to create server");
-        assert_eq!(server.redirect_uri(), "http://localhost:8020/callback");
+        let uri = server.redirect_uri();
+        assert!(
+            uri.starts_with("http://localhost:"),
+            "Expected URI starting with http://localhost:, got: {}",
+            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();
+        assert!(port > 0, "Expected a non-zero port, got: {}", port);
     }
 }
diff --git a/rust/src/auth/oauth/m2m.rs b/rust/src/auth/oauth/m2m.rs
@@ -170,77 +170,6 @@ impl ClientCredentialsProvider {
             scopes,
         })
     }
-
-    /// Exchanges client credentials for an access token.
-    ///
-    /// This method performs the OAuth 2.0 client credentials grant:
-    /// ```text
-    /// POST {token_endpoint}
-    /// Authorization: Basic base64(client_id:client_secret)
-    /// Content-Type: application/x-www-form-urlencoded
-    ///
-    /// grant_type=client_credentials&scope=all-apis
-    /// ```
-    ///
-    /// Uses the DatabricksHttpClient's inner reqwest::Client to execute the request.
-    /// The oauth2 crate adds Basic authentication header automatically.
-    ///
-    /// Note: This method is currently unused but kept for potential future use
-    /// (e.g., direct testing or alternative refresh strategies).
-    #[allow(dead_code)]
-    async fn fetch_token(&self) -> Result<OAuthToken> {
-        // Build the OAuth client
-        let oauth_client = BasicClient::new(ClientId::new(self.client_id.clone()))
-            .set_client_secret(ClientSecret::new(self.client_secret.clone()))
-            .set_auth_uri(AuthUrl::new(self.auth_endpoint.clone()).map_err(|e| {
-                DatabricksErrorHelper::invalid_argument()
-                    .message(format!("Invalid authorization endpoint URL: {}", e))
-            })?)
-            .set_token_uri(TokenUrl::new(self.token_endpoint.clone()).map_err(|e| {
-                DatabricksErrorHelper::invalid_argument()
-                    .message(format!("Invalid token endpoint URL: {}", e))
-            })?);
-
-        // Build the token request with scopes
-        let mut token_request = oauth_client.exchange_client_credentials();
-
-        for scope in &self.scopes {
-            token_request = token_request.add_scope(Scope::new(scope.clone()));
-        }
-
-        // Execute the token exchange using the inner reqwest client
-        // The oauth2 crate's reqwest::Client implements AsyncHttpClient
-        let token_response = token_request
-            .request_async(self.http_client.inner())
-            .await
-            .map_err(|e| {
-                DatabricksErrorHelper::io()
-                    .message(format!("M2M token exchange failed: {}", e))
-                    .context("client credentials grant")
-            })?;
-
-        // Convert oauth2 token response to our OAuthToken
-        let access_token = token_response.access_token().secret().to_string();
-        let token_type = token_response.token_type().as_ref().to_string();
-        let expires_in = token_response
-            .expires_in()
-            .map(|d| d.as_secs() as i64)
-            .unwrap_or(3600); // Default to 1 hour if not specified
-
-        let scopes = token_response
-            .scopes()
-            .map(|s| s.iter().map(|scope| scope.to_string()).collect())
-            .unwrap_or_else(|| self.scopes.clone());
-
-        // M2M tokens have no refresh_token
-        Ok(OAuthToken::new(
-            access_token,
-            token_type,
-            expires_in,
-            None, // No refresh token for M2M
-            scopes,
-        ))
-    }
 }
 
 impl AuthProvider for ClientCredentialsProvider {
diff --git a/rust/src/auth/oauth/token_store.rs b/rust/src/auth/oauth/token_store.rs
@@ -73,7 +73,7 @@ use std::sync::{Arc, RwLock};
 ///     ))
 /// })?;
 /// ```
-#[allow(dead_code)] // Used in Phase 2 (M2M) and Phase 3 (U2M)
+
 #[derive(Debug)]
 pub(crate) struct TokenStore {
     /// The current token, protected by a read-write lock.
@@ -83,7 +83,6 @@ pub(crate) struct TokenStore {
     refreshing: Arc<AtomicBool>,
 }
 
-#[allow(dead_code)] // Used in Phase 2 (M2M) and Phase 3 (U2M)
 impl TokenStore {
     /// Creates a new empty token store.
     pub fn new() -> Self {
@@ -284,7 +283,6 @@ impl TokenStore {
 }
 
 /// Internal enum representing the token's current state.
-#[allow(dead_code)] // Used internally by TokenStore
 #[derive(Debug)]
 enum TokenState {
     /// No token is present in the store.
diff --git a/rust/src/auth/oauth/u2m.rs b/rust/src/auth/oauth/u2m.rs
@@ -111,6 +111,8 @@ pub struct AuthorizationCodeProvider {
     callback_port: u16,
     /// Callback timeout
     callback_timeout: Duration,
+    /// Whether to open browser automatically (set to false in tests)
+    open_browser: bool,
 }
 
 impl AuthorizationCodeProvider {
@@ -179,9 +181,37 @@ impl AuthorizationCodeProvider {
         callback_port: u16,
         callback_timeout: Duration,
     ) -> Result<Self> {
-        // Discover OIDC endpoints
+        Self::new_with_full_config(
+            host,
+            client_id,
+            http_client,
+            scopes,
+            callback_port,
+            callback_timeout,
+            None, // No token_endpoint override
+        )
+        .await
+    }
+
+    /// Creates a new U2M provider with full configuration including optional token endpoint override.
+    ///
+    /// This is used by Database when a custom token endpoint is configured.
+    /// Most users should use `new()` or `new_with_config()` which use OIDC discovery.
+    pub async fn new_with_full_config(
+        host: &str,
+        client_id: &str,
+        http_client: Arc<DatabricksHttpClient>,
+        scopes: Vec<String>,
+        callback_port: u16,
+        callback_timeout: Duration,
+        token_endpoint_override: Option<String>,
+    ) -> Result<Self> {
+        // Discover OIDC endpoints (unless token_endpoint is overridden, we still need auth endpoint)
         let endpoints = OidcEndpoints::discover(host, &http_client).await?;
 
+        // Use override if provided, otherwise use discovered endpoint
+        let token_endpoint = token_endpoint_override.unwrap_or(endpoints.token_endpoint);
+
         // Create token store
         let token_store = TokenStore::new();
 
@@ -206,17 +236,29 @@ impl AuthorizationCodeProvider {
         Ok(Self {
             host: host.to_string(),
             client_id: client_id.to_string(),
-            token_endpoint: endpoints.token_endpoint,
+            token_endpoint,
             auth_endpoint: endpoints.authorization_endpoint,
             token_store,
             http_client,
             scopes,
             callback_port,
             callback_timeout,
+            open_browser: true,
         })
     }
 }
 
+impl AuthorizationCodeProvider {
+    /// Disables automatic browser opening during the U2M flow.
+    ///
+    /// This is primarily useful for testing scenarios where the browser
+    /// flow fallback should not launch a real browser.
+    #[doc(hidden)]
+    pub fn suppress_browser(&mut self) {
+        self.open_browser = false;
+    }
+}
+
 impl AuthProvider for AuthorizationCodeProvider {
     /// Returns a valid Bearer token for authentication.
     ///
@@ -246,6 +288,7 @@ impl AuthProvider for AuthorizationCodeProvider {
         let scopes = self.scopes.clone();
         let callback_port = self.callback_port;
         let callback_timeout = self.callback_timeout;
+        let open_browser = self.open_browser;
         let http_client = self.http_client.clone();
 
         // Get or refresh token via TokenStore
@@ -366,10 +409,14 @@ impl AuthProvider for AuthorizationCodeProvider {
 
                     let (auth_url, csrf_state) = auth_url_builder.url();
 
-                    // Launch browser
-                    tracing::info!("Opening browser for OAuth authorization: {}", auth_url);
-                    if let Err(e) = open::that(auth_url.as_str()) {
-                        tracing::warn!("Failed to automatically open browser: {}. Please manually navigate to: {}", e, auth_url);
+                    // Launch browser (unless suppressed for testing)
+                    if open_browser {
+                        tracing::info!("Opening browser for OAuth authorization: {}", auth_url);
+                        if let Err(e) = open::that(auth_url.as_str()) {
+                            tracing::warn!("Failed to automatically open browser: {}. Please manually navigate to: {}", e, auth_url);
+                        }
+                    } else {
+                        tracing::debug!("Browser opening suppressed, authorization URL: {}", auth_url);
                     }
 
                     // Wait for callback
@@ -607,17 +654,18 @@ mod tests {
             &cached_token,
         );
 
-        // Create provider
-        let provider = AuthorizationCodeProvider::new_with_config(
+        // Create provider with browser opening suppressed
+        let mut provider = AuthorizationCodeProvider::new_with_config(
             &mock_server.uri(),
             "test-client-id",
             http_client,
             vec!["all-apis".to_string(), "offline_access".to_string()],
-            8021,                   // Different port to avoid conflicts
+            0,                      // Use port 0 to avoid conflicts
             Duration::from_secs(5), // Short timeout for test
         )
         .await
         .expect("Failed to create provider");
+        provider.suppress_browser();
 
         // Wait to ensure token expires
         tokio::time::sleep(Duration::from_secs(2)).await;
diff --git a/rust/src/database.rs b/rust/src/database.rs
