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    PersonalAccessToken(Token),
19}
20
21struct DatabricksCliFlow {
22    native: OAuthFlow,
23    profile: String,
24}
25
26impl DatabricksCliFlow {
27    fn new(native: OAuthFlow, profile: String) -> Self {
28        Self { native, profile }
29    }
30
31    async fn token(&self, force_refresh: bool) -> Result<Token> {
32        let profile = self.profile.clone();
33        tokio::task::spawn_blocking(move || {
34            let output = dbx_tools_databricks::databricks_cli_token(&profile, force_refresh)
35                .map_err(|error| Error::OAuth(error.to_string()))?;
36            serde_json::from_slice(&output).map_err(Into::into)
37        })
38        .await
39        .map_err(|error| Error::OAuth(format!("databricks auth token task failed: {error}")))?
40    }
41}
42
43impl AuthClient {
44    pub fn new(
45        profile: Profile,
46        store: Arc<dyn CredentialStore>,
47        options: AuthOptions,
48        use_databricks_cli: bool,
49    ) -> Result<Self> {
50        let flow = match profile.auth_kind {
51            AuthKind::UserToMachine => {
52                let native = OAuthFlow::new(profile.clone())?
53                    .with_template(OAuthTemplate::new(options.callback_image_src.clone()));
54                if use_databricks_cli {
55                    AuthFlow::UserToMachineCli(DatabricksCliFlow::new(native, profile.name.clone()))
56                } else {
57                    AuthFlow::UserToMachine(native)
58                }
59            }
60            AuthKind::MachineToMachine => {
61                AuthFlow::MachineToMachine(MachineToMachineFlow::new(profile.clone())?)
62            }
63            AuthKind::PersonalAccessToken => AuthFlow::PersonalAccessToken(Token {
64                access_token: profile
65                    .access_token()
66                    .ok_or_else(|| Error::Config("pat requires token".into()))?
67                    .to_owned(),
68                token_type: "Bearer".into(),
69                refresh_token: None,
70                expires_at: None,
71                scopes: Vec::new(),
72            }),
73        };
74        let inner =
75            dbx_tools_auth::AuthClient::new(profile.cache_key(), Arc::new(flow), store, options);
76        Ok(Self { profile, inner })
77    }
78    pub fn profile(&self) -> &Profile {
79        &self.profile
80    }
81}
82
83impl AuthSession for AuthClient {
84    fn auth_client(&self) -> &dbx_tools_auth::AuthClient {
85        &self.inner
86    }
87}
88
89#[async_trait::async_trait]
90impl dbx_tools_auth::TokenProvider for AuthFlow {
91    async fn authenticate(&self, timeout: Duration) -> Result<Token> {
92        match self {
93            Self::UserToMachine(flow) => flow.login(timeout).await,
94            Self::UserToMachineCli(flow) => flow.token(false).await,
95            Self::MachineToMachine(flow) => flow.token().await,
96            Self::PersonalAccessToken(token) => Ok(token.clone()),
97        }
98    }
99    async fn login(&self, timeout: Duration) -> Result<Token> {
100        match self {
101            Self::UserToMachine(flow) => flow.login(timeout).await,
102            Self::UserToMachineCli(flow) => flow.native.login(timeout).await,
103            Self::MachineToMachine(flow) => flow.token().await,
104            Self::PersonalAccessToken(token) => Ok(token.clone()),
105        }
106    }
107    async fn refresh(&self, token: &Token) -> Result<Token> {
108        match self {
109            Self::UserToMachine(flow) => flow.refresh(token).await,
110            Self::UserToMachineCli(flow) => flow.token(true).await,
111            Self::MachineToMachine(flow) => flow.token().await,
112            Self::PersonalAccessToken(token) => Ok(token.clone()),
113        }
114    }
115    fn can_authenticate_silently(&self) -> bool {
116        matches!(
117            self,
118            Self::UserToMachineCli(_) | Self::MachineToMachine(_) | Self::PersonalAccessToken(_)
119        )
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126    use crate::{MemoryStore, TargetKind};
127    use url::Url;
128
129    #[test]
130    fn parses_databricks_cli_token_json() {
131        let token: Token = serde_json::from_slice(
132            br#"{"access_token":"access","token_type":"Bearer","refresh_token":"refresh","expiry":"2026-09-05T13:00:00Z","scopes":["all-apis","offline_access"]}"#,
133        )
134        .unwrap();
135        assert_eq!(token.access_token, "access");
136        assert_eq!(token.refresh_token.as_deref(), Some("refresh"));
137        assert_eq!(token.scopes, ["all-apis", "offline_access"]);
138    }
139
140    #[tokio::test]
141    async fn personal_access_token_authenticates_silently() {
142        let profile = Profile {
143            name: "DEFAULT".into(),
144            host: Url::parse("https://workspace.example").unwrap(),
145            account_id: None,
146            workspace_id: None,
147            client_id: String::new(),
148            group_id: None,
149            scopes: Vec::new(),
150            target: TargetKind::Workspace,
151            auth_kind: AuthKind::PersonalAccessToken,
152            client_secret: None,
153            access_token: Some("access".into()),
154        };
155        let client = AuthClient::new(
156            profile,
157            Arc::new(MemoryStore::new()),
158            AuthOptions::default(),
159            false,
160        )
161        .unwrap();
162
163        let token = client.token_with_login(Some(false)).await.unwrap();
164
165        assert_eq!(token.access_token, "access");
166    }
167}
dbx_tools_databricks_auth/client.rs

dbx_tools_databricks_auth/
client.rs
