1use std::{env, fmt, path::PathBuf};
2
3use configparser::ini::Ini;
4use directories::UserDirs;
5use sha2::{Digest, Sha256};
6use url::Url;
7
8use crate::{Error, Result};
9
10pub const DEFAULT_CLIENT_ID: &str = "databricks-cli";
11pub const DEFAULT_ACCOUNTS_HOST: &str = "https://accounts.cloud.databricks.com";
12pub const DEFAULT_CONFIG_FILE: &str = "~/.databrickscfg";
13const SETTINGS_SECTION: &str = "__settings__";
14const AUTH_TYPE_DATABRICKS_CLI: &str = "databricks-cli";
15const AUTH_TYPE_M2M: &str = "oauth-m2m";
16
17#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
19pub enum AuthKind {
20 #[default]
22 UserToMachine,
23 MachineToMachine,
25}
26
27#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
28pub enum TargetKind {
29 #[default]
30 Workspace,
31 Account,
32 Unified,
33}
34
35#[derive(Clone)]
36pub struct Profile {
37 pub name: String,
38 pub host: Url,
39 pub account_id: Option<String>,
40 pub workspace_id: Option<String>,
41 pub client_id: String,
42 pub group_id: Option<String>,
44 pub scopes: Vec<String>,
45 pub target: TargetKind,
46 pub auth_kind: AuthKind,
48 pub(crate) client_secret: Option<String>,
49}
50
51impl Profile {
52 pub fn from_sources(options: ProfileOptions) -> Result<Self> {
53 let config_file = resolve_config_file(options.config_file.as_deref())?;
54 let prefer_user_to_machine = options.prefer_user_to_machine;
55 let environment_profile = env_nonempty("DATABRICKS_CONFIG_PROFILE");
56 let explicit_profile = options.profile.is_some() || environment_profile.is_some();
57 let requested_name = options
58 .profile
59 .or(environment_profile)
60 .map(|value| value.trim().to_owned())
61 .filter(|value| !value.is_empty());
62 let config = load_config(&config_file)?;
63 let profile_name = resolve_auth_profile_name(
64 requested_name.as_deref(),
65 explicit_profile,
66 config.as_ref(),
67 prefer_user_to_machine,
68 )?;
69 let configured = config
70 .as_ref()
71 .map(|config| load_profile(config, &profile_name))
72 .unwrap_or_default();
73
74 let host = options
75 .host
76 .or_else(|| env_nonempty("DATABRICKS_HOST"))
77 .or(configured.host)
78 .ok_or_else(|| Error::Config(format!("profile {profile_name} has no host")))?;
79 let host = normalize_host(&host)?;
80 let account_id = options
81 .account_id
82 .or_else(|| env_nonempty("DATABRICKS_ACCOUNT_ID"))
83 .or(configured.account_id);
84 let workspace_id = options
85 .workspace_id
86 .or_else(|| env_nonempty("DATABRICKS_WORKSPACE_ID"))
87 .or(configured.workspace_id);
88 let client_id = options
89 .client_id
90 .or_else(|| env_nonempty("DATABRICKS_CLIENT_ID"))
91 .or(configured.client_id);
92 let client_secret = options
93 .client_secret
94 .or_else(|| env_nonempty("DATABRICKS_CLIENT_SECRET"))
95 .or(configured.client_secret);
96 let group_id = options
97 .group_id
98 .or_else(|| env_nonempty("DATABRICKS_GROUP_ID"))
99 .or(configured.group_id);
100 let auth_type = options
101 .auth_type
102 .or_else(|| env_nonempty("DATABRICKS_AUTH_TYPE"))
103 .or(configured.auth_type);
104 let auth_kind = resolve_auth_kind(
105 auth_type.as_deref(),
106 client_id.as_deref(),
107 client_secret.as_deref(),
108 )?;
109 let client_id = match auth_kind {
110 AuthKind::UserToMachine => client_id.unwrap_or_else(|| DEFAULT_CLIENT_ID.to_owned()),
111 AuthKind::MachineToMachine => client_id.ok_or_else(|| {
112 Error::Config(format!(
113 "profile {profile_name} requires client_id for oauth-m2m"
114 ))
115 })?,
116 };
117 let scopes = options
118 .scopes
119 .or_else(|| configured.scopes.map(split_list))
120 .unwrap_or_else(|| vec!["all-apis".to_owned()]);
121 let target = options.target.unwrap_or_else(|| {
122 if account_id.is_some() && host.host_str() == Some("accounts.cloud.databricks.com") {
123 TargetKind::Account
124 } else {
125 TargetKind::Workspace
126 }
127 });
128
129 Ok(Self {
130 name: profile_name,
131 host,
132 account_id,
133 workspace_id,
134 client_id,
135 group_id,
136 scopes,
137 target,
138 auth_kind,
139 client_secret,
140 })
141 }
142
143 pub fn cache_key(&self) -> String {
144 match self.auth_kind {
145 AuthKind::UserToMachine => self.name.clone(),
146 AuthKind::MachineToMachine => {
147 let scopes = self.machine_scopes();
148 let identity = format!(
149 "{}\0{}\0{}\0{}\0{}\0{}",
150 self.host,
151 self.account_id.as_deref().unwrap_or_default(),
152 self.workspace_id.as_deref().unwrap_or_default(),
153 self.client_id,
154 self.group_id.as_deref().unwrap_or_default(),
155 scopes.join(" "),
156 );
157 format!(
158 "{}-oauth-m2m-{:x}",
159 self.name,
160 Sha256::digest(identity.as_bytes())
161 )
162 }
163 }
164 }
165
166 pub(crate) fn client_secret(&self) -> Option<&str> {
167 self.client_secret.as_deref()
168 }
169
170 pub fn effective_scopes(&self) -> Vec<String> {
171 let mut scopes = vec!["offline_access".to_owned()];
172 for scope in &self.scopes {
173 if !scopes.contains(scope) {
174 scopes.push(scope.clone());
175 }
176 }
177 scopes
178 }
179
180 pub fn machine_scopes(&self) -> Vec<String> {
181 let mut scopes = self.scopes.clone();
182 if scopes.is_empty() {
183 scopes.push("all-apis".to_owned());
184 }
185 scopes.sort();
186 scopes.dedup();
187 scopes
188 }
189}
190
191impl fmt::Debug for Profile {
192 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
193 formatter
194 .debug_struct("Profile")
195 .field("name", &self.name)
196 .field("host", &self.host)
197 .field("account_id", &self.account_id)
198 .field("workspace_id", &self.workspace_id)
199 .field("client_id", &self.client_id)
200 .field("group_id", &self.group_id)
201 .field("scopes", &self.scopes)
202 .field("target", &self.target)
203 .field("auth_kind", &self.auth_kind)
204 .field(
205 "client_secret",
206 &self.client_secret.as_ref().map(|_| "[REDACTED]"),
207 )
208 .finish()
209 }
210}
211
212#[derive(Clone, Debug)]
213pub struct ProfileOptions {
214 pub profile: Option<String>,
215 pub host: Option<String>,
216 pub account_id: Option<String>,
217 pub workspace_id: Option<String>,
218 pub client_id: Option<String>,
219 pub client_secret: Option<String>,
221 pub group_id: Option<String>,
223 pub auth_type: Option<String>,
225 pub scopes: Option<Vec<String>>,
226 pub target: Option<TargetKind>,
227 pub config_file: Option<PathBuf>,
228 pub prefer_user_to_machine: bool,
230}
231
232impl Default for ProfileOptions {
233 fn default() -> Self {
234 Self {
235 profile: None,
236 host: None,
237 account_id: None,
238 workspace_id: None,
239 client_id: None,
240 client_secret: None,
241 group_id: None,
242 auth_type: None,
243 scopes: None,
244 target: None,
245 config_file: None,
246 prefer_user_to_machine: true,
247 }
248 }
249}
250
251#[derive(Clone, Debug, Default)]
252struct RawProfile {
253 host: Option<String>,
254 account_id: Option<String>,
255 workspace_id: Option<String>,
256 client_id: Option<String>,
257 client_secret: Option<String>,
258 group_id: Option<String>,
259 scopes: Option<String>,
260 auth_type: Option<String>,
261}
262
263pub fn resolve_config_file(explicit: Option<&std::path::Path>) -> Result<PathBuf> {
264 let path = explicit
265 .map(PathBuf::from)
266 .or_else(|| env_nonempty("DATABRICKS_CONFIG_FILE").map(PathBuf::from))
267 .unwrap_or_else(|| PathBuf::from(DEFAULT_CONFIG_FILE));
268 expand_home(path)
269}
270
271fn expand_home(path: PathBuf) -> Result<PathBuf> {
272 let value = path.to_string_lossy();
273 if value == "~" || value.starts_with("~/") || value.starts_with("~\\") {
274 let home = UserDirs::new()
275 .map(|dirs| dirs.home_dir().to_path_buf())
276 .ok_or_else(|| Error::Config("cannot find home directory".into()))?;
277 if value == "~" {
278 return Ok(home);
279 }
280 return Ok(home.join(&value[2..]));
281 }
282 Ok(path)
283}
284
285fn load_config(path: &std::path::Path) -> Result<Option<Ini>> {
286 if !path.exists() {
287 return Ok(None);
288 }
289 let mut ini = Ini::new_cs();
290 ini.load(path)
291 .map_err(|error| Error::Config(format!("could not read {}: {error}", path.display())))?;
292 Ok(Some(ini))
293}
294
295fn resolve_auth_profile_name(
296 requested: Option<&str>,
297 explicit: bool,
298 config: Option<&Ini>,
299 prefer_user_to_machine: bool,
300) -> Result<String> {
301 let selected = resolve_profile_name(requested, config)?;
302 if explicit || !prefer_user_to_machine {
303 return Ok(selected);
304 }
305 let Some(config) = config else {
306 return Ok(selected);
307 };
308 let selected_profile = load_profile(config, &selected);
309 if !is_m2m_profile(&selected_profile) {
310 return Ok(selected);
311 }
312 if selected_profile.host.is_none() {
313 return Ok(selected);
314 };
315 let mut matches = config
316 .sections()
317 .into_iter()
318 .filter(|name| name != SETTINGS_SECTION && name != &selected)
319 .filter(|name| {
320 let profile = load_profile(config, name);
321 profile.auth_type.as_deref().is_some_and(is_u2m_auth_type)
322 && same_auth_target(&selected_profile, &profile)
323 });
324 let Some(profile) = matches.next() else {
325 return Ok(selected);
326 };
327 if matches.next().is_some() {
328 return Ok(selected);
329 }
330 Ok(profile)
331}
332
333fn is_u2m_auth_type(auth_type: &str) -> bool {
334 auth_type == AUTH_TYPE_DATABRICKS_CLI
335}
336
337fn is_m2m_profile(profile: &RawProfile) -> bool {
338 profile.auth_type.as_deref() == Some(AUTH_TYPE_M2M)
339 || (profile.auth_type.is_none()
340 && profile.client_id.is_some()
341 && profile.client_secret.is_some())
342}
343
344fn same_auth_target(selected: &RawProfile, candidate: &RawProfile) -> bool {
345 let hosts_match = selected
346 .host
347 .as_deref()
348 .and_then(|host| normalize_host(host).ok())
349 .zip(
350 candidate
351 .host
352 .as_deref()
353 .and_then(|host| normalize_host(host).ok()),
354 )
355 .is_some_and(|(selected, candidate)| selected == candidate);
356 hosts_match
357 && selected
358 .account_id
359 .as_ref()
360 .is_none_or(|account_id| candidate.account_id.as_ref() == Some(account_id))
361 && selected
362 .workspace_id
363 .as_ref()
364 .is_none_or(|workspace_id| candidate.workspace_id.as_ref() == Some(workspace_id))
365}
366
367fn resolve_auth_kind(
368 auth_type: Option<&str>,
369 client_id: Option<&str>,
370 client_secret: Option<&str>,
371) -> Result<AuthKind> {
372 match auth_type {
373 Some(AUTH_TYPE_DATABRICKS_CLI) => Ok(AuthKind::UserToMachine),
374 Some(AUTH_TYPE_M2M) => {
375 if client_id.is_none() || client_secret.is_none() {
376 return Err(Error::Config(
377 "oauth-m2m requires client_id and client_secret".into(),
378 ));
379 }
380 Ok(AuthKind::MachineToMachine)
381 }
382 Some(auth_type) => Err(Error::Config(format!(
383 "authentication type {auth_type} is not supported"
384 ))),
385 None if client_id.is_some() && client_secret.is_some() => Ok(AuthKind::MachineToMachine),
386 None if client_secret.is_some() => Err(Error::Config(
387 "oauth-m2m client_secret requires client_id".into(),
388 )),
389 None => Ok(AuthKind::UserToMachine),
390 }
391}
392
393fn resolve_profile_name(requested: Option<&str>, config: Option<&Ini>) -> Result<String> {
394 if let Some(profile) = requested {
395 if profile == SETTINGS_SECTION {
396 return Err(Error::Config(format!(
397 "{SETTINGS_SECTION} is a reserved section name and cannot be used as a profile"
398 )));
399 }
400 return Ok(profile.to_owned());
401 }
402 if let Some(profile) = config
403 .and_then(|config| config.get(SETTINGS_SECTION, "default_profile"))
404 .map(|value| value.trim().to_owned())
405 .filter(|value| !value.is_empty())
406 {
407 if profile == SETTINGS_SECTION {
408 return Err(Error::Config(format!(
409 "{SETTINGS_SECTION} is a reserved section name and cannot be used as a profile"
410 )));
411 }
412 return Ok(profile);
413 }
414 Ok("DEFAULT".to_owned())
415}
416
417fn load_profile(ini: &Ini, name: &str) -> RawProfile {
418 RawProfile {
419 host: ini.get(name, "host"),
420 account_id: ini.get(name, "account_id"),
421 workspace_id: ini.get(name, "workspace_id"),
422 client_id: ini.get(name, "client_id"),
423 client_secret: ini.get(name, "client_secret"),
424 group_id: ini.get(name, "group_id"),
425 scopes: ini.get(name, "scopes"),
426 auth_type: ini
427 .get(name, "auth_type")
428 .map(|value| value.trim().to_ascii_lowercase())
429 .filter(|value| !value.is_empty()),
430 }
431}
432
433pub(crate) fn configured_auth_storage(path: &std::path::Path) -> Result<Option<String>> {
434 Ok(load_config(path)?
435 .and_then(|config| config.get(SETTINGS_SECTION, "auth_storage"))
436 .map(|value| value.trim().to_owned())
437 .filter(|value| !value.is_empty()))
438}
439
440fn normalize_host(value: &str) -> Result<Url> {
441 let value = value.trim().trim_end_matches('/');
442 let value = if value.starts_with("http://") || value.starts_with("https://") {
443 value.to_owned()
444 } else {
445 format!("https://{value}")
446 };
447 let url = Url::parse(&value)?;
448 if url.scheme() != "https"
449 && url
450 .host_str()
451 .is_none_or(|host| host != "127.0.0.1" && host != "localhost")
452 {
453 return Err(Error::Config("Databricks host must use HTTPS".into()));
454 }
455 Ok(url)
456}
457
458fn env_nonempty(name: &str) -> Option<String> {
459 env::var(name)
460 .ok()
461 .map(|value| value.trim().to_owned())
462 .filter(|value| !value.is_empty())
463}
464
465fn split_list(value: String) -> Vec<String> {
466 value
467 .split(',')
468 .map(str::trim)
469 .filter(|value| !value.is_empty())
470 .map(str::to_owned)
471 .collect()
472}
473
474#[cfg(test)]
475mod tests {
476 use super::*;
477
478 #[test]
479 fn configured_default_precedes_legacy_default() {
480 let mut config = Ini::new();
481 config.set(SETTINGS_SECTION, "default_profile", Some("selected".into()));
482 assert_eq!(
483 resolve_profile_name(None, Some(&config)).unwrap(),
484 "selected"
485 );
486 }
487
488 #[test]
489 fn default_is_the_legacy_fallback() {
490 assert_eq!(resolve_profile_name(None, None).unwrap(), "DEFAULT");
491 }
492
493 #[test]
494 fn settings_is_not_a_profile() {
495 assert!(resolve_profile_name(Some(SETTINGS_SECTION), None).is_err());
496 }
497
498 #[test]
499 fn reads_profile_auth_type() {
500 let mut config = Ini::new();
501 config.set("service", "auth_type", Some("oauth-m2m".into()));
502 config.set("service", "client_id", Some("client".into()));
503 config.set("service", "client_secret", Some("secret".into()));
504 config.set("service", "group_id", Some("group".into()));
505 let profile = load_profile(&config, "service");
506 assert_eq!(profile.auth_type.as_deref(), Some("oauth-m2m"));
507 assert_eq!(profile.client_id.as_deref(), Some("client"));
508 assert_eq!(profile.client_secret.as_deref(), Some("secret"));
509 assert_eq!(profile.group_id.as_deref(), Some("group"));
510 }
511
512 fn mixed_auth_config() -> Ini {
513 let mut config = Ini::new_cs();
514 config.set(SETTINGS_SECTION, "default_profile", Some("DEFAULT".into()));
515 config.set("DEFAULT", "host", Some("https://workspace.example".into()));
516 config.set("DEFAULT", "auth_type", Some("oauth-m2m".into()));
517 config.set("FEVM-AWS", "host", Some("https://workspace.example".into()));
518 config.set("FEVM-AWS", "auth_type", Some("databricks-cli".into()));
519 config
520 }
521
522 #[test]
523 fn implicit_m2m_default_maps_to_unique_u2m_profile_on_same_host() {
524 let config = mixed_auth_config();
525 assert_eq!(
526 resolve_auth_profile_name(None, false, Some(&config), true).unwrap(),
527 "FEVM-AWS"
528 );
529 }
530
531 #[test]
532 fn disabled_u2m_preference_keeps_the_m2m_default() {
533 let config = mixed_auth_config();
534 assert_eq!(
535 resolve_auth_profile_name(None, false, Some(&config), false).unwrap(),
536 "DEFAULT"
537 );
538 }
539
540 #[test]
541 fn implicit_credentials_without_auth_type_still_prefer_u2m() {
542 let mut config = mixed_auth_config();
543 config.remove_key("DEFAULT", "auth_type");
544 config.set("DEFAULT", "client_id", Some("client".into()));
545 config.set("DEFAULT", "client_secret", Some("secret".into()));
546 assert_eq!(
547 resolve_auth_profile_name(None, false, Some(&config), true).unwrap(),
548 "FEVM-AWS"
549 );
550 }
551
552 #[test]
553 fn only_databricks_cli_is_u2m_compatible() {
554 assert!(is_u2m_auth_type(AUTH_TYPE_DATABRICKS_CLI));
555 assert!(!is_u2m_auth_type(AUTH_TYPE_M2M));
556 assert!(!is_u2m_auth_type("external-browser"));
557 }
558
559 #[test]
560 fn non_m2m_default_is_not_remapped() {
561 let mut config = mixed_auth_config();
562 config.set("DEFAULT", "auth_type", Some("pat".into()));
563 assert_eq!(
564 resolve_auth_profile_name(None, false, Some(&config), true).unwrap(),
565 "DEFAULT"
566 );
567 }
568
569 #[test]
570 fn explicit_m2m_profile_is_not_remapped() {
571 let config = mixed_auth_config();
572 assert_eq!(
573 resolve_auth_profile_name(Some("DEFAULT"), true, Some(&config), true).unwrap(),
574 "DEFAULT"
575 );
576 }
577
578 #[test]
579 fn ambiguous_u2m_host_match_does_not_remap() {
580 let mut config = mixed_auth_config();
581 config.set(
582 "FEVM-AWS-2",
583 "host",
584 Some("https://workspace.example".into()),
585 );
586 config.set("FEVM-AWS-2", "auth_type", Some("databricks-cli".into()));
587 assert_eq!(
588 resolve_auth_profile_name(None, false, Some(&config), true).unwrap(),
589 "DEFAULT"
590 );
591 }
592
593 #[test]
594 fn different_host_does_not_remap() {
595 let mut config = mixed_auth_config();
596 config.set("FEVM-AWS", "host", Some("https://other.example".into()));
597 assert_eq!(
598 resolve_auth_profile_name(None, false, Some(&config), true).unwrap(),
599 "DEFAULT"
600 );
601 }
602
603 #[test]
604 fn equivalent_hosts_match_after_normalization() {
605 let mut config = mixed_auth_config();
606 config.set("DEFAULT", "host", Some("workspace.example/".into()));
607 config.set(
608 "FEVM-AWS",
609 "host",
610 Some("https://workspace.example".into()),
611 );
612 assert_eq!(
613 resolve_auth_profile_name(None, false, Some(&config), true).unwrap(),
614 "FEVM-AWS"
615 );
616 }
617
618 #[test]
619 fn different_account_does_not_remap() {
620 let mut config = mixed_auth_config();
621 config.set("DEFAULT", "account_id", Some("account-a".into()));
622 config.set("FEVM-AWS", "account_id", Some("account-b".into()));
623 assert_eq!(
624 resolve_auth_profile_name(None, false, Some(&config), true).unwrap(),
625 "DEFAULT"
626 );
627 }
628
629 #[test]
630 fn resolves_explicit_and_default_m2m_credentials() {
631 assert_eq!(
632 resolve_auth_kind(Some(AUTH_TYPE_M2M), Some("client"), Some("secret")).unwrap(),
633 AuthKind::MachineToMachine
634 );
635 assert_eq!(
636 resolve_auth_kind(None, Some("client"), Some("secret")).unwrap(),
637 AuthKind::MachineToMachine
638 );
639 assert_eq!(
640 resolve_auth_kind(None, Some("client"), None).unwrap(),
641 AuthKind::UserToMachine
642 );
643 assert!(resolve_auth_kind(Some(AUTH_TYPE_M2M), Some("client"), None).is_err());
644 assert!(resolve_auth_kind(None, None, Some("secret")).is_err());
645 }
646
647 #[test]
648 fn explicit_m2m_profile_builds_without_browser_auth() {
649 let directory = tempfile::tempdir().unwrap();
650 let profile = Profile::from_sources(ProfileOptions {
651 profile: Some("service".into()),
652 host: Some("http://127.0.0.1:8080".into()),
653 client_id: Some("client".into()),
654 client_secret: Some("secret".into()),
655 auth_type: Some(AUTH_TYPE_M2M.into()),
656 config_file: Some(directory.path().join("missing")),
657 ..ProfileOptions::default()
658 })
659 .unwrap();
660 assert_eq!(profile.auth_kind, AuthKind::MachineToMachine);
661 assert_eq!(profile.client_secret(), Some("secret"));
662 }
663
664 #[test]
665 fn client_credentials_inference_does_not_depend_on_profile_preference() {
666 let directory = tempfile::tempdir().unwrap();
667 let options = |prefer_user_to_machine| ProfileOptions {
668 profile: Some("service".into()),
669 host: Some("http://127.0.0.1:8080".into()),
670 client_id: Some("client".into()),
671 client_secret: Some("secret".into()),
672 config_file: Some(directory.path().join("missing")),
673 prefer_user_to_machine,
674 ..ProfileOptions::default()
675 };
676 assert_eq!(
677 Profile::from_sources(options(true)).unwrap().auth_kind,
678 AuthKind::MachineToMachine
679 );
680 assert_eq!(
681 Profile::from_sources(options(false)).unwrap().auth_kind,
682 AuthKind::MachineToMachine
683 );
684 }
685
686 #[test]
687 fn m2m_cache_keys_include_client_group_and_scopes() {
688 let profile = Profile {
689 name: "service".into(),
690 host: Url::parse("https://workspace.example").unwrap(),
691 account_id: None,
692 workspace_id: None,
693 client_id: "client".into(),
694 group_id: Some("group".into()),
695 scopes: vec!["jobs".into(), "files:read".into()],
696 target: TargetKind::Workspace,
697 auth_kind: AuthKind::MachineToMachine,
698 client_secret: Some("credential-value".into()),
699 };
700 let key = profile.cache_key();
701 assert!(key.starts_with("service-oauth-m2m-"));
702 assert!(!key.contains("credential-value"));
703 assert!(!format!("{profile:?}").contains("credential-value"));
704 }
705}