Bradley Axen
71 lines
2.1 KiB
Rust
71 lines
2.1 KiB
Rust
use anyhow::Result;
|
|
use async_trait::async_trait;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::sync::{Arc, Mutex};
|
|
|
|
use super::api_client::AuthProvider;
|
|
use super::oauth;
|
|
|
|
const DEFAULT_CLIENT_ID: &str = "databricks-cli";
|
|
const DEFAULT_REDIRECT_URL: &str = "http://localhost";
|
|
const DEFAULT_SCOPES: &[&str] = &["all-apis", "offline_access"];
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub enum DatabricksAuth {
|
|
Token(String),
|
|
OAuth {
|
|
host: String,
|
|
client_id: String,
|
|
redirect_url: String,
|
|
scopes: Vec<String>,
|
|
},
|
|
}
|
|
|
|
impl DatabricksAuth {
|
|
pub fn oauth(host: String) -> Self {
|
|
Self::OAuth {
|
|
host,
|
|
client_id: DEFAULT_CLIENT_ID.to_string(),
|
|
redirect_url: DEFAULT_REDIRECT_URL.to_string(),
|
|
scopes: DEFAULT_SCOPES.iter().map(|s| s.to_string()).collect(),
|
|
}
|
|
}
|
|
|
|
pub fn token(token: String) -> Self {
|
|
Self::Token(token)
|
|
}
|
|
}
|
|
|
|
pub(crate) struct DatabricksAuthProvider {
|
|
pub auth: DatabricksAuth,
|
|
pub token_cache: Arc<Mutex<Option<String>>>,
|
|
}
|
|
|
|
#[async_trait]
|
|
impl AuthProvider for DatabricksAuthProvider {
|
|
async fn get_auth_header(&self) -> Result<(String, String)> {
|
|
let token = match &self.auth {
|
|
DatabricksAuth::Token(original) => {
|
|
let cached = self.token_cache.lock().unwrap().clone();
|
|
match cached {
|
|
Some(t) => t,
|
|
None => {
|
|
let fresh = crate::config::Config::global()
|
|
.get_secret::<String>("DATABRICKS_TOKEN")
|
|
.unwrap_or_else(|_| original.clone());
|
|
*self.token_cache.lock().unwrap() = Some(fresh.clone());
|
|
fresh
|
|
}
|
|
}
|
|
}
|
|
DatabricksAuth::OAuth {
|
|
host,
|
|
client_id,
|
|
redirect_url,
|
|
scopes,
|
|
} => oauth::get_oauth_token_async(host, client_id, redirect_url, scopes).await?,
|
|
};
|
|
Ok(("Authorization".to_string(), format!("Bearer {}", token)))
|
|
}
|
|
}
|