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 in_app = is_databricks_app();
113    let profile = resolve_profile(&options, in_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, is_databricks_app())?;
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, in_app: bool) -> 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        access_token: None,
181        group_id: options.group_id.clone(),
182        auth_type: options.auth_type.clone(),
183        scopes: options.scopes.clone(),
184        target: options.target.as_deref().map(parse_target).transpose()?,
185        config_file: options.config_file.as_deref().map(PathBuf::from),
186        prefer_user_to_machine: options.prefer_user_to_machine,
187        skip_implicit_pat: in_app,
188    })
189    .map_err(binding_error)
190}
191
192#[uniffi::export(async_runtime = "tokio")]
193impl PersistentAuth {
194    /// Start an explicit login and persist the resulting credential.
195    pub async fn challenge(&self) -> BindingResult<()> {
196        self.inner.login().await.map(|_| ()).map_err(binding_error)
197    }
198
199    /// True forces login, false forbids interactive login, and omission permits missing-token login.
200    #[uniffi::method(default(login = None))]
201    pub async fn token(&self, login: Option<bool>) -> BindingResult<AccessToken> {
202        self.inner
203            .token_with_login(login)
204            .await
205            .map(Into::into)
206            .map_err(binding_error)
207    }
208
209    /// Renew the stored credential even before its refresh window.
210    pub async fn force_refresh_token(&self) -> BindingResult<AccessToken> {
211        self.inner
212            .force_refresh()
213            .await
214            .map(Into::into)
215            .map_err(binding_error)
216    }
217
218    /// Reuse another caller's replacement or renew the rejected token.
219    pub async fn refresh_rejected_token(
220        &self,
221        stale_access_token: String,
222    ) -> BindingResult<AccessToken> {
223        self.inner
224            .refresh_rejected_token(&stale_access_token)
225            .await
226            .map(Into::into)
227            .map_err(binding_error)
228    }
229
230    /// Delete the credential while holding the store's refresh lock.
231    pub async fn logout(&self) -> BindingResult<()> {
232        self.inner.logout().await.map_err(binding_error)
233    }
234
235    /// Return the resolved identity and active built-in storage backend.
236    pub fn status(&self) -> DatabricksAuthStatus {
237        DatabricksAuthStatus {
238            profile: self.inner.profile().name.clone(),
239            host: self.inner.profile().host.to_string(),
240            storage: storage_from_name(self.inner.store_name()),
241        }
242    }
243}
244
245async fn open_binding_store(
246    options: &DatabricksAuthOptions,
247    storage: Storage,
248) -> BindingResult<Arc<dyn CredentialStore>> {
249    open_store(StoreOptions {
250        backend: Some(storage),
251        cache_dir: options.cache_dir.as_deref().map(PathBuf::from),
252    })
253    .await
254    .map_err(binding_error)
255}
256
257fn parse_target(value: &str) -> BindingResult<TargetKind> {
258    match value.trim().to_ascii_lowercase().as_str() {
259        "workspace" => Ok(TargetKind::Workspace),
260        "account" => Ok(TargetKind::Account),
261        "unified" => Ok(TargetKind::Unified),
262        _ => Err(DatabricksAuthError::Failure {
263            message: "target must be workspace, account, or unified".into(),
264        }),
265    }
266}
267
268fn storage_from_name(name: &str) -> Storage {
269    match name {
270        "memory" => Storage::Memory,
271        _ => Storage::File,
272    }
273}
274
275fn binding_error(error: impl std::fmt::Display) -> DatabricksAuthError {
276    DatabricksAuthError::Failure {
277        message: error.to_string(),
278    }
279}
280
281uniffi::setup_scaffolding!();
282
283#[cfg(test)]
284mod tests {
285    use super::*;
286
287    #[test]
288    fn cli_refresh_requires_available_automatic_u2m() {
289        assert!(should_use_databricks_cli(
290            AuthKind::UserToMachine,
291            None,
292            false,
293            true
294        ));
295        assert!(should_use_databricks_cli(
296            AuthKind::UserToMachine,
297            Some(Storage::Auto),
298            false,
299            true
300        ));
301        assert!(!should_use_databricks_cli(
302            AuthKind::UserToMachine,
303            Some(Storage::File),
304            false,
305            true
306        ));
307        assert!(!should_use_databricks_cli(
308            AuthKind::UserToMachine,
309            Some(Storage::Memory),
310            false,
311            true
312        ));
313        assert!(!should_use_databricks_cli(
314            AuthKind::MachineToMachine,
315            None,
316            false,
317            true
318        ));
319        assert!(!should_use_databricks_cli(
320            AuthKind::PersonalAccessToken,
321            None,
322            false,
323            true
324        ));
325        assert!(!should_use_databricks_cli(
326            AuthKind::UserToMachine,
327            None,
328            false,
329            false
330        ));
331        assert!(!should_use_databricks_cli(
332            AuthKind::UserToMachine,
333            None,
334            true,
335            true
336        ));
337    }
338
339    #[test]
340    fn automatic_storage_tracks_the_runtime_and_explicit_storage_is_preserved() {
341        assert_eq!(storage_backend(None, false), Storage::File);
342        assert_eq!(storage_backend(Some(Storage::Auto), false), Storage::File);
343        assert_eq!(storage_backend(None, true), Storage::Memory);
344        assert_eq!(storage_backend(Some(Storage::Auto), true), Storage::Memory);
345        assert_eq!(storage_backend(Some(Storage::File), true), Storage::File);
346        assert_eq!(
347            storage_backend(Some(Storage::Memory), true),
348            Storage::Memory
349        );
350    }
351}
dbx_tools_databricks_auth/lib.rs

dbx_tools_databricks_auth/
lib.rs
