dbx_tools_databricks_auth/
lib.rs

1mod client;
2mod error;
3mod m2m;
4mod oauth;
5mod oauth_endpoints;
6mod oauth_template;
7mod profile;
8mod storage;
9mod token;
10
11pub use client::{AuthClient, AuthOptions};
12pub use error::{Error, Result};
13pub use m2m::MachineToMachineFlow;
14pub use oauth::OAuthFlow;
15pub use oauth_template::{default_callback_image_src, OAuthTemplate, OAuthTemplateContext};
16pub use profile::{
17    resolve_config_file, AuthKind, Profile, ProfileOptions, TargetKind, DEFAULT_ACCOUNTS_HOST,
18    DEFAULT_CLIENT_ID, DEFAULT_CONFIG_FILE,
19};
20use std::{path::PathBuf, sync::Arc, time::Duration};
21#[cfg(feature = "keyring")]
22pub use storage::KeyringStore;
23pub use storage::{
24    open_store, CredentialStore, FileStore, MemoryStore, StorageLock, StoreBackend, StoreOptions,
25};
26pub use token::Token;
27
28use time::Duration as TimeDuration;
29
30#[uniffi::export(with_foreign)]
31#[async_trait::async_trait]
32pub trait StorageAdapter: Send + Sync {
33    async fn load(&self, profile: String) -> BindingResult<Option<String>>;
34    async fn prepare_write(&self) -> BindingResult<()>;
35    async fn save(&self, profile: String, token: String) -> BindingResult<()>;
36    async fn remove(&self, profile: String) -> BindingResult<()>;
37    async fn acquire_lock(&self, profile: String, timeout_millis: u64) -> BindingResult<String>;
38    async fn release_lock(&self, lease: String) -> BindingResult<()>;
39    fn name(&self) -> String;
40}
41
42/// Configuration shared by the generated Node and Python auth bindings.
43#[derive(Clone, uniffi::Record)]
44pub struct DatabricksAuthOptions {
45    #[uniffi(default = None)]
46    pub profile: Option<String>,
47    #[uniffi(default = None)]
48    pub host: Option<String>,
49    #[uniffi(default = None)]
50    pub account_id: Option<String>,
51    #[uniffi(default = None)]
52    pub workspace_id: Option<String>,
53    #[uniffi(default = None)]
54    pub config_file: Option<String>,
55    #[uniffi(default = None)]
56    pub client_id: Option<String>,
57    /// Optional group role requested by M2M token generation.
58    #[uniffi(default = None)]
59    pub group_id: Option<String>,
60    /// Explicit Databricks authentication strategy.
61    #[uniffi(default = None)]
62    pub auth_type: Option<String>,
63    #[uniffi(default = None)]
64    pub scopes: Option<Vec<String>>,
65    #[uniffi(default = None)]
66    pub target: Option<String>,
67    #[uniffi(default = None)]
68    pub cache_dir: Option<String>,
69    /// Logo URL or data URI displayed by the browser callback page.
70    #[uniffi(default = None)]
71    pub callback_image_src: Option<String>,
72    #[uniffi(default = 30)]
73    pub lock_timeout_seconds: u64,
74    #[uniffi(default = 3600)]
75    pub login_timeout_seconds: u64,
76    #[uniffi(default = 300)]
77    pub refresh_buffer_seconds: i64,
78    /// Whether implicit M2M defaults should select one matching U2M profile.
79    #[uniffi(default = true)]
80    pub prefer_user_to_machine: bool,
81}
82
83impl Default for DatabricksAuthOptions {
84    fn default() -> Self {
85        Self {
86            profile: None,
87            host: None,
88            account_id: None,
89            workspace_id: None,
90            config_file: None,
91            client_id: None,
92            group_id: None,
93            auth_type: None,
94            scopes: None,
95            target: None,
96            cache_dir: None,
97            callback_image_src: None,
98            lock_timeout_seconds: 30,
99            login_timeout_seconds: 3600,
100            refresh_buffer_seconds: 300,
101            prefer_user_to_machine: true,
102        }
103    }
104}
105
106#[derive(Clone, Copy, Debug, PartialEq, Eq, uniffi::Enum)]
107pub enum Storage {
108    Auto,
109    Memory,
110    File,
111    Keyring,
112}
113
114#[derive(Clone, uniffi::Record)]
115pub struct AccessToken {
116    pub access_token: String,
117    pub token_type: String,
118    pub expiry: Option<String>,
119    pub scopes: Vec<String>,
120}
121
122#[derive(Clone, uniffi::Record)]
123pub struct DatabricksAuthStatus {
124    pub profile: String,
125    pub host: String,
126    pub storage: Storage,
127}
128
129#[derive(Debug, thiserror::Error, uniffi::Error)]
130pub enum DatabricksAuthError {
131    #[error("{message}")]
132    Failure { message: String },
133}
134
135impl From<uniffi::UnexpectedUniFFICallbackError> for DatabricksAuthError {
136    fn from(error: uniffi::UnexpectedUniFFICallbackError) -> Self {
137        Self::Failure {
138            message: error.to_string(),
139        }
140    }
141}
142
143type BindingResult<T> = std::result::Result<T, DatabricksAuthError>;
144
145#[derive(uniffi::Object)]
146pub struct PersistentAuth {
147    inner: AuthClient,
148}
149
150#[uniffi::export(async_runtime = "tokio", default(storage = None))]
151pub async fn create_persistent_auth(
152    options: DatabricksAuthOptions,
153    storage: Option<Storage>,
154) -> BindingResult<Arc<PersistentAuth>> {
155    let store = open_binding_store(&options, storage).await?;
156    create_persistent_auth_with_store(options, store).await
157}
158
159#[uniffi::export(async_runtime = "tokio")]
160pub async fn create_persistent_auth_with_storage(
161    options: DatabricksAuthOptions,
162    storage: Arc<dyn StorageAdapter>,
163) -> BindingResult<Arc<PersistentAuth>> {
164    create_persistent_auth_with_store(options, Arc::new(ForeignStore { storage })).await
165}
166
167async fn create_persistent_auth_with_store(
168    options: DatabricksAuthOptions,
169    store: Arc<dyn CredentialStore>,
170) -> BindingResult<Arc<PersistentAuth>> {
171    let profile = Profile::from_sources(ProfileOptions {
172        profile: options.profile.clone(),
173        host: options.host.clone(),
174        account_id: options.account_id.clone(),
175        workspace_id: options.workspace_id.clone(),
176        client_id: options.client_id.clone(),
177        client_secret: None,
178        group_id: options.group_id.clone(),
179        auth_type: options.auth_type.clone(),
180        scopes: options.scopes.clone(),
181        target: options.target.as_deref().map(parse_target).transpose()?,
182        config_file: options.config_file.as_deref().map(PathBuf::from),
183        prefer_user_to_machine: options.prefer_user_to_machine,
184    })
185    .map_err(binding_error)?;
186    let inner = AuthClient::new(
187        profile,
188        store,
189        AuthOptions {
190            refresh_buffer: TimeDuration::seconds(options.refresh_buffer_seconds),
191            lock_timeout: Duration::from_secs(options.lock_timeout_seconds),
192            login_timeout: Duration::from_secs(options.login_timeout_seconds),
193            callback_image_src: options.callback_image_src.clone(),
194        },
195    )
196    .map_err(binding_error)?;
197    Ok(Arc::new(PersistentAuth { inner }))
198}
199
200struct ForeignStore {
201    storage: Arc<dyn StorageAdapter>,
202}
203
204struct ForeignLock {
205    storage: Arc<dyn StorageAdapter>,
206    lease: String,
207}
208
209#[async_trait::async_trait]
210impl StorageLock for ForeignLock {
211    async fn release(self: Box<Self>) -> Result<()> {
212        self.storage
213            .release_lock(self.lease)
214            .await
215            .map_err(|error| Error::Storage(error.to_string()))
216    }
217}
218
219#[async_trait::async_trait]
220impl CredentialStore for ForeignStore {
221    async fn load(&self, profile: &str) -> Result<Option<Token>> {
222        self.storage
223            .load(profile.to_owned())
224            .await
225            .map_err(|error| Error::Storage(error.to_string()))?
226            .map(|token| serde_json::from_str(&token).map_err(Into::into))
227            .transpose()
228    }
229
230    async fn prepare_write(&self) -> Result<()> {
231        self.storage
232            .prepare_write()
233            .await
234            .map_err(|error| Error::Storage(error.to_string()))
235    }
236
237    async fn save(&self, profile: &str, token: &Token) -> Result<()> {
238        self.storage
239            .save(profile.to_owned(), serde_json::to_string(token)?)
240            .await
241            .map_err(|error| Error::Storage(error.to_string()))
242    }
243
244    async fn delete(&self, profile: &str) -> Result<()> {
245        self.storage
246            .remove(profile.to_owned())
247            .await
248            .map_err(|error| Error::Storage(error.to_string()))
249    }
250
251    async fn lock(&self, profile: &str, timeout: Duration) -> Result<Box<dyn StorageLock>> {
252        let timeout_millis = u64::try_from(timeout.as_millis()).unwrap_or(u64::MAX);
253        let lease = self
254            .storage
255            .acquire_lock(profile.to_owned(), timeout_millis)
256            .await
257            .map_err(|error| Error::Storage(error.to_string()))?;
258        Ok(Box::new(ForeignLock {
259            storage: Arc::clone(&self.storage),
260            lease,
261        }))
262    }
263
264    fn name(&self) -> &'static str {
265        "custom"
266    }
267}
268
269#[uniffi::export(async_runtime = "tokio")]
270impl PersistentAuth {
271    pub async fn challenge(&self) -> BindingResult<()> {
272        self.inner.login().await.map(|_| ()).map_err(binding_error)
273    }
274
275    #[uniffi::method(default(login = None))]
276    pub async fn token(&self, login: Option<bool>) -> BindingResult<AccessToken> {
277        let token = match login {
278            Some(true) => self.inner.login().await,
279            None => self.inner.token_or_login().await,
280            Some(false) => self.inner.token().await,
281        };
282        token.map(Into::into).map_err(binding_error)
283    }
284
285    pub async fn force_refresh_token(&self) -> BindingResult<AccessToken> {
286        self.inner
287            .force_refresh()
288            .await
289            .map(Into::into)
290            .map_err(binding_error)
291    }
292
293    pub async fn logout(&self) -> BindingResult<()> {
294        self.inner.logout().await.map_err(binding_error)
295    }
296
297    pub fn status(&self) -> DatabricksAuthStatus {
298        DatabricksAuthStatus {
299            profile: self.inner.profile().name.clone(),
300            host: self.inner.profile().host.to_string(),
301            storage: storage_from_name(self.inner.store_name()),
302        }
303    }
304}
305
306impl From<Token> for AccessToken {
307    fn from(token: Token) -> Self {
308        Self {
309            access_token: token.access_token,
310            token_type: token.token_type,
311            expiry: token.expires_at.map(|value| value.to_string()),
312            scopes: token.scopes,
313        }
314    }
315}
316
317async fn open_binding_store(
318    options: &DatabricksAuthOptions,
319    storage: Option<Storage>,
320) -> BindingResult<Arc<dyn CredentialStore>> {
321    open_store(StoreOptions {
322        backend: storage.map(Into::into),
323        cache_dir: options.cache_dir.as_deref().map(PathBuf::from),
324        config_file: options.config_file.as_deref().map(PathBuf::from),
325    })
326    .await
327    .map_err(binding_error)
328}
329
330fn parse_target(value: &str) -> BindingResult<TargetKind> {
331    match value.trim().to_ascii_lowercase().as_str() {
332        "workspace" => Ok(TargetKind::Workspace),
333        "account" => Ok(TargetKind::Account),
334        "unified" => Ok(TargetKind::Unified),
335        _ => Err(DatabricksAuthError::Failure {
336            message: "target must be workspace, account, or unified".into(),
337        }),
338    }
339}
340
341impl From<Storage> for StoreBackend {
342    fn from(storage: Storage) -> Self {
343        match storage {
344            Storage::Auto => Self::Auto,
345            Storage::Memory => Self::Memory,
346            Storage::File => Self::File,
347            Storage::Keyring => Self::Keyring,
348        }
349    }
350}
351
352fn storage_from_name(name: &str) -> Storage {
353    match name {
354        "memory" => Storage::Memory,
355        "keyring" => Storage::Keyring,
356        _ => Storage::File,
357    }
358}
359
360fn binding_error(error: impl std::fmt::Display) -> DatabricksAuthError {
361    DatabricksAuthError::Failure {
362        message: error.to_string(),
363    }
364}
365
366uniffi::setup_scaffolding!();
dbx_tools_databricks_auth/lib.rs

dbx_tools_databricks_auth/
lib.rs
