olai_http/databricks/
builder.rs

1use std::str::FromStr;
2use std::sync::Arc;
3
4use futures::future::BoxFuture;
5use serde::{Deserialize, Serialize};
6use tokio::runtime::Handle;
7
8use super::cfg_file::load_cfg_profile;
9use super::credential::{
10    DatabricksCredential, DatabricksGcpTokenExchangeProvider, DatabricksM2MProvider,
11    OidcEnvTokenProvider, OidcFileTokenProvider,
12};
13use crate::azure::credential::{AZURE_DATABRICKS_RESOURCE, ImdsManagedIdentityProvider};
14// AZURE_DATABRICKS_SCOPE is used by the Azure SP two-token flow (Phase 2.5)
15#[allow(unused_imports)]
16use crate::azure::credential::AZURE_DATABRICKS_SCOPE;
17use crate::service::make_service;
18use crate::{
19    ClientConfigKey, ClientOptions, CredentialProvider, RequestSigner, Result, RetryConfig,
20    StaticCredentialProvider, TokenCredentialProvider,
21};
22
23/// Default OAuth scope for Databricks Unity Catalog (least privilege).
24const DEFAULT_DATABRICKS_SCOPE: &str = "unity-catalog";
25
26/// Default env var name holding an OIDC token.
27const DEFAULT_OIDC_TOKEN_ENV: &str = "DATABRICKS_OIDC_TOKEN";
28
29type DatabricksCredentialProvider = Arc<dyn CredentialProvider<Credential = DatabricksCredential>>;
30
31/// Unified Databricks auth configuration (result of building a [`DatabricksBuilder`]).
32#[derive(Debug, Clone)]
33pub struct DatabricksConfig {
34    pub credentials: DatabricksCredentialProvider,
35    pub retry_config: RetryConfig,
36    pub client_options: ClientOptions,
37}
38
39impl RequestSigner for DatabricksConfig {
40    fn sign<'a>(
41        &'a self,
42        req: reqwest::RequestBuilder,
43    ) -> BoxFuture<'a, Result<reqwest::RequestBuilder>> {
44        Box::pin(async move {
45            let cred = self.credentials.get_credential().await?;
46            Ok(req.bearer_auth(&cred.bearer))
47        })
48    }
49}
50
51/// Configuration keys for [`DatabricksBuilder`].
52#[derive(PartialEq, Eq, Hash, Clone, Debug, Copy, Deserialize, Serialize)]
53#[non_exhaustive]
54pub enum DatabricksConfigKey {
55    /// Databricks workspace host URL.
56    ///
57    /// Supported keys: `databricks_host`, `host`
58    Host,
59
60    /// Personal Access Token or static token.
61    ///
62    /// Supported keys: `databricks_token`, `token`
63    Token,
64
65    /// OAuth M2M service principal client ID.
66    ///
67    /// Supported keys: `databricks_client_id`, `client_id`
68    ClientId,
69
70    /// OAuth M2M service principal client secret.
71    ///
72    /// Supported keys: `databricks_client_secret`, `client_secret`
73    ClientSecret,
74
75    /// Account ID (for account-level operations).
76    ///
77    /// Supported keys: `databricks_account_id`, `account_id`
78    AccountId,
79
80    /// Azure resource ID; presence triggers Azure MSI fallback path.
81    ///
82    /// Supported keys: `databricks_azure_resource_id`, `azure_resource_id`
83    AzureResourceId,
84
85    /// OAuth scope override (defaults to `unity-catalog`).
86    ///
87    /// Supported keys: `databricks_oauth_scope`, `oauth_scope`
88    Scope,
89
90    /// Path to a non-default `.databrickscfg` file.
91    ///
92    /// Supported keys: `databricks_config_file`, `config_file`
93    ConfigFile,
94
95    /// Profile name within `.databrickscfg`.
96    ///
97    /// Supported keys: `databricks_config_profile`, `config_profile`, `profile`
98    ConfigProfile,
99
100    /// Force a specific auth type (e.g. `pat`, `oauth-m2m`, `env-oidc`, `file-oidc`).
101    ///
102    /// Supported keys: `databricks_auth_type`, `auth_type`
103    AuthType,
104
105    /// Name of the environment variable holding an OIDC token.
106    /// Defaults to `DATABRICKS_OIDC_TOKEN`.
107    ///
108    /// Supported keys: `databricks_oidc_token_env`, `oidc_token_env`
109    OidcTokenEnv,
110
111    /// File path to an OIDC token.
112    ///
113    /// Supported keys: `databricks_oidc_token_filepath`, `oidc_token_filepath`
114    OidcTokenFilepath,
115
116    /// Delegated HTTP client config keys.
117    Client(ClientConfigKey),
118}
119
120impl AsRef<str> for DatabricksConfigKey {
121    fn as_ref(&self) -> &str {
122        match self {
123            Self::Host => "databricks_host",
124            Self::Token => "databricks_token",
125            Self::ClientId => "databricks_client_id",
126            Self::ClientSecret => "databricks_client_secret",
127            Self::AccountId => "databricks_account_id",
128            Self::AzureResourceId => "databricks_azure_resource_id",
129            Self::Scope => "databricks_oauth_scope",
130            Self::ConfigFile => "databricks_config_file",
131            Self::ConfigProfile => "databricks_config_profile",
132            Self::AuthType => "databricks_auth_type",
133            Self::OidcTokenEnv => "databricks_oidc_token_env",
134            Self::OidcTokenFilepath => "databricks_oidc_token_filepath",
135            Self::Client(key) => key.as_ref(),
136        }
137    }
138}
139
140impl FromStr for DatabricksConfigKey {
141    type Err = crate::Error;
142
143    fn from_str(s: &str) -> Result<Self, Self::Err> {
144        match s {
145            "databricks_host" | "host" => Ok(Self::Host),
146            "databricks_token" | "token" => Ok(Self::Token),
147            "databricks_client_id" | "client_id" => Ok(Self::ClientId),
148            "databricks_client_secret" | "client_secret" => Ok(Self::ClientSecret),
149            "databricks_account_id" | "account_id" => Ok(Self::AccountId),
150            "databricks_azure_resource_id" | "azure_resource_id" => Ok(Self::AzureResourceId),
151            "databricks_oauth_scope" | "oauth_scope" => Ok(Self::Scope),
152            "databricks_config_file" | "config_file" => Ok(Self::ConfigFile),
153            "databricks_config_profile" | "config_profile" | "profile" => Ok(Self::ConfigProfile),
154            "databricks_auth_type" | "auth_type" => Ok(Self::AuthType),
155            "databricks_oidc_token_env" | "oidc_token_env" => Ok(Self::OidcTokenEnv),
156            "databricks_oidc_token_filepath" | "oidc_token_filepath" => Ok(Self::OidcTokenFilepath),
157            _ => match s.strip_prefix("databricks_").unwrap_or(s).parse() {
158                Ok(key) => Ok(Self::Client(key)),
159                Err(_) => Err(crate::Error::UnknownConfigurationKey { key: s.into() }),
160            },
161        }
162    }
163}
164
165/// Builder for Databricks authentication configuration.
166///
167/// Auth resolution order (when no `auth_type` is forced):
168/// 1. Explicit `with_credentials(provider)` override.
169/// 2. PAT / static token (`DATABRICKS_TOKEN` / `with_token()`).
170/// 3. OAuth M2M (`DATABRICKS_CLIENT_ID` + `DATABRICKS_CLIENT_SECRET`).
171/// 4. `env-oidc` — OIDC token from env var (`DATABRICKS_OIDC_TOKEN_ENV`, default `DATABRICKS_OIDC_TOKEN`).
172/// 5. `file-oidc` — OIDC token from file (`DATABRICKS_OIDC_TOKEN_FILEPATH`).
173/// 6. Azure MSI fallback (if `DATABRICKS_AZURE_RESOURCE_ID` is set).
174/// 7. GCP service account token exchange (if `GOOGLE_APPLICATION_CREDENTIALS` is set).
175///
176/// If `DATABRICKS_AUTH_TYPE` is set (or `with_auth_type()` called), only the named
177/// auth type is attempted; an error is returned immediately if its required fields are absent.
178///
179/// `.databrickscfg` profile values are loaded as the lowest-priority source — env vars
180/// and code-level setters override them.
181#[derive(Default, Clone)]
182pub struct DatabricksBuilder {
183    host: Option<String>,
184    token: Option<String>,
185    client_id: Option<String>,
186    client_secret: Option<String>,
187    account_id: Option<String>,
188    azure_resource_id: Option<String>,
189    scope: Option<String>,
190    config_file: Option<String>,
191    config_profile: Option<String>,
192    auth_type: Option<String>,
193    oidc_token_env: Option<String>,
194    oidc_token_filepath: Option<String>,
195    retry_config: RetryConfig,
196    client_options: ClientOptions,
197    credentials: Option<DatabricksCredentialProvider>,
198}
199
200impl std::fmt::Debug for DatabricksBuilder {
201    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
202        write!(f, "DatabricksBuilder {{ host: {:?} }}", self.host)
203    }
204}
205
206impl DatabricksBuilder {
207    /// Create a new [`DatabricksBuilder`] with default values.
208    pub fn new() -> Self {
209        Default::default()
210    }
211
212    /// Populate from environment variables and optionally a `.databrickscfg` profile.
213    ///
214    /// Profile values from `.databrickscfg` (or `DATABRICKS_CONFIG_FILE`) are loaded first
215    /// as the lowest-priority source. Environment variables override profile values.
216    ///
217    /// Variables read: `DATABRICKS_HOST`, `DATABRICKS_TOKEN`, `DATABRICKS_CLIENT_ID`,
218    /// `DATABRICKS_CLIENT_SECRET`, `DATABRICKS_ACCOUNT_ID`, `DATABRICKS_AZURE_RESOURCE_ID`,
219    /// `DATABRICKS_OAUTH_SCOPE`, `DATABRICKS_CONFIG_FILE`, `DATABRICKS_CONFIG_PROFILE`,
220    /// `DATABRICKS_AUTH_TYPE`, `DATABRICKS_OIDC_TOKEN_ENV`, `DATABRICKS_OIDC_TOKEN_FILEPATH`,
221    /// `GOOGLE_APPLICATION_CREDENTIALS`.
222    pub fn from_env() -> Self {
223        let mut builder = Self::default();
224
225        // First pass: collect DATABRICKS_* env vars into a map so we can read
226        // CONFIG_FILE / CONFIG_PROFILE before applying other values.
227        let mut env_map: Vec<(DatabricksConfigKey, String)> = Vec::new();
228        let mut config_file_env: Option<String> = None;
229        let mut config_profile_env: Option<String> = None;
230
231        for (os_key, os_value) in std::env::vars_os() {
232            if let (Some(key), Some(value)) = (os_key.to_str(), os_value.to_str()) {
233                if key.starts_with("DATABRICKS_") || key == "GOOGLE_APPLICATION_CREDENTIALS" {
234                    if let Ok(config_key) = key.to_ascii_lowercase().parse::<DatabricksConfigKey>()
235                    {
236                        match config_key {
237                            DatabricksConfigKey::ConfigFile => {
238                                config_file_env = Some(value.to_owned())
239                            }
240                            DatabricksConfigKey::ConfigProfile => {
241                                config_profile_env = Some(value.to_owned())
242                            }
243                            _ => env_map.push((config_key, value.to_owned())),
244                        }
245                    }
246                }
247            }
248        }
249
250        // Apply profile values first (lowest priority).
251        if let Ok(profile_values) =
252            load_cfg_profile(config_file_env.as_deref(), config_profile_env.as_deref())
253        {
254            for (k, v) in profile_values {
255                builder = builder.with_config(k, v);
256            }
257        }
258
259        // Apply env vars (override profile).
260        for (key, value) in env_map {
261            builder = builder.with_config(key, value);
262        }
263
264        // Restore config_file / config_profile from env (they were extracted above).
265        if let Some(v) = config_file_env {
266            builder.config_file = Some(v);
267        }
268        if let Some(v) = config_profile_env {
269            builder.config_profile = Some(v);
270        }
271
272        builder
273    }
274
275    /// Set an option via a key-value pair.
276    pub fn with_config(mut self, key: DatabricksConfigKey, value: impl Into<String>) -> Self {
277        match key {
278            DatabricksConfigKey::Host => self.host = Some(value.into()),
279            DatabricksConfigKey::Token => self.token = Some(value.into()),
280            DatabricksConfigKey::ClientId => self.client_id = Some(value.into()),
281            DatabricksConfigKey::ClientSecret => self.client_secret = Some(value.into()),
282            DatabricksConfigKey::AccountId => self.account_id = Some(value.into()),
283            DatabricksConfigKey::AzureResourceId => self.azure_resource_id = Some(value.into()),
284            DatabricksConfigKey::Scope => self.scope = Some(value.into()),
285            DatabricksConfigKey::ConfigFile => self.config_file = Some(value.into()),
286            DatabricksConfigKey::ConfigProfile => self.config_profile = Some(value.into()),
287            DatabricksConfigKey::AuthType => self.auth_type = Some(value.into()),
288            DatabricksConfigKey::OidcTokenEnv => self.oidc_token_env = Some(value.into()),
289            DatabricksConfigKey::OidcTokenFilepath => self.oidc_token_filepath = Some(value.into()),
290            DatabricksConfigKey::Client(key) => {
291                self.client_options = self.client_options.with_config(key, value)
292            }
293        }
294        self
295    }
296
297    /// Override credential resolution with a custom provider.
298    pub fn with_credentials(mut self, credentials: DatabricksCredentialProvider) -> Self {
299        self.credentials = Some(credentials);
300        self
301    }
302
303    /// Set the Databricks host URL.
304    pub fn with_host(mut self, host: impl Into<String>) -> Self {
305        self.host = Some(host.into());
306        self
307    }
308
309    /// Set a static PAT or bearer token.
310    pub fn with_token(mut self, token: impl Into<String>) -> Self {
311        self.token = Some(token.into());
312        self
313    }
314
315    /// Force a specific auth type (e.g. `"pat"`, `"oauth-m2m"`, `"env-oidc"`, `"file-oidc"`).
316    pub fn with_auth_type(mut self, auth_type: impl Into<String>) -> Self {
317        self.auth_type = Some(auth_type.into());
318        self
319    }
320
321    /// Set the env var name that holds an OIDC token (default: `DATABRICKS_OIDC_TOKEN`).
322    pub fn with_oidc_token_env(mut self, env_var: impl Into<String>) -> Self {
323        self.oidc_token_env = Some(env_var.into());
324        self
325    }
326
327    /// Set the file path to an OIDC token.
328    pub fn with_oidc_token_filepath(mut self, path: impl Into<String>) -> Self {
329        self.oidc_token_filepath = Some(path.into());
330        self
331    }
332
333    /// Set the retry configuration.
334    pub fn with_retry(mut self, retry_config: RetryConfig) -> Self {
335        self.retry_config = retry_config;
336        self
337    }
338
339    /// Set the HTTP client options.
340    pub fn with_client_options(mut self, options: ClientOptions) -> Self {
341        self.client_options = options;
342        self
343    }
344
345    /// Build a [`DatabricksConfig`].
346    ///
347    /// If `runtime` is provided, HTTP I/O for credential refresh is spawned on that runtime.
348    pub fn build(self, runtime: Option<&Handle>) -> Result<DatabricksConfig> {
349        let scope = self
350            .scope
351            .as_deref()
352            .unwrap_or(DEFAULT_DATABRICKS_SCOPE)
353            .to_owned();
354
355        // Determine which auth types are candidates based on available fields.
356        let has_pat = self.token.is_some();
357        let has_m2m = self.client_id.is_some() && self.client_secret.is_some();
358        let has_env_oidc =
359            self.oidc_token_env.is_some() || std::env::var(DEFAULT_OIDC_TOKEN_ENV).is_ok();
360        let has_file_oidc = self.oidc_token_filepath.is_some();
361        let has_azure_msi = self.azure_resource_id.is_some();
362        let has_gcp = std::env::var("GOOGLE_APPLICATION_CREDENTIALS").is_ok();
363
364        // Extract fields we need after auth resolution.
365        let retry_config = self.retry_config.clone();
366        let client_options = self.client_options.clone();
367
368        let credentials: DatabricksCredentialProvider = if let Some(creds) = self.credentials {
369            // 1. Explicit override — always wins regardless of auth_type.
370            creds
371        } else if let Some(auth_type) = self.auth_type.clone() {
372            // Forced auth type — only attempt the named type.
373            self.build_forced_auth(
374                auth_type,
375                scope,
376                runtime,
377                has_pat,
378                has_m2m,
379                has_env_oidc,
380                has_file_oidc,
381                has_azure_msi,
382                has_gcp,
383            )?
384        } else if has_pat {
385            // 2. PAT / static token
386            Arc::new(StaticCredentialProvider::new(DatabricksCredential {
387                bearer: self.token.unwrap(),
388            }))
389        } else if has_m2m {
390            // 3. OAuth M2M (AWS/GCP Databricks-hosted)
391            let host = self.host.as_deref().unwrap_or("").trim_end_matches('/');
392            let token_url = format!("{host}/oidc/v1/token");
393            let provider = DatabricksM2MProvider {
394                token_url,
395                client_id: self.client_id.unwrap(),
396                client_secret: self.client_secret.unwrap(),
397                scope,
398            };
399            let client = client_options.client()?;
400            let service = make_service(client.clone(), runtime);
401            Arc::new(TokenCredentialProvider::new(
402                provider,
403                client,
404                service,
405                retry_config.clone(),
406            )) as _
407        } else if has_env_oidc {
408            // 4. env-oidc
409            let env_var = self
410                .oidc_token_env
411                .unwrap_or_else(|| DEFAULT_OIDC_TOKEN_ENV.to_owned());
412            let provider = OidcEnvTokenProvider { env_var };
413            let client = client_options.client()?;
414            let service = make_service(client.clone(), runtime);
415            Arc::new(TokenCredentialProvider::new(
416                provider,
417                client,
418                service,
419                retry_config.clone(),
420            )) as _
421        } else if has_file_oidc {
422            // 5. file-oidc
423            let provider = OidcFileTokenProvider {
424                filepath: self.oidc_token_filepath.unwrap(),
425            };
426            let client = client_options.client()?;
427            let service = make_service(client.clone(), runtime);
428            Arc::new(TokenCredentialProvider::new(
429                provider,
430                client,
431                service,
432                retry_config.clone(),
433            )) as _
434        } else if has_azure_msi {
435            // 6. Azure MSI fallback
436            let msi_provider = ImdsManagedIdentityProvider::new_with_resource(
437                None,
438                None,
439                None,
440                None,
441                AZURE_DATABRICKS_RESOURCE,
442            );
443            let client = client_options.metadata_client()?;
444            let service = make_service(client.clone(), runtime);
445            let azure_provider =
446                TokenCredentialProvider::new(msi_provider, client, service, retry_config.clone());
447            Arc::new(AzureToDatabricksBridge {
448                inner: Arc::new(azure_provider),
449            }) as _
450        } else if has_gcp {
451            // 7. GCP SA token exchange
452            let host = self.host.as_deref().unwrap_or("").trim_end_matches('/');
453            let token_url = format!("{host}/oidc/v1/token");
454            let gcp_config = crate::gcp::GoogleBuilder::new().build(runtime)?;
455            let provider = DatabricksGcpTokenExchangeProvider {
456                token_url,
457                gcp_provider: gcp_config.credentials,
458            };
459            let client = client_options.client()?;
460            let service = make_service(client.clone(), runtime);
461            Arc::new(TokenCredentialProvider::new(
462                provider,
463                client,
464                service,
465                retry_config.clone(),
466            )) as _
467        } else {
468            return Err(crate::Error::Generic {
469                source: "No Databricks credentials configured. Set DATABRICKS_TOKEN, \
470                         DATABRICKS_CLIENT_ID + DATABRICKS_CLIENT_SECRET, \
471                         DATABRICKS_OIDC_TOKEN (env-oidc), DATABRICKS_OIDC_TOKEN_FILEPATH (file-oidc), \
472                         DATABRICKS_AZURE_RESOURCE_ID, or GOOGLE_APPLICATION_CREDENTIALS."
473                    .into(),
474            });
475        };
476
477        Ok(DatabricksConfig {
478            credentials,
479            retry_config,
480            client_options,
481        })
482    }
483
484    /// Build credentials when `auth_type` forces a specific method.
485    #[allow(clippy::too_many_arguments)]
486    fn build_forced_auth(
487        self,
488        auth_type: String,
489        scope: String,
490        runtime: Option<&Handle>,
491        has_pat: bool,
492        has_m2m: bool,
493        has_env_oidc: bool,
494        has_file_oidc: bool,
495        has_azure_msi: bool,
496        has_gcp: bool,
497    ) -> Result<DatabricksCredentialProvider> {
498        match auth_type.as_str() {
499            "pat" => {
500                if !has_pat {
501                    return Err(crate::Error::Generic {
502                        source: "auth_type=pat requires DATABRICKS_TOKEN to be set".into(),
503                    });
504                }
505                Ok(Arc::new(StaticCredentialProvider::new(DatabricksCredential {
506                    bearer: self.token.unwrap(),
507                })))
508            }
509            "oauth-m2m" => {
510                if !has_m2m {
511                    return Err(crate::Error::Generic {
512                        source: "auth_type=oauth-m2m requires DATABRICKS_CLIENT_ID and \
513                                 DATABRICKS_CLIENT_SECRET to be set"
514                            .into(),
515                    });
516                }
517                let host = self.host.as_deref().unwrap_or("").trim_end_matches('/');
518                let token_url = format!("{host}/oidc/v1/token");
519                let provider = DatabricksM2MProvider {
520                    token_url,
521                    client_id: self.client_id.unwrap(),
522                    client_secret: self.client_secret.unwrap(),
523                    scope,
524                };
525                let client = self.client_options.client()?;
526                let service = make_service(client.clone(), runtime);
527                Ok(Arc::new(TokenCredentialProvider::new(
528                    provider,
529                    client,
530                    service,
531                    self.retry_config,
532                )) as _)
533            }
534            "env-oidc" => {
535                if !has_env_oidc {
536                    return Err(crate::Error::Generic {
537                        source: "auth_type=env-oidc requires DATABRICKS_OIDC_TOKEN or \
538                                 DATABRICKS_OIDC_TOKEN_ENV to be set"
539                            .into(),
540                    });
541                }
542                let env_var = self
543                    .oidc_token_env
544                    .unwrap_or_else(|| DEFAULT_OIDC_TOKEN_ENV.to_owned());
545                let provider = OidcEnvTokenProvider { env_var };
546                let client = self.client_options.client()?;
547                let service = make_service(client.clone(), runtime);
548                Ok(Arc::new(TokenCredentialProvider::new(
549                    provider,
550                    client,
551                    service,
552                    self.retry_config,
553                )) as _)
554            }
555            "file-oidc" => {
556                if !has_file_oidc {
557                    return Err(crate::Error::Generic {
558                        source: "auth_type=file-oidc requires DATABRICKS_OIDC_TOKEN_FILEPATH \
559                                 to be set"
560                            .into(),
561                    });
562                }
563                let provider = OidcFileTokenProvider {
564                    filepath: self.oidc_token_filepath.unwrap(),
565                };
566                let client = self.client_options.client()?;
567                let service = make_service(client.clone(), runtime);
568                Ok(Arc::new(TokenCredentialProvider::new(
569                    provider,
570                    client,
571                    service,
572                    self.retry_config,
573                )) as _)
574            }
575            "azure-msi" => {
576                if !has_azure_msi {
577                    return Err(crate::Error::Generic {
578                        source: "auth_type=azure-msi requires DATABRICKS_AZURE_RESOURCE_ID \
579                                 to be set"
580                            .into(),
581                    });
582                }
583                let msi_provider = ImdsManagedIdentityProvider::new_with_resource(
584                    None,
585                    None,
586                    None,
587                    None,
588                    AZURE_DATABRICKS_RESOURCE,
589                );
590                let client = self.client_options.metadata_client()?;
591                let service = make_service(client.clone(), runtime);
592                let azure_provider = TokenCredentialProvider::new(
593                    msi_provider,
594                    client,
595                    service,
596                    self.retry_config,
597                );
598                Ok(Arc::new(AzureToDatabricksBridge {
599                    inner: Arc::new(azure_provider),
600                }) as _)
601            }
602            "gcp" => {
603                if !has_gcp {
604                    return Err(crate::Error::Generic {
605                        source: "auth_type=gcp requires GOOGLE_APPLICATION_CREDENTIALS to be set"
606                            .into(),
607                    });
608                }
609                let host = self.host.as_deref().unwrap_or("").trim_end_matches('/');
610                let token_url = format!("{host}/oidc/v1/token");
611                let gcp_config = crate::gcp::GoogleBuilder::new().build(runtime)?;
612                let provider = DatabricksGcpTokenExchangeProvider {
613                    token_url,
614                    gcp_provider: gcp_config.credentials,
615                };
616                let client = self.client_options.client()?;
617                let service = make_service(client.clone(), runtime);
618                Ok(Arc::new(TokenCredentialProvider::new(
619                    provider,
620                    client,
621                    service,
622                    self.retry_config,
623                )) as _)
624            }
625            other => Err(crate::Error::Generic {
626                source: format!("Unknown auth_type: {other:?}. Valid values: pat, oauth-m2m, env-oidc, file-oidc, azure-msi, gcp").into(),
627            }),
628        }
629    }
630}
631
632/// Bridges `AzureCredential::BearerToken` → `DatabricksCredential`.
633///
634/// Azure MSI returns an Azure AD bearer token for the Databricks resource ID;
635/// that token is used directly as the Databricks API bearer token.
636#[derive(Debug)]
637struct AzureToDatabricksBridge {
638    inner: Arc<dyn CredentialProvider<Credential = crate::azure::credential::AzureCredential>>,
639}
640
641#[async_trait::async_trait]
642impl CredentialProvider for AzureToDatabricksBridge {
643    type Credential = DatabricksCredential;
644
645    async fn get_credential(&self) -> crate::Result<Arc<DatabricksCredential>> {
646        use crate::azure::credential::AzureCredential;
647        let cred = self.inner.get_credential().await?;
648        let AzureCredential::BearerToken(token) = cred.as_ref();
649        Ok(Arc::new(DatabricksCredential {
650            bearer: token.clone(),
651        }))
652    }
653}
654
655#[cfg(test)]
656mod tests {
657    use std::collections::HashMap;
658
659    use super::*;
660
661    #[test]
662    fn test_databricks_config_from_map() {
663        let options = HashMap::from([
664            (
665                "databricks_host",
666                "https://my-workspace.azuredatabricks.net",
667            ),
668            ("databricks_token", "dapi1234567890"),
669            ("databricks_oauth_scope", "all-apis"),
670        ]);
671
672        let builder = options
673            .into_iter()
674            .fold(DatabricksBuilder::new(), |b, (k, v)| {
675                b.with_config(k.parse().unwrap(), v)
676            });
677
678        assert_eq!(
679            builder.host.as_deref(),
680            Some("https://my-workspace.azuredatabricks.net")
681        );
682        assert_eq!(builder.token.as_deref(), Some("dapi1234567890"));
683        assert_eq!(builder.scope.as_deref(), Some("all-apis"));
684    }
685
686    #[test]
687    fn test_databricks_config_key_roundtrip() {
688        let cases = [
689            ("databricks_host", DatabricksConfigKey::Host),
690            ("host", DatabricksConfigKey::Host),
691            ("databricks_token", DatabricksConfigKey::Token),
692            ("token", DatabricksConfigKey::Token),
693            ("databricks_client_id", DatabricksConfigKey::ClientId),
694            ("client_id", DatabricksConfigKey::ClientId),
695            (
696                "databricks_client_secret",
697                DatabricksConfigKey::ClientSecret,
698            ),
699            ("client_secret", DatabricksConfigKey::ClientSecret),
700            ("databricks_account_id", DatabricksConfigKey::AccountId),
701            ("account_id", DatabricksConfigKey::AccountId),
702            (
703                "databricks_azure_resource_id",
704                DatabricksConfigKey::AzureResourceId,
705            ),
706            ("azure_resource_id", DatabricksConfigKey::AzureResourceId),
707            ("databricks_oauth_scope", DatabricksConfigKey::Scope),
708            ("oauth_scope", DatabricksConfigKey::Scope),
709            ("databricks_config_file", DatabricksConfigKey::ConfigFile),
710            ("config_file", DatabricksConfigKey::ConfigFile),
711            (
712                "databricks_config_profile",
713                DatabricksConfigKey::ConfigProfile,
714            ),
715            ("config_profile", DatabricksConfigKey::ConfigProfile),
716            ("profile", DatabricksConfigKey::ConfigProfile),
717            ("databricks_auth_type", DatabricksConfigKey::AuthType),
718            ("auth_type", DatabricksConfigKey::AuthType),
719            (
720                "databricks_oidc_token_env",
721                DatabricksConfigKey::OidcTokenEnv,
722            ),
723            ("oidc_token_env", DatabricksConfigKey::OidcTokenEnv),
724            (
725                "databricks_oidc_token_filepath",
726                DatabricksConfigKey::OidcTokenFilepath,
727            ),
728            (
729                "oidc_token_filepath",
730                DatabricksConfigKey::OidcTokenFilepath,
731            ),
732        ];
733
734        for (s, expected) in cases {
735            assert_eq!(
736                s.parse::<DatabricksConfigKey>().unwrap(),
737                expected,
738                "failed for {s}"
739            );
740        }
741    }
742
743    #[tokio::test]
744    async fn test_build_with_pat() {
745        let config = DatabricksBuilder::new()
746            .with_host("https://my.azuredatabricks.net")
747            .with_token("dapimytoken")
748            .build(None)
749            .unwrap();
750
751        let cred = config.credentials.get_credential().await.unwrap();
752        assert_eq!(cred.bearer, "dapimytoken");
753    }
754
755    #[test]
756    fn test_auth_type_enforcement_missing_fields() {
757        let err = DatabricksBuilder::new()
758            .with_auth_type("pat")
759            .build(None)
760            .unwrap_err();
761        assert!(err.to_string().contains("pat"));
762    }
763}
olai_http/databricks/builder.rs

olai_http/databricks/
builder.rs
