dbx_tools_databricks_auth/
lib.rs

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

dbx_tools_databricks_auth/
lib.rs
