diff --git a/bkapi-client/src/lib.rs b/bkapi-client/src/lib.rs index 8b7d05c..2cc8af7 100644 --- a/bkapi-client/src/lib.rs +++ b/bkapi-client/src/lib.rs @@ -111,6 +111,77 @@ impl InjectContext for reqwest::RequestBuilder { } } +/// The BKApi client, operating over NATS instead of HTTP. +#[derive(Clone)] +pub struct BKApiNatsClient { + client: async_nats::Client, +} + +/// A hash and distance. +#[derive(serde::Serialize, serde::Deserialize)] +pub struct HashDistance { + hash: i64, + distance: u32, +} + +impl BKApiNatsClient { + const NATS_SUBJECT: &str = "bkapi.search"; + + /// Create a new client with a given NATS client. + pub fn new(client: async_nats::Client) -> Self { + Self { client } + } + + /// Search for a single hash. + pub async fn search( + &self, + hash: i64, + distance: i32, + ) -> Result { + let hashes = [HashDistance { + hash, + distance: distance as u32, + }]; + + self.search_many(&hashes) + .await + .map(|mut results| results.remove(0)) + } + + /// Search many hashes at once. + pub async fn search_many( + &self, + hashes: &[HashDistance], + ) -> Result, async_nats::Error> { + let payload = serde_json::to_vec(hashes).unwrap(); + + let message = self + .client + .request(Self::NATS_SUBJECT.to_string(), payload.into()) + .await?; + + let results: Vec> = serde_json::from_slice(&message.payload).unwrap(); + + let results = results + .into_iter() + .zip(hashes) + .map(|(results, search)| SearchResults { + hash: search.hash, + distance: search.distance as u64, + hashes: results + .into_iter() + .map(|result| SearchResult { + hash: result.hash, + distance: result.distance as u64, + }) + .collect(), + }) + .collect(); + + Ok(results) + } +} + #[cfg(test)] mod tests { fn get_test_endpoint() -> String {