use std::collections::HashMap;
use crate::config::{Config, UserState};
use crate::types::app::{AppGetResponse, AppsListResponse};
use crate::types::job::{
CloudJobCreatedResponse, CreatedJobDict, ExperimentCreateResponse,
ExperimentsPaginatedResponse, JobDict, JobStatusResponse, JobsPaginatedResponse,
PresignedDownloadResponse, PresignedUploadResponse, ResourceDetailedDict,
TokenExchangeResponse, TokenRefreshResponse, UserInfo,
};
use super::http::{HttpClient, HttpResponse};
const CLIENT_TYPE: &str = "biolib-rs";
const CLIENT_VERSION: &str = env!("CARGO_PKG_VERSION");
pub struct ApiClient {
http: HttpClient,
base_url: String,
access_token: Option<String>,
refresh_token: Option<String>,
deploy_key: Option<String>,
is_public_biolib: bool,
cloud_base_url: Option<String>,
cloud_job_storage_base_url: Option<String>,
}
impl ApiClient {
pub fn new(config: &Config) -> crate::Result<Self> {
let mut client = Self {
http: HttpClient::shared(),
base_url: config.base_url.clone(),
access_token: None,
refresh_token: None,
deploy_key: None,
is_public_biolib: config.is_public_biolib,
cloud_base_url: config.cloud_base_url.clone(),
cloud_job_storage_base_url: config.cloud_job_storage_base_url.clone(),
};
client.attempt_sign_in(config)?;
Ok(client)
}
pub fn new_unauthenticated(base_url: &str) -> Self {
Self {
http: HttpClient::shared(),
base_url: base_url.to_string(),
access_token: None,
refresh_token: None,
deploy_key: None,
is_public_biolib: base_url.ends_with("biolib.com"),
cloud_base_url: None,
cloud_job_storage_base_url: None,
}
}
pub fn base_url(&self) -> &str {
&self.base_url
}
pub fn cloud_base_url(&self) -> Option<&str> {
self.cloud_base_url.as_deref()
}
fn get_client_opt_user_info(&self) -> Option<String> {
if self.is_public_biolib {
return None;
}
std::env::var("BIOLIB_OPT_USER")
.ok()
.or_else(|| std::env::var("USER").ok())
}
pub fn is_signed_in(&self) -> bool {
self.refresh_token.is_some() || self.deploy_key.is_some()
}
fn attempt_sign_in(&mut self, config: &Config) -> crate::Result<()> {
if let Some(ref token) = config.api_token {
if token.starts_with("bld_") {
self.deploy_key = Some(token.clone());
} else {
crate::logging::debug("ApiClient: Signing in with BIOLIB_TOKEN...");
if let Err(err) = self.sign_in_with_api_token(token) {
crate::logging::error("Sign in with API token failed");
return Err(err);
}
}
} else {
let user_state = Config::load_user_state();
if let Some(ref refresh_token) = user_state.refresh_token {
crate::logging::debug(
"ApiClient: Signing in with refresh token from user state...",
);
self.refresh_token = Some(refresh_token.clone());
if self.refresh_access_token().is_err() {
crate::logging::error("Sign in with refresh token failed");
self.refresh_token = None;
Config::save_user_state(&UserState::default())?;
}
}
}
Ok(())
}
pub fn sign_in_with_api_token(&mut self, api_token: &str) -> crate::Result<()> {
let url = format!("{}/api/user/api_tokens/exchange/", self.base_url);
let body = serde_json::json!({"token": api_token});
let body_bytes = serde_json::to_vec(&body)?;
let response = self
.http
.request("POST", &url, Some(&body_bytes), None, 5, None)?;
let token_response: TokenExchangeResponse = response.json()?;
self.access_token = Some(token_response.access_token);
self.refresh_token = Some(token_response.refresh_token);
Ok(())
}
pub fn sign_in_with_credentials(&mut self, email: &str, password: &str) -> crate::Result<()> {
let url = format!("{}/api/user/token/", self.base_url);
let body = serde_json::json!({"email": email, "password": password});
let body_bytes = serde_json::to_vec(&body)?;
let response = self
.http
.request("POST", &url, Some(&body_bytes), None, 5, None)?;
let token_response: TokenRefreshResponse = response.json()?;
self.access_token = Some(token_response.access.clone());
let refresh_body = serde_json::json!({"email": email, "password": password});
let refresh_bytes = serde_json::to_vec(&refresh_body)?;
let refresh_response =
self.http
.request("POST", &url, Some(&refresh_bytes), None, 5, None)?;
#[derive(serde::Deserialize)]
struct FullTokenResponse {
access: String,
refresh: String,
}
let full: FullTokenResponse = refresh_response.json()?;
self.access_token = Some(full.access);
self.refresh_token = Some(full.refresh.clone());
Config::save_user_state(&UserState {
refresh_token: Some(full.refresh),
})?;
Ok(())
}
pub fn sign_out(&mut self) -> crate::Result<()> {
self.access_token = None;
self.refresh_token = None;
self.deploy_key = None;
Config::save_user_state(&UserState::default())?;
Ok(())
}
fn refresh_access_token(&mut self) -> crate::Result<()> {
if self.deploy_key.is_some() {
return Ok(());
}
let refresh_token = match &self.refresh_token {
Some(t) => t.clone(),
None => return Ok(()),
};
if let Some(ref access) = self.access_token {
if !is_token_expired(access) {
return Ok(());
}
}
let url = format!("{}/api/user/token/refresh/", self.base_url);
let body = serde_json::json!({"refresh": refresh_token});
let body_bytes = serde_json::to_vec(&body)?;
let response = self
.http
.request("POST", &url, Some(&body_bytes), None, 5, None)?;
let token_response: TokenRefreshResponse = response.json()?;
self.access_token = Some(token_response.access);
Ok(())
}
fn get_auth_headers(&mut self) -> crate::Result<HashMap<String, String>> {
let mut headers = HashMap::new();
headers.insert("client-type".to_string(), CLIENT_TYPE.to_string());
headers.insert("client-version".to_string(), CLIENT_VERSION.to_string());
if let Some(ref key) = self.deploy_key {
headers.insert("authorization".to_string(), format!("Token {key}"));
} else if self.refresh_token.is_some() {
self.refresh_access_token()?;
if let Some(ref token) = self.access_token {
headers.insert("authorization".to_string(), format!("Bearer {token}"));
}
}
Ok(headers)
}
fn api_url(&self, path: &str) -> String {
let trimmed = path.trim_matches('/');
format!("{}/api/{}/", self.base_url, trimmed)
}
pub fn get(
&mut self,
path: &str,
params: Option<&[(&str, &str)]>,
) -> crate::Result<HttpResponse> {
let mut url = self.api_url(path);
if let Some(params) = params {
let query: Vec<String> = params
.iter()
.map(|(k, v)| format!("{}={}", urlencoding(k), urlencoding(v)))
.collect();
if !query.is_empty() {
url = format!("{}?{}", url, query.join("&"));
}
}
let headers = self.get_auth_headers()?;
self.http
.request("GET", &url, None, Some(&headers), 10, None)
}
pub fn post(
&mut self,
path: &str,
data: Option<&[u8]>,
extra_headers: Option<&HashMap<String, String>>,
) -> crate::Result<HttpResponse> {
let url = self.api_url(path);
let mut headers = self.get_auth_headers()?;
if let Some(extra) = extra_headers {
for (k, v) in extra {
headers.insert(k.clone(), v.clone());
}
}
self.http
.request("POST", &url, data, Some(&headers), 5, None)
}
pub fn post_json<T: serde::Serialize>(
&mut self,
path: &str,
data: &T,
) -> crate::Result<HttpResponse> {
let body = serde_json::to_vec(data)?;
self.post(path, Some(&body), None)
}
pub fn patch_json<T: serde::Serialize>(
&mut self,
path: &str,
data: &T,
) -> crate::Result<HttpResponse> {
let url = self.api_url(path);
let headers = self.get_auth_headers()?;
let body = serde_json::to_vec(data)?;
self.http
.request("PATCH", &url, Some(&body), Some(&headers), 5, None)
}
pub fn patch_json_with_headers<T: serde::Serialize>(
&mut self,
path: &str,
data: &T,
extra_headers: &HashMap<String, String>,
) -> crate::Result<HttpResponse> {
let url = self.api_url(path);
let mut headers = self.get_auth_headers()?;
for (k, v) in extra_headers {
headers.insert(k.clone(), v.clone());
}
let body = serde_json::to_vec(data)?;
self.http
.request("PATCH", &url, Some(&body), Some(&headers), 5, None)
}
pub fn delete(&mut self, path: &str) -> crate::Result<HttpResponse> {
let url = self.api_url(path);
let headers = self.get_auth_headers()?;
self.http
.request("DELETE", &url, None, Some(&headers), 0, None)
}
pub fn get_app_by_uri(&self, uri: &str) -> crate::Result<AppGetResponse> {
let url = format!("{}/api/app/?uri={}", self.base_url, urlencoding(uri));
let mut headers = HashMap::new();
headers.insert("client-type".to_string(), CLIENT_TYPE.to_string());
headers.insert("client-version".to_string(), CLIENT_VERSION.to_string());
if let Some(ref key) = self.deploy_key {
headers.insert("authorization".to_string(), format!("Token {key}"));
} else if let Some(ref token) = self.access_token {
headers.insert("authorization".to_string(), format!("Bearer {token}"));
}
let response = self
.http
.request("GET", &url, None, Some(&headers), 10, None)?;
response.json()
}
pub fn search_apps(
&self,
query: Option<&str>,
team: Option<&str>,
count: usize,
) -> crate::Result<Vec<String>> {
let mut url = format!("{}/api/apps/?page_size={}", self.base_url, count);
if let Some(q) = query {
url = format!("{}&search={}", url, urlencoding(q));
}
if let Some(t) = team {
let handle = if t.starts_with('@') {
t.to_string()
} else {
format!("@biolib.com/{t}")
};
url = format!("{}&account_handle={}", url, urlencoding(&handle));
}
let mut headers = HashMap::new();
headers.insert("client-type".to_string(), CLIENT_TYPE.to_string());
headers.insert("client-version".to_string(), CLIENT_VERSION.to_string());
let response = self
.http
.request("GET", &url, None, Some(&headers), 10, None)?;
let apps_response: AppsListResponse = response.json()?;
Ok(apps_response
.results
.into_iter()
.map(|r| r.resource_uri)
.collect())
}
#[allow(clippy::too_many_arguments)]
pub fn create_job(
&mut self,
app_version_id: &str,
app_resource_name_prefix: Option<&str>,
override_command: bool,
machine: Option<&str>,
experiment_uuid: Option<&str>,
timeout: Option<u64>,
notify: bool,
requested_machine_count: Option<u32>,
) -> crate::Result<CreatedJobDict> {
let mut data = serde_json::json!({
"app_version_id": app_version_id,
"client_type": CLIENT_TYPE,
"client_version": CLIENT_VERSION,
"notify": notify,
"client_opt_user_info": self.get_client_opt_user_info(),
});
if let Some(prefix) = app_resource_name_prefix {
data["app_resource_name_prefix"] = serde_json::Value::String(prefix.to_string());
}
if override_command {
data["arguments_override_command"] = serde_json::Value::Bool(true);
}
if let Some(m) = machine {
data["requested_machine"] = serde_json::Value::String(m.to_string());
}
if let Some(exp) = experiment_uuid {
data["experiment_uuid"] = serde_json::Value::String(exp.to_string());
}
if let Some(t) = timeout {
data["requested_timeout_seconds"] = serde_json::Value::Number(t.into());
}
if let Some(count) = requested_machine_count {
data["requested_machine_count"] = serde_json::Value::Number(count.into());
}
let response = self.post_json("/jobs/", &data)?;
response.json()
}
#[allow(clippy::too_many_arguments)]
pub fn create_job_with_data(
&mut self,
app_version_uuid: &str,
app_resource_name_prefix: Option<&str>,
module_input: &[u8],
override_command: bool,
experiment_uuid: Option<&str>,
machine: Option<&str>,
result_name_prefix: Option<&str>,
timeout: Option<u64>,
notify: bool,
requested_machine_count: Option<u32>,
) -> crate::Result<JobDict> {
let url = self.api_url("/jobs/create_job_with_data/");
let mut headers = self.get_auth_headers()?;
headers.insert(
"content-type".to_string(),
"application/octet-stream".to_string(),
);
headers.insert("app-version-uuid".to_string(), app_version_uuid.to_string());
if let Some(prefix) = app_resource_name_prefix {
headers.insert("app-resource-name-prefix".to_string(), prefix.to_string());
}
headers.insert(
"arguments-override-command".to_string(),
override_command.to_string(),
);
headers.insert("client-type".to_string(), CLIENT_TYPE.to_string());
headers.insert("client-version".to_string(), CLIENT_VERSION.to_string());
if let Some(user_info) = self.get_client_opt_user_info() {
headers.insert("client-opt-user-info".to_string(), user_info);
}
if let Some(exp) = experiment_uuid {
headers.insert("experiment-uuid".to_string(), exp.to_string());
}
if let Some(m) = machine {
headers.insert("requested-machine".to_string(), m.to_string());
}
if let Some(prefix) = result_name_prefix {
headers.insert("result-name-prefix".to_string(), prefix.to_string());
}
if let Some(t) = timeout {
headers.insert("requested-timeout-seconds".to_string(), t.to_string());
}
headers.insert(
"notify".to_string(),
if notify { "true" } else { "false" }.to_string(),
);
if let Some(count) = requested_machine_count {
headers.insert("requested-machine-count".to_string(), count.to_string());
}
let response =
self.http
.request("POST", &url, Some(module_input), Some(&headers), 5, None)?;
response.json()
}
pub fn create_cloud_job(
&mut self,
job_id: &str,
result_name_prefix: Option<&str>,
) -> crate::Result<CloudJobCreatedResponse> {
let mut data = serde_json::json!({"job_id": job_id});
if let Some(prefix) = result_name_prefix {
data["result_name_prefix"] = serde_json::Value::String(prefix.to_string());
}
let response = self.post_json("/jobs/cloud/", &data)?;
response.json()
}
pub fn get_job(&mut self, uuid: &str, auth_token: Option<&str>) -> crate::Result<JobDict> {
let url = self.api_url(&format!("/jobs/{uuid}/"));
let mut headers = self.get_auth_headers()?;
if let Some(token) = auth_token {
headers.insert("job-auth-token".to_string(), token.to_string());
}
let response = self
.http
.request("GET", &url, None, Some(&headers), 10, None)?;
response.json()
}
pub fn get_jobs(
&mut self,
page_size: usize,
status: Option<&str>,
page: Option<u64>,
) -> crate::Result<JobsPaginatedResponse> {
let mut url = format!("{}/api/jobs/?page_size={}", self.base_url, page_size);
if let Some(s) = status {
url = format!("{}&state={}", url, s);
}
if let Some(p) = page {
url = format!("{}&page={}", url, p);
}
let headers = self.get_auth_headers()?;
let response = self
.http
.request("GET", &url, None, Some(&headers), 10, None)?;
response.json()
}
pub fn get_job_storage_download_url(
&mut self,
job_uuid: &str,
job_auth_token: &str,
storage_type: &str,
) -> crate::Result<String> {
let url = self.api_url(&format!(
"/jobs/{job_uuid}/storage/{storage_type}/download/"
));
let mut headers = self.get_auth_headers()?;
headers.insert("job-auth-token".to_string(), job_auth_token.to_string());
let response = self
.http
.request("GET", &url, None, Some(&headers), 10, None)?;
let download_response: PresignedDownloadResponse = response.json()?;
Ok(self.rewrite_presigned_url(&download_response.presigned_download_url))
}
fn rewrite_presigned_url(&self, presigned_url: &str) -> String {
if let Some(ref proxy_base) = self.cloud_job_storage_base_url {
if let Some(path_start) = presigned_url.find("://").and_then(|scheme_end| {
presigned_url[scheme_end + 3..]
.find('/')
.map(|p| scheme_end + 3 + p)
}) {
return format!("{}{}", proxy_base, &presigned_url[path_start..]);
}
}
presigned_url.to_string()
}
pub fn patch_job(
&mut self,
job_uuid: &str,
data: &serde_json::Value,
auth_token: Option<&str>,
) -> crate::Result<HttpResponse> {
let url = self.api_url(&format!("/jobs/{job_uuid}/"));
let mut headers = self.get_auth_headers()?;
if let Some(token) = auth_token {
headers.insert("job-auth-token".to_string(), token.to_string());
}
let body = serde_json::to_vec(data)?;
self.http
.request("PATCH", &url, Some(&body), Some(&headers), 5, None)
}
pub fn delete_job(&mut self, job_uuid: &str) -> crate::Result<HttpResponse> {
self.delete(&format!("/jobs/{job_uuid}/"))
}
pub fn get_job_status_from_compute_node(
&self,
compute_node_url: &str,
job_uuid: &str,
logs: Option<&str>,
) -> crate::Result<JobStatusResponse> {
let mut url = format!("{compute_node_url}/v1/job/{job_uuid}/status/");
if let Some(l) = logs {
url = format!("{url}?logs={l}");
}
let response = HttpClient::shared().request("GET", &url, None, None, 5, Some(30))?;
response.json()
}
pub fn start_multipart_upload(
&mut self,
path: &str,
extra_headers: Option<&HashMap<String, String>>,
) -> crate::Result<()> {
self.post(path, None, extra_headers)?;
Ok(())
}
pub fn get_presigned_upload_url(
&mut self,
path: &str,
part_number: u32,
extra_headers: Option<&HashMap<String, String>>,
) -> crate::Result<String> {
let url = format!("{}?part_number={}", self.api_url(path), part_number,);
let mut headers = self.get_auth_headers()?;
if let Some(extra) = extra_headers {
for (k, v) in extra {
headers.insert(k.clone(), v.clone());
}
}
let response = self
.http
.request("GET", &url, None, Some(&headers), 5, None)?;
let presigned: PresignedUploadResponse = response.json()?;
Ok(self.rewrite_presigned_url(&presigned.presigned_upload_url))
}
pub fn complete_multipart_upload(
&mut self,
path: &str,
parts: &[serde_json::Value],
extra_headers: Option<&HashMap<String, String>>,
) -> crate::Result<()> {
let data = serde_json::json!({
"parts": parts,
});
let body = serde_json::to_vec(&data)?;
self.post(path, Some(&body), extra_headers)?;
Ok(())
}
pub fn upload_to_presigned_url(&self, url: &str, data: &[u8]) -> crate::Result<String> {
let mut headers = HashMap::new();
headers.insert(
"content-type".to_string(),
"application/octet-stream".to_string(),
);
let response = self
.http
.request("PUT", url, Some(data), Some(&headers), 5, None)?;
let etag = response.headers.get("etag").cloned().unwrap_or_default();
Ok(etag)
}
pub fn get_current_user(&mut self) -> crate::Result<UserInfo> {
let response = self.get("/users/me/", None)?;
response.json()
}
pub fn create_or_get_experiment(
&mut self,
uri_or_name: &str,
) -> crate::Result<ExperimentCreateResponse> {
let key = if uri_or_name.contains('/') {
"uri"
} else {
"name"
};
let data = serde_json::json!({key: uri_or_name});
let response = self.post_json("/experiments/", &data)?;
response.json()
}
pub fn get_resource(&mut self, uri_or_name: &str) -> crate::Result<ResourceDetailedDict> {
let key = if uri_or_name.contains('/') {
"uri"
} else {
"name"
};
let response = self.get("/resource/", Some(&[(key, uri_or_name)]))?;
response.json()
}
pub fn get_resource_by_uuid(&mut self, uuid: &str) -> crate::Result<ResourceDetailedDict> {
let response = self.get(&format!("/resources/{uuid}/"), None)?;
response.json()
}
pub fn get_experiments(
&mut self,
page_size: usize,
) -> crate::Result<ExperimentsPaginatedResponse> {
let response = self.get(
"/experiments/",
Some(&[("page_size", &page_size.to_string())]),
)?;
response.json()
}
pub fn add_job_to_experiment(
&mut self,
experiment_uuid: &str,
job_uuid: &str,
) -> crate::Result<()> {
let data = serde_json::json!({"job_uuid": job_uuid});
self.post_json(&format!("/experiments/{experiment_uuid}/jobs/"), &data)?;
Ok(())
}
pub fn get_experiment_jobs(
&mut self,
experiment_uuid: &str,
page_size: usize,
status: Option<&str>,
page: Option<u64>,
) -> crate::Result<JobsPaginatedResponse> {
let mut params: Vec<(&str, String)> = vec![("page_size", page_size.to_string())];
if let Some(s) = status {
params.push(("status", s.to_string()));
}
if let Some(p) = page {
params.push(("page", p.to_string()));
}
let param_refs: Vec<(&str, &str)> = params.iter().map(|(k, v)| (*k, v.as_str())).collect();
let response = self.get(
&format!("/experiments/{experiment_uuid}/jobs/"),
Some(¶m_refs),
)?;
response.json()
}
}
fn urlencoding(s: &str) -> String {
let mut result = String::with_capacity(s.len());
for b in s.bytes() {
match b {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
result.push(b as char);
}
_ => {
result.push('%');
result.push(char::from(b"0123456789ABCDEF"[(b >> 4) as usize]));
result.push(char::from(b"0123456789ABCDEF"[(b & 0x0f) as usize]));
}
}
}
result
}
fn base64_decode(input: &str) -> std::result::Result<Vec<u8>, String> {
const TABLE: &[u8; 128] = b"\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\
\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\
\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\x3e\xff\x3e\xff\x3f\
\x34\x35\x36\x37\x38\x39\x3a\x3b\x3c\x3d\xff\xff\xff\xff\xff\xff\
\xff\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\
\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\xff\xff\xff\xff\x3f\
\xff\x1a\x1b\x1c\x1d\x1e\x1f\x20\x21\x22\x23\x24\x25\x26\x27\x28\
\x29\x2a\x2b\x2c\x2d\x2e\x2f\x30\x31\x32\x33\xff\xff\xff\xff\xff";
let input = input.trim_end_matches('=');
let mut out = Vec::with_capacity(input.len() * 3 / 4);
let mut buf: u32 = 0;
let mut bits: u32 = 0;
for &b in input.as_bytes() {
if b >= 128 {
return Err("invalid base64".to_string());
}
let val = TABLE[b as usize];
if val == 0xff {
return Err("invalid base64".to_string());
}
buf = (buf << 6) | val as u32;
bits += 6;
if bits >= 8 {
bits -= 8;
out.push((buf >> bits) as u8);
buf &= (1 << bits) - 1;
}
}
Ok(out)
}
fn is_token_expired(token: &str) -> bool {
let parts: Vec<&str> = token.split('.').collect();
if parts.len() != 3 {
return true;
}
let payload = match base64_decode(parts[1]) {
Ok(bytes) => bytes,
Err(_) => return true,
};
#[derive(serde::Deserialize)]
struct JwtPayload {
exp: f64,
}
let payload: JwtPayload = match serde_json::from_slice(&payload) {
Ok(p) => p,
Err(_) => return true,
};
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs_f64();
now >= payload.exp - 60.0
}