dbx_tools_databricks_auth/
client.rs

1use std::{sync::Arc, time::Duration};
2
3use time::{Duration as TimeDuration, OffsetDateTime};
4
5use crate::{
6    AuthKind, CredentialStore, Error, MachineToMachineFlow, OAuthFlow, OAuthTemplate, Profile,
7    Result, Token,
8};
9
10#[derive(Clone, Debug)]
11pub struct AuthOptions {
12    pub refresh_buffer: TimeDuration,
13    pub lock_timeout: Duration,
14    pub login_timeout: Duration,
15    /// Logo URL or data URI displayed by the browser callback page.
16    pub callback_image_src: Option<String>,
17}
18
19impl Default for AuthOptions {
20    fn default() -> Self {
21        Self {
22            refresh_buffer: TimeDuration::minutes(5),
23            lock_timeout: Duration::from_secs(30),
24            login_timeout: Duration::from_secs(3600),
25            callback_image_src: None,
26        }
27    }
28}
29
30pub struct AuthClient {
31    profile: Profile,
32    flow: AuthFlow,
33    store: Arc<dyn CredentialStore>,
34    options: AuthOptions,
35}
36
37enum AuthFlow {
38    UserToMachine(OAuthFlow),
39    MachineToMachine(MachineToMachineFlow),
40}
41
42impl AuthClient {
43    pub fn new(
44        profile: Profile,
45        store: Arc<dyn CredentialStore>,
46        options: AuthOptions,
47    ) -> Result<Self> {
48        let flow = match profile.auth_kind {
49            AuthKind::UserToMachine => AuthFlow::UserToMachine(
50                OAuthFlow::new(profile.clone())?
51                    .with_template(OAuthTemplate::new(options.callback_image_src.clone())),
52            ),
53            AuthKind::MachineToMachine => {
54                AuthFlow::MachineToMachine(MachineToMachineFlow::new(profile.clone())?)
55            }
56        };
57        Ok(Self {
58            profile,
59            flow,
60            store,
61            options,
62        })
63    }
64
65    pub fn profile(&self) -> &Profile {
66        &self.profile
67    }
68
69    pub fn store_name(&self) -> &'static str {
70        self.store.name()
71    }
72
73    pub async fn login(&self) -> Result<Token> {
74        let cache_key = self.profile.cache_key();
75        let lock = self
76            .store
77            .lock(&cache_key, self.options.lock_timeout)
78            .await?;
79        let result = async {
80            let (_, token) = tokio::try_join!(
81                self.store.prepare_write(),
82                self.flow.authenticate(self.options.login_timeout)
83            )?;
84            self.store.save(&cache_key, &token).await?;
85            Ok(public_token(token))
86        }
87        .await;
88        release(lock, result).await
89    }
90
91    pub async fn token(&self) -> Result<Token> {
92        let cache_key = self.profile.cache_key();
93        let now = OffsetDateTime::now_utc();
94        if let Some(token) = self.store.load(&cache_key).await? {
95            if !token.needs_refresh(now, self.options.refresh_buffer) {
96                return Ok(public_token(token));
97            }
98        }
99
100        let lock = self
101            .store
102            .lock(&cache_key, self.options.lock_timeout)
103            .await?;
104        let result = async {
105            let token = self.store.load(&cache_key).await?;
106            let now = OffsetDateTime::now_utc();
107            if let Some(token) = token.as_ref() {
108                if !token.needs_refresh(now, self.options.refresh_buffer) {
109                    return Ok(public_token(token.clone()));
110                }
111            }
112            let loaded = match token {
113                Some(token) => {
114                    self.store.prepare_write().await?;
115                    self.flow.refresh(&token).await?
116                }
117                None if self.flow.is_machine_to_machine() => {
118                    self.store.prepare_write().await?;
119                    self.flow.authenticate(self.options.login_timeout).await?
120                }
121                None => return Err(Error::LoginRequired(self.profile.name.clone())),
122            };
123            self.store.save(&cache_key, &loaded).await?;
124            Ok(public_token(loaded))
125        }
126        .await;
127        release(lock, result).await
128    }
129
130    pub async fn token_or_login(&self) -> Result<Token> {
131        match self.token().await {
132            Err(Error::LoginRequired(_)) => self.login().await,
133            result => result,
134        }
135    }
136
137    pub async fn force_refresh(&self) -> Result<Token> {
138        let cache_key = self.profile.cache_key();
139        let lock = self
140            .store
141            .lock(&cache_key, self.options.lock_timeout)
142            .await?;
143        let result = async {
144            let token = self.store.load(&cache_key).await?;
145            let loaded = match token {
146                Some(token) => {
147                    self.store.prepare_write().await?;
148                    self.flow.refresh(&token).await?
149                }
150                None if self.flow.is_machine_to_machine() => {
151                    self.store.prepare_write().await?;
152                    self.flow.authenticate(self.options.login_timeout).await?
153                }
154                None => return Err(Error::LoginRequired(self.profile.name.clone())),
155            };
156            self.store.save(&cache_key, &loaded).await?;
157            Ok(public_token(loaded))
158        }
159        .await;
160        release(lock, result).await
161    }
162
163    pub async fn logout(&self) -> Result<()> {
164        let cache_key = self.profile.cache_key();
165        let lock = self
166            .store
167            .lock(&cache_key, self.options.lock_timeout)
168            .await?;
169        let result = self.store.delete(&cache_key).await;
170        release(lock, result).await
171    }
172}
173
174impl AuthFlow {
175    async fn authenticate(&self, login_timeout: Duration) -> Result<Token> {
176        match self {
177            Self::UserToMachine(flow) => flow.login(login_timeout).await,
178            Self::MachineToMachine(flow) => flow.token().await,
179        }
180    }
181
182    async fn refresh(&self, token: &Token) -> Result<Token> {
183        match self {
184            Self::UserToMachine(flow) => flow.refresh(token).await,
185            Self::MachineToMachine(flow) => flow.token().await,
186        }
187    }
188
189    fn is_machine_to_machine(&self) -> bool {
190        matches!(self, Self::MachineToMachine(_))
191    }
192}
193
194async fn release<T>(lock: Box<dyn crate::StorageLock>, result: Result<T>) -> Result<T> {
195    let released = lock.release().await;
196    match result {
197        Err(error) => Err(error),
198        Ok(value) => released.map(|()| value),
199    }
200}
201
202fn public_token(mut token: Token) -> Token {
203    token.refresh_token = None;
204    token
205}
dbx_tools_databricks_auth/client.rs

dbx_tools_databricks_auth/
client.rs
