dbx_tools_databricks_auth/
client.rs

1use crate::{
2    databricks_cli::DatabricksCliFlow, AuthKind, AuthSession, CredentialStore,
3    MachineToMachineFlow, OAuthFlow, OAuthTemplate, Profile, Result, Token,
4};
5pub use dbx_tools_auth::AuthOptions;
6use std::{sync::Arc, time::Duration};
7
8/// Databricks profile and acquisition policy over the shared `AuthSession` lifecycle.
9pub struct AuthClient {
10    profile: Profile,
11    inner: dbx_tools_auth::AuthClient,
12}
13
14enum AuthFlow {
15    UserToMachine(OAuthFlow),
16    UserToMachineCli(DatabricksCliFlow),
17    MachineToMachine(MachineToMachineFlow),
18}
19
20impl AuthClient {
21    pub fn new(
22        profile: Profile,
23        store: Arc<dyn CredentialStore>,
24        options: AuthOptions,
25        use_databricks_cli: bool,
26    ) -> Result<Self> {
27        let flow = match profile.auth_kind {
28            AuthKind::UserToMachine => {
29                let native = OAuthFlow::new(profile.clone())?
30                    .with_template(OAuthTemplate::new(options.callback_image_src.clone()));
31                if use_databricks_cli {
32                    AuthFlow::UserToMachineCli(DatabricksCliFlow::new(native, profile.name.clone()))
33                } else {
34                    AuthFlow::UserToMachine(native)
35                }
36            }
37            AuthKind::MachineToMachine => {
38                AuthFlow::MachineToMachine(MachineToMachineFlow::new(profile.clone())?)
39            }
40        };
41        let inner =
42            dbx_tools_auth::AuthClient::new(profile.cache_key(), Arc::new(flow), store, options);
43        Ok(Self { profile, inner })
44    }
45    pub fn profile(&self) -> &Profile {
46        &self.profile
47    }
48}
49
50impl AuthSession for AuthClient {
51    fn auth_client(&self) -> &dbx_tools_auth::AuthClient {
52        &self.inner
53    }
54}
55
56#[async_trait::async_trait]
57impl dbx_tools_auth::TokenProvider for AuthFlow {
58    async fn authenticate(&self, timeout: Duration) -> Result<Token> {
59        match self {
60            Self::UserToMachine(flow) => flow.login(timeout).await,
61            Self::UserToMachineCli(flow) => flow.authenticate(timeout).await,
62            Self::MachineToMachine(flow) => flow.token().await,
63        }
64    }
65    async fn login(&self, timeout: Duration) -> Result<Token> {
66        match self {
67            Self::UserToMachine(flow) => flow.login(timeout).await,
68            Self::UserToMachineCli(flow) => flow.login(timeout).await,
69            Self::MachineToMachine(flow) => flow.token().await,
70        }
71    }
72    async fn refresh(&self, token: &Token) -> Result<Token> {
73        match self {
74            Self::UserToMachine(flow) => flow.refresh(token).await,
75            Self::UserToMachineCli(flow) => flow.refresh(token).await,
76            Self::MachineToMachine(flow) => flow.token().await,
77        }
78    }
79    fn can_authenticate_silently(&self) -> bool {
80        matches!(self, Self::UserToMachineCli(_) | Self::MachineToMachine(_))
81    }
82}
dbx_tools_databricks_auth/client.rs

dbx_tools_databricks_auth/
client.rs
