use serde::{Deserialize, Serialize};
use std::env;
use std::time::Duration;
fn timeout() -> Duration {
std::env::var("PACKSET_TIMEOUT_MS")
.ok()
.and_then(|v| v.trim().parse::<u64>().ok())
.filter(|ms| *ms > 0)
.map_or(Duration::from_secs(30), Duration::from_millis)
}
fn path_seg(id: &str) -> String {
let mut out = String::with_capacity(id.len());
for b in id.bytes() {
match b {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
out.push(b as char)
}
_ => out.push_str(&format!("%{b:02X}")),
}
}
out
}
pub const DEFAULT_PORT: u16 = 8761;
#[must_use]
pub fn default_port() -> u16 {
env::var("PACKSET_PORT")
.or_else(|_| env::var("GROK_MEM_PORT"))
.ok()
.and_then(|raw| raw.trim().parse().ok())
.unwrap_or(DEFAULT_PORT)
}
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("packset url missing")]
NoUrl,
#[error("http: {0}")]
Http(#[from] Box<ureq::Error>),
#[error("io: {0}")]
Io(#[from] std::io::Error),
#[error("json: {0}")]
Json(#[from] serde_json::Error),
#[error("bad response: {0}")]
Bad(String),
}
#[derive(Debug, Clone)]
pub struct PacksetClient {
base: String,
workspace: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Hit {
pub id: Option<String>,
pub text: String,
#[serde(default)]
pub score: f64,
#[serde(default)]
pub kind: String,
#[serde(default)]
pub ts: Option<String>,
#[serde(default)]
pub ballots: Option<u32>,
#[serde(default)]
pub of: Option<u32>,
}
fn refused(url: &str, e: ureq::Error) -> Error {
match e {
ureq::Error::Status(code, response) => {
let text = response.into_string().unwrap_or_default();
let reason = serde_json::from_str::<serde_json::Value>(&text)
.ok()
.and_then(|v| v.get("error").and_then(|r| r.as_str()).map(str::to_string))
.unwrap_or(text);
let reason = reason.trim();
if reason.is_empty() {
Error::Bad(format!("{url}: status code {code}"))
} else {
Error::Bad(format!("{url}: {code}: {reason}"))
}
}
other => Error::Http(Box::new(other)),
}
}
impl PacksetClient {
pub fn new(base: impl Into<String>) -> Self {
let mut base = base.into();
while base.ends_with('/') {
base.pop();
}
Self {
base,
workspace: None,
}
}
#[must_use]
pub fn with_workspace(mut self, workspace: impl Into<String>) -> Self {
let workspace = workspace.into();
self.workspace = (!workspace.is_empty()).then_some(workspace);
self
}
pub fn from_env() -> Result<Self, Error> {
let url = env::var("PACKSET_URL")
.or_else(|_| env::var("INSIDE_MEMORY_URL"))
.ok()
.filter(|url| !url.is_empty());
match url {
Some(url) if url == "off" => Err(Error::NoUrl),
Some(url) => Ok(Self::new(url)),
None => Ok(Self::new(format!("http://127.0.0.1:{}", default_port()))),
}
}
pub fn base(&self) -> &str {
&self.base
}
pub fn workspace(&self) -> String {
if let Some(w) = &self.workspace {
return w.clone();
}
if let Ok(w) = env::var("PACKSET_WORKSPACE") {
if !w.is_empty() {
return w;
}
}
let cwd = env::var("GROKOS_WORKSPACE")
.ok()
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.map(std::path::PathBuf::from)
.or_else(|| env::current_dir().ok())
.unwrap_or_else(|| std::path::PathBuf::from("."));
self.workspace_for_cwd(&cwd)
}
pub fn workspace_for_cwd(&self, cwd: &std::path::Path) -> String {
let abs = cwd.canonicalize().unwrap_or_else(|_| cwd.to_path_buf());
let url = format!("{}/v1/identity", self.base);
let body = ureq::get(&url)
.query("cwd", abs.to_string_lossy().as_ref())
.timeout(timeout())
.call()
.ok()
.and_then(|r| r.into_string().ok());
if let Some(body) = body {
if let Ok(val) = serde_json::from_str::<serde_json::Value>(&body) {
if let Some(ws) = val.get("workspace").and_then(|v| v.as_str()) {
if !ws.is_empty() {
return ws.to_string();
}
}
}
}
format!("dir:{}", abs.display())
}
pub fn health(&self) -> Result<String, Error> {
let url = format!("{}/health", self.base);
let body = ureq::get(&url)
.timeout(timeout())
.call()
.map_err(|e| refused(&url, e))?
.into_string()?;
Ok(body)
}
pub fn get_atom(&self, workspace: &str, id: &str) -> Result<serde_json::Value, Error> {
let encoded = path_seg(id);
let url = format!("{}/v1/atoms/{encoded}", self.base);
let resp = match ureq::get(&url)
.query("workspace", workspace)
.timeout(timeout())
.call()
{
Ok(resp) => resp,
Err(ureq::Error::Status(404, _)) => {
return Err(Error::Bad(format!("no atom {id}")));
}
Err(e) => return Err(Error::Http(Box::new(e))),
};
Ok(resp.into_json()?)
}
pub fn list_atoms(&self, workspace: &str) -> Result<Vec<serde_json::Value>, Error> {
self.atoms_as_of(workspace, None)
}
pub fn atoms_as_of(
&self,
workspace: &str,
as_of: Option<&str>,
) -> Result<Vec<serde_json::Value>, Error> {
let url = format!("{}/v1/atoms", self.base);
let mut req = ureq::get(&url)
.query("workspace", workspace)
.timeout(timeout());
if let Some(at) = as_of {
req = req.query("as_of", at);
}
let body: serde_json::Value = req.call().map_err(|e| refused(&url, e))?.into_json()?;
let atoms = body
.get("atoms")
.cloned()
.unwrap_or(serde_json::Value::Array(vec![]));
Ok(serde_json::from_value(atoms)?)
}
pub fn search(&self, workspace: &str, q: &str, limit: u32) -> Result<Vec<Hit>, Error> {
self.search_as_of(workspace, q, limit, None)
}
pub fn search_as_of(
&self,
workspace: &str,
q: &str,
limit: u32,
as_of: Option<&str>,
) -> Result<Vec<Hit>, Error> {
self.search_opts(workspace, q, limit, as_of, false)
}
pub fn search_opts(
&self,
workspace: &str,
q: &str,
limit: u32,
as_of: Option<&str>,
rerank: bool,
) -> Result<Vec<Hit>, Error> {
let url = format!("{}/v1/search", self.base);
let budget = if rerank {
Duration::from_secs(60).max(timeout())
} else {
timeout()
};
let mut req = ureq::get(&url)
.query("workspace", workspace)
.query("q", q)
.query("limit", &limit.to_string())
.timeout(budget);
if let Some(at) = as_of {
req = req.query("as_of", at);
}
if rerank {
req = req.query("rerank", "1");
}
let body: serde_json::Value = req.call().map_err(|e| refused(&url, e))?.into_json()?;
let hits = body
.get("hits")
.cloned()
.unwrap_or(serde_json::Value::Array(vec![]));
Ok(serde_json::from_value(hits)?)
}
pub fn status(&self, workspace: Option<&str>) -> Result<serde_json::Value, Error> {
let url = format!("{}/v1/status", self.base);
let mut req = ureq::get(&url).timeout(timeout());
if let Some(workspace) = workspace {
req = req.query("workspace", workspace);
}
Ok(req.call().map_err(|e| refused(&url, e))?.into_json()?)
}
pub fn pin(&self, workspace: &str) -> Result<serde_json::Value, Error> {
let url = format!("{}/v1/pin", self.base);
Ok(ureq::get(&url)
.query("workspace", workspace)
.timeout(timeout())
.call()
.map_err(|e| refused(&url, e))?
.into_json()?)
}
pub fn set_pin(&self, workspace: &str, name: &str) -> Result<serde_json::Value, Error> {
let url = format!("{}/v1/pin", self.base);
Ok(ureq::put(&url)
.timeout(timeout())
.send_json(serde_json::json!({ "workspace": workspace, "name": name }))
.map_err(|e| refused(&url, e))?
.into_json()?)
}
pub fn accessions(&self, workspace: &str) -> Result<Vec<String>, Error> {
let url = format!("{}/v1/accessions", self.base);
let body: serde_json::Value = ureq::get(&url)
.query("workspace", workspace)
.timeout(timeout())
.call()
.map_err(|e| refused(&url, e))?
.into_json()?;
let found = body
.get("accessions")
.cloned()
.unwrap_or(serde_json::Value::Array(vec![]));
Ok(serde_json::from_value(found)?)
}
pub fn atoms(&self, workspace: &str) -> Result<Vec<serde_json::Value>, Error> {
self.atoms_as_of(workspace, None)
}
pub fn citers(
&self,
workspace: &str,
accession: &str,
) -> Result<Vec<serde_json::Value>, Error> {
let url = format!("{}/v1/citers", self.base);
let body: serde_json::Value = ureq::get(&url)
.query("workspace", workspace)
.query("accession", accession)
.timeout(timeout())
.call()
.map_err(|e| refused(&url, e))?
.into_json()?;
let found = body
.get("atoms")
.cloned()
.unwrap_or(serde_json::Value::Array(vec![]));
Ok(serde_json::from_value(found)?)
}
pub fn delete_atom(
&self,
workspace: &str,
id: &str,
why: Option<&str>,
) -> Result<serde_json::Value, Error> {
let url = format!("{}/v1/atoms/delete", self.base);
let mut body = serde_json::json!({
"workspace": workspace,
"id": id,
});
if let Some(accession) = why {
body["why"] = serde_json::Value::String(accession.to_string());
}
let resp = match ureq::post(&url).timeout(timeout()).send_json(body) {
Ok(resp) => resp,
Err(ureq::Error::Status(404, _)) => {
return Err(Error::Bad(format!("no atom {id}")));
}
Err(e) => return Err(refused(&url, e)),
};
Ok(resp.into_json()?)
}
pub fn grade(
&self,
workspace: &str,
id: &str,
recalled: bool,
) -> Result<serde_json::Value, Error> {
let url = format!("{}/v1/grade", self.base);
let body: serde_json::Value = ureq::post(&url)
.timeout(timeout())
.send_json(serde_json::json!({
"workspace": workspace,
"id": id,
"recalled": recalled,
}))
.map_err(|e| refused(&url, e))?
.into_json()?;
Ok(body)
}
pub fn hubs(&self, workspace: &str, limit: usize) -> Result<serde_json::Value, Error> {
let url = format!("{}/v1/hubs", self.base);
let body: serde_json::Value = ureq::get(&url)
.query("workspace", workspace)
.query("limit", &limit.to_string())
.timeout(timeout())
.call()
.map_err(|e| refused(&url, e))?
.into_json()?;
Ok(body)
}
pub fn islands(&self, workspace: &str) -> Result<serde_json::Value, Error> {
let url = format!("{}/v1/islands", self.base);
let body: serde_json::Value = ureq::get(&url)
.query("workspace", workspace)
.timeout(timeout())
.call()
.map_err(|e| refused(&url, e))?
.into_json()?;
Ok(body)
}
pub fn fire(&self, workspace: &str, ids: &[String]) -> Result<serde_json::Value, Error> {
let url = format!("{}/v1/fire", self.base);
let body: serde_json::Value = ureq::post(&url)
.timeout(timeout())
.send_json(serde_json::json!({"workspace": workspace, "ids": ids}))
.map_err(|e| refused(&url, e))?
.into_json()?;
Ok(body)
}
pub fn consolidate(&self, workspace: &str, apply: bool) -> Result<serde_json::Value, Error> {
let url = format!("{}/v1/consolidate", self.base);
let body: serde_json::Value = ureq::post(&url)
.timeout(timeout())
.send_json(serde_json::json!({"workspace": workspace, "apply": apply}))
.map_err(|e| refused(&url, e))?
.into_json()?;
Ok(body)
}
pub fn activate(
&self,
workspace: &str,
q: &str,
limit: u32,
fire: bool,
) -> Result<serde_json::Value, Error> {
let url = format!("{}/v1/activate", self.base);
let body: serde_json::Value = ureq::get(&url)
.query("workspace", workspace)
.query("q", q)
.query("limit", &limit.to_string())
.query("fire", if fire { "1" } else { "0" })
.timeout(timeout())
.call()
.map_err(|e| refused(&url, e))?
.into_json()?;
Ok(body)
}
pub fn post_atom(&self, atom: &serde_json::Value) -> Result<serde_json::Value, Error> {
let url = format!("{}/v1/atoms", self.base);
let body: serde_json::Value = ureq::post(&url)
.timeout(timeout())
.send_json(atom.clone())
.map_err(|e| refused(&url, e))?
.into_json()?;
Ok(body)
}
}