dbx_tools_databricks_auth/
client.rs

1use crate::{
2    AuthKind, AuthSession, CredentialStore, Error, MachineToMachineFlow, OAuthFlow, OAuthTemplate,
3    Profile, Result, Token,
4};
5pub use dbx_tools_auth::AuthOptions;
6use std::{sync::Arc, time::Duration};
7
8/// Databricks profile and acquisition policy over the shared `AuthSession` lifecycle.
9pub struct AuthClient {
10    profile: Profile,
11    inner: dbx_tools_auth::AuthClient,
12}
13
14enum AuthFlow {
15    UserToMachine(OAuthFlow),
16    UserToMachineCli(DatabricksCliFlow),
17    MachineToMachine(MachineToMachineFlow),
18}
19
20struct DatabricksCliFlow {
21    native: OAuthFlow,
22    profile: String,
23}
24
25impl DatabricksCliFlow {
26    fn new(native: OAuthFlow, profile: String) -> Self {
27        Self { native, profile }
28    }
29
30    async fn token(&self, force_refresh: bool) -> Result<Token> {
31        let profile = self.profile.clone();
32        tokio::task::spawn_blocking(move || {
33            let output = dbx_tools_databricks::databricks_cli_token(&profile, force_refresh)
34                .map_err(|error| Error::OAuth(error.to_string()))?;
35            serde_json::from_slice(&output).map_err(Into::into)
36        })
37        .await
38        .map_err(|error| Error::OAuth(format!("databricks auth token task failed: {error}")))?
39    }
40}
41
42impl AuthClient {
43    pub fn new(
44        profile: Profile,
45        store: Arc<dyn CredentialStore>,
46        options: AuthOptions,
47        use_databricks_cli: bool,
48    ) -> Result<Self> {
49        let flow = match profile.auth_kind {
50            AuthKind::UserToMachine => {
51                let native = OAuthFlow::new(profile.clone())?
52                    .with_template(OAuthTemplate::new(options.callback_image_src.clone()));
53                if use_databricks_cli {
54                    AuthFlow::UserToMachineCli(DatabricksCliFlow::new(native, profile.name.clone()))
55                } else {
56                    AuthFlow::UserToMachine(native)
57                }
58            }
59            AuthKind::MachineToMachine => {
60                AuthFlow::MachineToMachine(MachineToMachineFlow::new(profile.clone())?)
61            }
62        };
63        let inner =
64            dbx_tools_auth::AuthClient::new(profile.cache_key(), Arc::new(flow), store, options);
65        Ok(Self { profile, inner })
66    }
67    pub fn profile(&self) -> &Profile {
68        &self.profile
69    }
70}
71
72impl AuthSession for AuthClient {
73    fn auth_client(&self) -> &dbx_tools_auth::AuthClient {
74        &self.inner
75    }
76}
77
78#[async_trait::async_trait]
79impl dbx_tools_auth::TokenProvider for AuthFlow {
80    async fn authenticate(&self, timeout: Duration) -> Result<Token> {
81        match self {
82            Self::UserToMachine(flow) => flow.login(timeout).await,
83            Self::UserToMachineCli(flow) => flow.token(false).await,
84            Self::MachineToMachine(flow) => flow.token().await,
85        }
86    }
87    async fn login(&self, timeout: Duration) -> Result<Token> {
88        match self {
89            Self::UserToMachine(flow) => flow.login(timeout).await,
90            Self::UserToMachineCli(flow) => flow.native.login(timeout).await,
91            Self::MachineToMachine(flow) => flow.token().await,
92        }
93    }
94    async fn refresh(&self, token: &Token) -> Result<Token> {
95        match self {
96            Self::UserToMachine(flow) => flow.refresh(token).await,
97            Self::UserToMachineCli(flow) => flow.token(true).await,
98            Self::MachineToMachine(flow) => flow.token().await,
99        }
100    }
101    fn can_authenticate_silently(&self) -> bool {
102        matches!(self, Self::UserToMachineCli(_) | Self::MachineToMachine(_))
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    #[test]
111    fn parses_databricks_cli_token_json() {
112        let token: Token = serde_json::from_slice(
113            br#"{"access_token":"access","token_type":"Bearer","refresh_token":"refresh","expiry":"2026-09-05T13:00:00Z","scopes":["all-apis","offline_access"]}"#,
114        )
115        .unwrap();
116        assert_eq!(token.access_token, "access");
117        assert_eq!(token.refresh_token.as_deref(), Some("refresh"));
118        assert_eq!(token.scopes, ["all-apis", "offline_access"]);
119    }
120}
dbx_tools_databricks_auth/client.rs

dbx_tools_databricks_auth/
client.rs
