Skip to main content

dbx_tools_databricks_auth/storage/
mod.rs

1mod file;
2mod memory;
3
4#[cfg(feature = "keyring")]
5mod keyring;
6
7use std::{path::PathBuf, sync::Arc, time::Duration};
8
9use async_trait::async_trait;
10use directories::UserDirs;
11
12use crate::{profile::configured_auth_storage, resolve_config_file, Error, Result, Token};
13
14pub use file::FileStore;
15#[cfg(feature = "keyring")]
16pub use keyring::KeyringStore;
17pub use memory::MemoryStore;
18
19#[async_trait]
20pub trait StorageLock: Send + Sync {
21    async fn release(self: Box<Self>) -> Result<()> {
22        Ok(())
23    }
24}
25
26#[async_trait]
27pub trait CredentialStore: Send + Sync {
28    async fn load(&self, profile: &str) -> Result<Option<Token>>;
29    async fn prepare_write(&self) -> Result<()> {
30        Ok(())
31    }
32    async fn save(&self, profile: &str, token: &Token) -> Result<()>;
33    async fn delete(&self, profile: &str) -> Result<()>;
34    async fn lock(&self, profile: &str, timeout: Duration) -> Result<Box<dyn StorageLock>>;
35    fn name(&self) -> &'static str;
36}
37
38#[derive(Clone, Debug, Default)]
39pub struct StoreOptions {
40    pub backend: Option<StoreBackend>,
41    pub cache_dir: Option<PathBuf>,
42    pub config_file: Option<PathBuf>,
43}
44
45#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
46pub enum StoreBackend {
47    #[default]
48    Auto,
49    Memory,
50    File,
51    Keyring,
52}
53
54pub async fn open_store(options: StoreOptions) -> Result<Arc<dyn CredentialStore>> {
55    let cache_dir = options.cache_dir.unwrap_or_else(default_cache_dir);
56    let backend = resolve_backend(options.backend, options.config_file.as_deref())?;
57    match backend {
58        StoreBackend::Memory => Ok(Arc::new(MemoryStore::new(cache_dir.join("memory-locks"))?)),
59        StoreBackend::File => Ok(Arc::new(FileStore::new(cache_dir)?)),
60        StoreBackend::Keyring => open_keyring_for_read(cache_dir).await,
61        StoreBackend::Auto => {
62            #[cfg(feature = "keyring")]
63            {
64                if let Ok(store) = KeyringStore::open_for_read(cache_dir.clone()).await {
65                    return Ok(Arc::new(store));
66                }
67            }
68            Ok(Arc::new(FileStore::new(cache_dir)?))
69        }
70    }
71}
72
73fn default_cache_dir() -> PathBuf {
74    UserDirs::new()
75        .map(|dirs| dirs.home_dir().join(".databricks"))
76        .unwrap_or_else(|| PathBuf::from(".dbx-tools-auth-u2m"))
77}
78
79fn resolve_backend(
80    explicit: Option<StoreBackend>,
81    config_file: Option<&std::path::Path>,
82) -> Result<StoreBackend> {
83    if let Some(backend) = explicit.filter(|backend| *backend != StoreBackend::Auto) {
84        return Ok(backend);
85    }
86    if let Ok(value) = std::env::var("DATABRICKS_AUTH_STORAGE") {
87        if !value.trim().is_empty() {
88            return parse_databricks_storage(&value, "DATABRICKS_AUTH_STORAGE");
89        }
90    }
91    let path = resolve_config_file(config_file)?;
92    if let Some(value) = configured_auth_storage(&path)? {
93        return parse_databricks_storage(&value, "auth_storage");
94    }
95    Ok(StoreBackend::Auto)
96}
97
98fn parse_databricks_storage(value: &str, source: &str) -> Result<StoreBackend> {
99    match value.trim().to_ascii_lowercase().as_str() {
100        "secure" => Ok(StoreBackend::Keyring),
101        "plaintext" => Ok(StoreBackend::File),
102        value => Err(Error::Config(format!(
103            "{source}: unknown storage mode {value:?} (want plaintext or secure)"
104        ))),
105    }
106}
107
108async fn open_keyring_for_read(cache_dir: PathBuf) -> Result<Arc<dyn CredentialStore>> {
109    #[cfg(feature = "keyring")]
110    {
111        Ok(Arc::new(KeyringStore::open_for_read(cache_dir).await?))
112    }
113    #[cfg(not(feature = "keyring"))]
114    {
115        let _ = cache_dir;
116        Err(crate::Error::Config(
117            "keyring support was not compiled in".into(),
118        ))
119    }
120}