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

olai_http/databricks/
builder.rs
