dbx_tools_databricks_auth/
lib.rs

1mod client;
2mod m2m;
3mod oauth;
4mod oauth_endpoints;
5mod profile;
6mod storage;
7
8pub use client::{AuthClient, AuthOptions};
9pub use dbx_tools_auth::AuthSession;
10pub use dbx_tools_auth::{default_callback_image_src, OAuthTemplate, OAuthTemplateContext};
11pub use dbx_tools_auth::{
12    AccessToken, AuthError as DatabricksAuthError, Storage, StorageAdapter, Token,
13};
14use dbx_tools_auth::{BindingResult, StorageHandle};
15pub use dbx_tools_auth::{Error, Result};
16use dbx_tools_databricks::{databricks_cli_available, is_databricks_app};
17pub use m2m::MachineToMachineFlow;
18pub use oauth::OAuthFlow;
19pub use profile::{
20    resolve_config_file, AuthKind, Profile, ProfileOptions, TargetKind, DEFAULT_ACCOUNTS_HOST,
21    DEFAULT_CLIENT_ID, DEFAULT_CONFIG_FILE,
22};
23use std::{path::PathBuf, sync::Arc};
24pub use storage::{
25    open_store, CredentialStore, FileStore, MemoryStore, StorageLock, StoreBackend, StoreOptions,
26};
27
28/// Configuration shared by the generated Node and Python auth bindings.
29#[derive(Clone, uniffi::Record)]
30pub struct DatabricksAuthOptions {
31    /// Explicit profile name; explicit choices are never remapped to another profile.
32    #[uniffi(default = None)]
33    pub profile: Option<String>,
34    /// Override the workspace or account host.
35    #[uniffi(default = None)]
36    pub host: Option<String>,
37    /// Account identifier for account-scoped authentication.
38    #[uniffi(default = None)]
39    pub account_id: Option<String>,
40    /// Workspace identifier for unified authentication.
41    #[uniffi(default = None)]
42    pub workspace_id: Option<String>,
43    /// Override the Databricks CLI configuration file.
44    #[uniffi(default = None)]
45    pub config_file: Option<String>,
46    /// Override the OAuth application identifier.
47    #[uniffi(default = None)]
48    pub client_id: Option<String>,
49    /// Optional group role requested by M2M token generation.
50    #[uniffi(default = None)]
51    pub group_id: Option<String>,
52    /// Explicit Databricks authentication strategy.
53    #[uniffi(default = None)]
54    pub auth_type: Option<String>,
55    /// Override profile scopes; omission preserves profile/default scope resolution.
56    #[uniffi(default = None)]
57    pub scopes: Option<Vec<String>>,
58    /// Target kind: workspace, account, or unified.
59    #[uniffi(default = None)]
60    pub target: Option<String>,
61    /// Override the directory containing the shared CLI token cache.
62    #[uniffi(default = None)]
63    pub cache_dir: Option<String>,
64    /// Shared lifecycle configuration; omission uses `AuthOptions::default()`.
65    #[uniffi(default = None)]
66    pub auth: Option<AuthOptions>,
67    /// Whether implicit M2M defaults should select one matching U2M profile.
68    #[uniffi(default = true)]
69    pub prefer_user_to_machine: bool,
70}
71
72impl Default for DatabricksAuthOptions {
73    fn default() -> Self {
74        Self {
75            profile: None,
76            host: None,
77            account_id: None,
78            workspace_id: None,
79            config_file: None,
80            client_id: None,
81            group_id: None,
82            auth_type: None,
83            scopes: None,
84            target: None,
85            cache_dir: None,
86            auth: None,
87            prefer_user_to_machine: true,
88        }
89    }
90}
91
92#[derive(Clone, uniffi::Record)]
93/// Resolved Databricks identity and active storage backend.
94pub struct DatabricksAuthStatus {
95    pub profile: String,
96    pub host: String,
97    pub storage: Storage,
98}
99
100#[derive(uniffi::Object)]
101/// Databricks binding facade over the shared persistent authentication lifecycle.
102pub struct PersistentAuth {
103    inner: AuthClient,
104}
105
106#[uniffi::export(async_runtime = "tokio", default(storage = None))]
107/// Resolve a Databricks profile and open built-in credential storage.
108pub async fn create_persistent_auth(
109    options: DatabricksAuthOptions,
110    storage: Option<Storage>,
111) -> BindingResult<Arc<PersistentAuth>> {
112    let profile = resolve_profile(&options)?;
113    let in_app = is_databricks_app();
114    let use_databricks_cli = should_use_databricks_cli(
115        profile.auth_kind,
116        storage,
117        in_app,
118        databricks_cli_available(),
119    );
120    let backend = storage_backend(storage, in_app);
121    let store = open_binding_store(&options, backend).await?;
122    create_persistent_auth_with_store(options, profile, store, use_databricks_cli).await
123}
124
125#[uniffi::export(async_runtime = "tokio")]
126/// Resolve a Databricks profile using a shared owning-library storage handle.
127pub async fn create_persistent_auth_with_storage(
128    options: DatabricksAuthOptions,
129    storage: Arc<StorageHandle>,
130) -> BindingResult<Arc<PersistentAuth>> {
131    let profile = resolve_profile(&options)?;
132    create_persistent_auth_with_store(options, profile, storage.store.clone(), false).await
133}
134
135async fn create_persistent_auth_with_store(
136    options: DatabricksAuthOptions,
137    profile: Profile,
138    store: Arc<dyn CredentialStore>,
139    use_databricks_cli: bool,
140) -> BindingResult<Arc<PersistentAuth>> {
141    let inner = AuthClient::new(
142        profile,
143        store,
144        options.auth.unwrap_or_default(),
145        use_databricks_cli,
146    )
147    .map_err(binding_error)?;
148    Ok(Arc::new(PersistentAuth { inner }))
149}
150
151fn should_use_databricks_cli(
152    auth_kind: AuthKind,
153    storage: Option<Storage>,
154    in_app: bool,
155    available: bool,
156) -> bool {
157    auth_kind == AuthKind::UserToMachine
158        && !in_app
159        && available
160        && storage.is_none_or(|storage| storage == Storage::Auto)
161}
162
163fn storage_backend(storage: Option<Storage>, in_app: bool) -> Storage {
164    match storage {
165        Some(Storage::Memory) => Storage::Memory,
166        Some(Storage::File) => Storage::File,
167        Some(Storage::Auto) | None if in_app => Storage::Memory,
168        Some(Storage::Auto) | None => Storage::File,
169    }
170}
171
172fn resolve_profile(options: &DatabricksAuthOptions) -> BindingResult<Profile> {
173    Profile::from_sources(ProfileOptions {
174        profile: options.profile.clone(),
175        host: options.host.clone(),
176        account_id: options.account_id.clone(),
177        workspace_id: options.workspace_id.clone(),
178        client_id: options.client_id.clone(),
179        client_secret: None,
180        group_id: options.group_id.clone(),
181        auth_type: options.auth_type.clone(),
182        scopes: options.scopes.clone(),
183        target: options.target.as_deref().map(parse_target).transpose()?,
184        config_file: options.config_file.as_deref().map(PathBuf::from),
185        prefer_user_to_machine: options.prefer_user_to_machine,
186    })
187    .map_err(binding_error)
188}
189
190#[uniffi::export(async_runtime = "tokio")]
191impl PersistentAuth {
192    /// Start an explicit login and persist the resulting credential.
193    pub async fn challenge(&self) -> BindingResult<()> {
194        self.inner.login().await.map(|_| ()).map_err(binding_error)
195    }
196
197    /// True forces login, false forbids interactive login, and omission permits missing-token login.
198    #[uniffi::method(default(login = None))]
199    pub async fn token(&self, login: Option<bool>) -> BindingResult<AccessToken> {
200        self.inner
201            .token_with_login(login)
202            .await
203            .map(Into::into)
204            .map_err(binding_error)
205    }
206
207    /// Renew the stored credential even before its refresh window.
208    pub async fn force_refresh_token(&self) -> BindingResult<AccessToken> {
209        self.inner
210            .force_refresh()
211            .await
212            .map(Into::into)
213            .map_err(binding_error)
214    }
215
216    /// Reuse another caller's replacement or renew the rejected token.
217    pub async fn refresh_rejected_token(
218        &self,
219        stale_access_token: String,
220    ) -> BindingResult<AccessToken> {
221        self.inner
222            .refresh_rejected_token(&stale_access_token)
223            .await
224            .map(Into::into)
225            .map_err(binding_error)
226    }
227
228    /// Delete the credential while holding the store's refresh lock.
229    pub async fn logout(&self) -> BindingResult<()> {
230        self.inner.logout().await.map_err(binding_error)
231    }
232
233    /// Return the resolved identity and active built-in storage backend.
234    pub fn status(&self) -> DatabricksAuthStatus {
235        DatabricksAuthStatus {
236            profile: self.inner.profile().name.clone(),
237            host: self.inner.profile().host.to_string(),
238            storage: storage_from_name(self.inner.store_name()),
239        }
240    }
241}
242
243async fn open_binding_store(
244    options: &DatabricksAuthOptions,
245    storage: Storage,
246) -> BindingResult<Arc<dyn CredentialStore>> {
247    open_store(StoreOptions {
248        backend: Some(storage),
249        cache_dir: options.cache_dir.as_deref().map(PathBuf::from),
250    })
251    .await
252    .map_err(binding_error)
253}
254
255fn parse_target(value: &str) -> BindingResult<TargetKind> {
256    match value.trim().to_ascii_lowercase().as_str() {
257        "workspace" => Ok(TargetKind::Workspace),
258        "account" => Ok(TargetKind::Account),
259        "unified" => Ok(TargetKind::Unified),
260        _ => Err(DatabricksAuthError::Failure {
261            message: "target must be workspace, account, or unified".into(),
262        }),
263    }
264}
265
266fn storage_from_name(name: &str) -> Storage {
267    match name {
268        "memory" => Storage::Memory,
269        _ => Storage::File,
270    }
271}
272
273fn binding_error(error: impl std::fmt::Display) -> DatabricksAuthError {
274    DatabricksAuthError::Failure {
275        message: error.to_string(),
276    }
277}
278
279uniffi::setup_scaffolding!();
280
281#[cfg(test)]
282mod tests {
283    use super::*;
284
285    #[test]
286    fn cli_refresh_requires_available_automatic_u2m() {
287        assert!(should_use_databricks_cli(
288            AuthKind::UserToMachine,
289            None,
290            false,
291            true
292        ));
293        assert!(should_use_databricks_cli(
294            AuthKind::UserToMachine,
295            Some(Storage::Auto),
296            false,
297            true
298        ));
299        assert!(!should_use_databricks_cli(
300            AuthKind::UserToMachine,
301            Some(Storage::File),
302            false,
303            true
304        ));
305        assert!(!should_use_databricks_cli(
306            AuthKind::UserToMachine,
307            Some(Storage::Memory),
308            false,
309            true
310        ));
311        assert!(!should_use_databricks_cli(
312            AuthKind::MachineToMachine,
313            None,
314            false,
315            true
316        ));
317        assert!(!should_use_databricks_cli(
318            AuthKind::UserToMachine,
319            None,
320            false,
321            false
322        ));
323        assert!(!should_use_databricks_cli(
324            AuthKind::UserToMachine,
325            None,
326            true,
327            true
328        ));
329    }
330
331    #[test]
332    fn automatic_storage_tracks_the_runtime_and_explicit_storage_is_preserved() {
333        assert_eq!(storage_backend(None, false), Storage::File);
334        assert_eq!(storage_backend(Some(Storage::Auto), false), Storage::File);
335        assert_eq!(storage_backend(None, true), Storage::Memory);
336        assert_eq!(storage_backend(Some(Storage::Auto), true), Storage::Memory);
337        assert_eq!(storage_backend(Some(Storage::File), true), Storage::File);
338        assert_eq!(
339            storage_backend(Some(Storage::Memory), true),
340            Storage::Memory
341        );
342    }
343}
dbx_tools_databricks_auth/lib.rs

dbx_tools_databricks_auth/
lib.rs
