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                && (key.starts_with("DATABRICKS_") || key == "GOOGLE_APPLICATION_CREDENTIALS")
234                && let Ok(config_key) = key.to_ascii_lowercase().parse::<DatabricksConfigKey>()
235            {
236                match config_key {
237                    DatabricksConfigKey::ConfigFile => config_file_env = Some(value.to_owned()),
238                    DatabricksConfigKey::ConfigProfile => {
239                        config_profile_env = Some(value.to_owned())
240                    }
241                    _ => env_map.push((config_key, value.to_owned())),
242                }
243            }
244        }
245
246        // Apply profile values first (lowest priority).
247        if let Ok(profile_values) =
248            load_cfg_profile(config_file_env.as_deref(), config_profile_env.as_deref())
249        {
250            for (k, v) in profile_values {
251                builder = builder.with_config(k, v);
252            }
253        }
254
255        // Apply env vars (override profile).
256        for (key, value) in env_map {
257            builder = builder.with_config(key, value);
258        }
259
260        // Restore config_file / config_profile from env (they were extracted above).
261        if let Some(v) = config_file_env {
262            builder.config_file = Some(v);
263        }
264        if let Some(v) = config_profile_env {
265            builder.config_profile = Some(v);
266        }
267
268        builder
269    }
270
271    /// Set an option via a key-value pair.
272    pub fn with_config(mut self, key: DatabricksConfigKey, value: impl Into<String>) -> Self {
273        match key {
274            DatabricksConfigKey::Host => self.host = Some(value.into()),
275            DatabricksConfigKey::Token => self.token = Some(value.into()),
276            DatabricksConfigKey::ClientId => self.client_id = Some(value.into()),
277            DatabricksConfigKey::ClientSecret => self.client_secret = Some(value.into()),
278            DatabricksConfigKey::AccountId => self.account_id = Some(value.into()),
279            DatabricksConfigKey::AzureResourceId => self.azure_resource_id = Some(value.into()),
280            DatabricksConfigKey::Scope => self.scope = Some(value.into()),
281            DatabricksConfigKey::ConfigFile => self.config_file = Some(value.into()),
282            DatabricksConfigKey::ConfigProfile => self.config_profile = Some(value.into()),
283            DatabricksConfigKey::AuthType => self.auth_type = Some(value.into()),
284            DatabricksConfigKey::OidcTokenEnv => self.oidc_token_env = Some(value.into()),
285            DatabricksConfigKey::OidcTokenFilepath => self.oidc_token_filepath = Some(value.into()),
286            DatabricksConfigKey::Client(key) => {
287                self.client_options = self.client_options.with_config(key, value)
288            }
289        }
290        self
291    }
292
293    /// Override credential resolution with a custom provider.
294    pub fn with_credentials(mut self, credentials: DatabricksCredentialProvider) -> Self {
295        self.credentials = Some(credentials);
296        self
297    }
298
299    /// Set the Databricks host URL.
300    pub fn with_host(mut self, host: impl Into<String>) -> Self {
301        self.host = Some(host.into());
302        self
303    }
304
305    /// Set a static PAT or bearer token.
306    pub fn with_token(mut self, token: impl Into<String>) -> Self {
307        self.token = Some(token.into());
308        self
309    }
310
311    /// Force a specific auth type (e.g. `"pat"`, `"oauth-m2m"`, `"env-oidc"`, `"file-oidc"`).
312    pub fn with_auth_type(mut self, auth_type: impl Into<String>) -> Self {
313        self.auth_type = Some(auth_type.into());
314        self
315    }
316
317    /// Set the env var name that holds an OIDC token (default: `DATABRICKS_OIDC_TOKEN`).
318    pub fn with_oidc_token_env(mut self, env_var: impl Into<String>) -> Self {
319        self.oidc_token_env = Some(env_var.into());
320        self
321    }
322
323    /// Set the file path to an OIDC token.
324    pub fn with_oidc_token_filepath(mut self, path: impl Into<String>) -> Self {
325        self.oidc_token_filepath = Some(path.into());
326        self
327    }
328
329    /// Set the retry configuration.
330    pub fn with_retry(mut self, retry_config: RetryConfig) -> Self {
331        self.retry_config = retry_config;
332        self
333    }
334
335    /// Set the HTTP client options.
336    pub fn with_client_options(mut self, options: ClientOptions) -> Self {
337        self.client_options = options;
338        self
339    }
340
341    /// Build a [`DatabricksConfig`].
342    ///
343    /// If `runtime` is provided, HTTP I/O for credential refresh is spawned on that runtime.
344    pub fn build(self, runtime: Option<&Handle>) -> Result<DatabricksConfig> {
345        let scope = self
346            .scope
347            .as_deref()
348            .unwrap_or(DEFAULT_DATABRICKS_SCOPE)
349            .to_owned();
350
351        // Determine which auth types are candidates based on available fields.
352        let has_pat = self.token.is_some();
353        let has_m2m = self.client_id.is_some() && self.client_secret.is_some();
354        let has_env_oidc =
355            self.oidc_token_env.is_some() || std::env::var(DEFAULT_OIDC_TOKEN_ENV).is_ok();
356        let has_file_oidc = self.oidc_token_filepath.is_some();
357        let has_azure_msi = self.azure_resource_id.is_some();
358        let has_gcp = std::env::var("GOOGLE_APPLICATION_CREDENTIALS").is_ok();
359
360        // Extract fields we need after auth resolution.
361        let retry_config = self.retry_config.clone();
362        let client_options = self.client_options.clone();
363
364        let credentials: DatabricksCredentialProvider = if let Some(creds) = self.credentials {
365            // 1. Explicit override — always wins regardless of auth_type.
366            creds
367        } else if let Some(auth_type) = self.auth_type.clone() {
368            // Forced auth type — only attempt the named type.
369            self.build_forced_auth(
370                auth_type,
371                scope,
372                runtime,
373                has_pat,
374                has_m2m,
375                has_env_oidc,
376                has_file_oidc,
377                has_azure_msi,
378                has_gcp,
379            )?
380        } else if has_pat {
381            // 2. PAT / static token
382            Arc::new(StaticCredentialProvider::new(DatabricksCredential {
383                bearer: self.token.unwrap(),
384            }))
385        } else if has_m2m {
386            // 3. OAuth M2M (AWS/GCP Databricks-hosted)
387            let host = self.host.as_deref().unwrap_or("").trim_end_matches('/');
388            let token_url = format!("{host}/oidc/v1/token");
389            let provider = DatabricksM2MProvider {
390                token_url,
391                client_id: self.client_id.unwrap(),
392                client_secret: self.client_secret.unwrap(),
393                scope,
394            };
395            let client = client_options.client()?;
396            let service = make_service(client.clone(), runtime);
397            Arc::new(TokenCredentialProvider::new(
398                provider,
399                client,
400                service,
401                retry_config.clone(),
402            )) as _
403        } else if has_env_oidc {
404            // 4. env-oidc
405            let env_var = self
406                .oidc_token_env
407                .unwrap_or_else(|| DEFAULT_OIDC_TOKEN_ENV.to_owned());
408            let provider = OidcEnvTokenProvider { env_var };
409            let client = client_options.client()?;
410            let service = make_service(client.clone(), runtime);
411            Arc::new(TokenCredentialProvider::new(
412                provider,
413                client,
414                service,
415                retry_config.clone(),
416            )) as _
417        } else if has_file_oidc {
418            // 5. file-oidc
419            let provider = OidcFileTokenProvider {
420                filepath: self.oidc_token_filepath.unwrap(),
421            };
422            let client = client_options.client()?;
423            let service = make_service(client.clone(), runtime);
424            Arc::new(TokenCredentialProvider::new(
425                provider,
426                client,
427                service,
428                retry_config.clone(),
429            )) as _
430        } else if has_azure_msi {
431            // 6. Azure MSI fallback
432            let msi_provider = ImdsManagedIdentityProvider::new_with_resource(
433                None,
434                None,
435                None,
436                None,
437                AZURE_DATABRICKS_RESOURCE,
438            );
439            let client = client_options.metadata_client()?;
440            let service = make_service(client.clone(), runtime);
441            let azure_provider =
442                TokenCredentialProvider::new(msi_provider, client, service, retry_config.clone());
443            Arc::new(AzureToDatabricksBridge {
444                inner: Arc::new(azure_provider),
445            }) as _
446        } else if has_gcp {
447            // 7. GCP SA token exchange
448            let host = self.host.as_deref().unwrap_or("").trim_end_matches('/');
449            let token_url = format!("{host}/oidc/v1/token");
450            let gcp_config = crate::gcp::GoogleBuilder::new().build(runtime)?;
451            let provider = DatabricksGcpTokenExchangeProvider {
452                token_url,
453                gcp_provider: gcp_config.credentials,
454            };
455            let client = client_options.client()?;
456            let service = make_service(client.clone(), runtime);
457            Arc::new(TokenCredentialProvider::new(
458                provider,
459                client,
460                service,
461                retry_config.clone(),
462            )) as _
463        } else {
464            return Err(crate::Error::Generic {
465                source: "No Databricks credentials configured. Set DATABRICKS_TOKEN, \
466                         DATABRICKS_CLIENT_ID + DATABRICKS_CLIENT_SECRET, \
467                         DATABRICKS_OIDC_TOKEN (env-oidc), DATABRICKS_OIDC_TOKEN_FILEPATH (file-oidc), \
468                         DATABRICKS_AZURE_RESOURCE_ID, or GOOGLE_APPLICATION_CREDENTIALS."
469                    .into(),
470            });
471        };
472
473        Ok(DatabricksConfig {
474            credentials,
475            retry_config,
476            client_options,
477        })
478    }
479
480    /// Build credentials when `auth_type` forces a specific method.
481    #[allow(clippy::too_many_arguments)]
482    fn build_forced_auth(
483        self,
484        auth_type: String,
485        scope: String,
486        runtime: Option<&Handle>,
487        has_pat: bool,
488        has_m2m: bool,
489        has_env_oidc: bool,
490        has_file_oidc: bool,
491        has_azure_msi: bool,
492        has_gcp: bool,
493    ) -> Result<DatabricksCredentialProvider> {
494        match auth_type.as_str() {
495            "pat" => {
496                if !has_pat {
497                    return Err(crate::Error::Generic {
498                        source: "auth_type=pat requires DATABRICKS_TOKEN to be set".into(),
499                    });
500                }
501                Ok(Arc::new(StaticCredentialProvider::new(DatabricksCredential {
502                    bearer: self.token.unwrap(),
503                })))
504            }
505            "oauth-m2m" => {
506                if !has_m2m {
507                    return Err(crate::Error::Generic {
508                        source: "auth_type=oauth-m2m requires DATABRICKS_CLIENT_ID and \
509                                 DATABRICKS_CLIENT_SECRET to be set"
510                            .into(),
511                    });
512                }
513                let host = self.host.as_deref().unwrap_or("").trim_end_matches('/');
514                let token_url = format!("{host}/oidc/v1/token");
515                let provider = DatabricksM2MProvider {
516                    token_url,
517                    client_id: self.client_id.unwrap(),
518                    client_secret: self.client_secret.unwrap(),
519                    scope,
520                };
521                let client = self.client_options.client()?;
522                let service = make_service(client.clone(), runtime);
523                Ok(Arc::new(TokenCredentialProvider::new(
524                    provider,
525                    client,
526                    service,
527                    self.retry_config,
528                )) as _)
529            }
530            "env-oidc" => {
531                if !has_env_oidc {
532                    return Err(crate::Error::Generic {
533                        source: "auth_type=env-oidc requires DATABRICKS_OIDC_TOKEN or \
534                                 DATABRICKS_OIDC_TOKEN_ENV to be set"
535                            .into(),
536                    });
537                }
538                let env_var = self
539                    .oidc_token_env
540                    .unwrap_or_else(|| DEFAULT_OIDC_TOKEN_ENV.to_owned());
541                let provider = OidcEnvTokenProvider { env_var };
542                let client = self.client_options.client()?;
543                let service = make_service(client.clone(), runtime);
544                Ok(Arc::new(TokenCredentialProvider::new(
545                    provider,
546                    client,
547                    service,
548                    self.retry_config,
549                )) as _)
550            }
551            "file-oidc" => {
552                if !has_file_oidc {
553                    return Err(crate::Error::Generic {
554                        source: "auth_type=file-oidc requires DATABRICKS_OIDC_TOKEN_FILEPATH \
555                                 to be set"
556                            .into(),
557                    });
558                }
559                let provider = OidcFileTokenProvider {
560                    filepath: self.oidc_token_filepath.unwrap(),
561                };
562                let client = self.client_options.client()?;
563                let service = make_service(client.clone(), runtime);
564                Ok(Arc::new(TokenCredentialProvider::new(
565                    provider,
566                    client,
567                    service,
568                    self.retry_config,
569                )) as _)
570            }
571            "azure-msi" => {
572                if !has_azure_msi {
573                    return Err(crate::Error::Generic {
574                        source: "auth_type=azure-msi requires DATABRICKS_AZURE_RESOURCE_ID \
575                                 to be set"
576                            .into(),
577                    });
578                }
579                let msi_provider = ImdsManagedIdentityProvider::new_with_resource(
580                    None,
581                    None,
582                    None,
583                    None,
584                    AZURE_DATABRICKS_RESOURCE,
585                );
586                let client = self.client_options.metadata_client()?;
587                let service = make_service(client.clone(), runtime);
588                let azure_provider = TokenCredentialProvider::new(
589                    msi_provider,
590                    client,
591                    service,
592                    self.retry_config,
593                );
594                Ok(Arc::new(AzureToDatabricksBridge {
595                    inner: Arc::new(azure_provider),
596                }) as _)
597            }
598            "gcp" => {
599                if !has_gcp {
600                    return Err(crate::Error::Generic {
601                        source: "auth_type=gcp requires GOOGLE_APPLICATION_CREDENTIALS to be set"
602                            .into(),
603                    });
604                }
605                let host = self.host.as_deref().unwrap_or("").trim_end_matches('/');
606                let token_url = format!("{host}/oidc/v1/token");
607                let gcp_config = crate::gcp::GoogleBuilder::new().build(runtime)?;
608                let provider = DatabricksGcpTokenExchangeProvider {
609                    token_url,
610                    gcp_provider: gcp_config.credentials,
611                };
612                let client = self.client_options.client()?;
613                let service = make_service(client.clone(), runtime);
614                Ok(Arc::new(TokenCredentialProvider::new(
615                    provider,
616                    client,
617                    service,
618                    self.retry_config,
619                )) as _)
620            }
621            other => Err(crate::Error::Generic {
622                source: format!("Unknown auth_type: {other:?}. Valid values: pat, oauth-m2m, env-oidc, file-oidc, azure-msi, gcp").into(),
623            }),
624        }
625    }
626}
627
628/// Bridges `AzureCredential::BearerToken` → `DatabricksCredential`.
629///
630/// Azure MSI returns an Azure AD bearer token for the Databricks resource ID;
631/// that token is used directly as the Databricks API bearer token.
632#[derive(Debug)]
633struct AzureToDatabricksBridge {
634    inner: Arc<dyn CredentialProvider<Credential = crate::azure::credential::AzureCredential>>,
635}
636
637#[async_trait::async_trait]
638impl CredentialProvider for AzureToDatabricksBridge {
639    type Credential = DatabricksCredential;
640
641    async fn get_credential(&self) -> crate::Result<Arc<DatabricksCredential>> {
642        use crate::azure::credential::AzureCredential;
643        let cred = self.inner.get_credential().await?;
644        let AzureCredential::BearerToken(token) = cred.as_ref();
645        Ok(Arc::new(DatabricksCredential {
646            bearer: token.clone(),
647        }))
648    }
649}
650
651#[cfg(test)]
652mod tests {
653    use std::collections::HashMap;
654
655    use super::*;
656
657    #[test]
658    fn test_databricks_config_from_map() {
659        let options = HashMap::from([
660            (
661                "databricks_host",
662                "https://my-workspace.azuredatabricks.net",
663            ),
664            ("databricks_token", "dapi1234567890"),
665            ("databricks_oauth_scope", "all-apis"),
666        ]);
667
668        let builder = options
669            .into_iter()
670            .fold(DatabricksBuilder::new(), |b, (k, v)| {
671                b.with_config(k.parse().unwrap(), v)
672            });
673
674        assert_eq!(
675            builder.host.as_deref(),
676            Some("https://my-workspace.azuredatabricks.net")
677        );
678        assert_eq!(builder.token.as_deref(), Some("dapi1234567890"));
679        assert_eq!(builder.scope.as_deref(), Some("all-apis"));
680    }
681
682    #[test]
683    fn test_databricks_config_key_roundtrip() {
684        let cases = [
685            ("databricks_host", DatabricksConfigKey::Host),
686            ("host", DatabricksConfigKey::Host),
687            ("databricks_token", DatabricksConfigKey::Token),
688            ("token", DatabricksConfigKey::Token),
689            ("databricks_client_id", DatabricksConfigKey::ClientId),
690            ("client_id", DatabricksConfigKey::ClientId),
691            (
692                "databricks_client_secret",
693                DatabricksConfigKey::ClientSecret,
694            ),
695            ("client_secret", DatabricksConfigKey::ClientSecret),
696            ("databricks_account_id", DatabricksConfigKey::AccountId),
697            ("account_id", DatabricksConfigKey::AccountId),
698            (
699                "databricks_azure_resource_id",
700                DatabricksConfigKey::AzureResourceId,
701            ),
702            ("azure_resource_id", DatabricksConfigKey::AzureResourceId),
703            ("databricks_oauth_scope", DatabricksConfigKey::Scope),
704            ("oauth_scope", DatabricksConfigKey::Scope),
705            ("databricks_config_file", DatabricksConfigKey::ConfigFile),
706            ("config_file", DatabricksConfigKey::ConfigFile),
707            (
708                "databricks_config_profile",
709                DatabricksConfigKey::ConfigProfile,
710            ),
711            ("config_profile", DatabricksConfigKey::ConfigProfile),
712            ("profile", DatabricksConfigKey::ConfigProfile),
713            ("databricks_auth_type", DatabricksConfigKey::AuthType),
714            ("auth_type", DatabricksConfigKey::AuthType),
715            (
716                "databricks_oidc_token_env",
717                DatabricksConfigKey::OidcTokenEnv,
718            ),
719            ("oidc_token_env", DatabricksConfigKey::OidcTokenEnv),
720            (
721                "databricks_oidc_token_filepath",
722                DatabricksConfigKey::OidcTokenFilepath,
723            ),
724            (
725                "oidc_token_filepath",
726                DatabricksConfigKey::OidcTokenFilepath,
727            ),
728        ];
729
730        for (s, expected) in cases {
731            assert_eq!(
732                s.parse::<DatabricksConfigKey>().unwrap(),
733                expected,
734                "failed for {s}"
735            );
736        }
737    }
738
739    #[tokio::test]
740    async fn test_build_with_pat() {
741        let config = DatabricksBuilder::new()
742            .with_host("https://my.azuredatabricks.net")
743            .with_token("dapimytoken")
744            .build(None)
745            .unwrap();
746
747        let cred = config.credentials.get_credential().await.unwrap();
748        assert_eq!(cred.bearer, "dapimytoken");
749    }
750
751    #[test]
752    fn test_auth_type_enforcement_missing_fields() {
753        let err = DatabricksBuilder::new()
754            .with_auth_type("pat")
755            .build(None)
756            .unwrap_err();
757        assert!(err.to_string().contains("pat"));
758    }
759}
olai_http/databricks/builder.rs

olai_http/databricks/
builder.rs
