use crate::error::{Result, SearchError};
use crate::query::Search;
use crate::response::SearchResults;
use rustlavel_client::Client;
use rustlavel_core::{Config, Json};
use rustlavel_http::Method;
use std::time::Duration;
#[derive(Clone)]
enum Auth {
None,
Header(String),
}
#[derive(Clone)]
pub struct SearchClient {
base_url: String,
auth: Auth,
http: Client,
retries: u32,
}
impl SearchClient {
pub fn new(base_url: impl Into<String>) -> SearchClient {
SearchClient {
base_url: base_url.into().trim_end_matches('/').to_string(),
auth: Auth::None,
http: Client::new().timeout(Duration::from_secs(30)),
retries: 2,
}
}
pub fn from_config(config: &Config) -> SearchClient {
let mut client = SearchClient::new(config.string("search.url", "http://localhost:9200"));
let username = config.string("search.username", "");
let api_key = config.string("search.api_key", "");
let token = config.string("search.token", "");
if !username.is_empty() {
client = client.basic_auth(&username, &config.string("search.password", ""));
} else if !api_key.is_empty() {
client = client.api_key(api_key);
} else if !token.is_empty() {
client = client.bearer(token);
}
let timeout_ms = config.int("search.timeout_ms", 0);
if timeout_ms > 0 {
client = client.timeout(Duration::from_millis(timeout_ms as u64));
}
let retries = config.int("search.retries", -1);
if retries >= 0 {
client = client.retries(retries.min(10) as u32);
}
client
}
pub fn basic_auth(mut self, username: &str, password: &str) -> SearchClient {
let encoded = base64(format!("{username}:{password}").as_bytes());
self.auth = Auth::Header(format!("Basic {encoded}"));
self
}
pub fn api_key(mut self, encoded: impl Into<String>) -> SearchClient {
self.auth = Auth::Header(format!("ApiKey {}", encoded.into()));
self
}
pub fn bearer(mut self, token: impl Into<String>) -> SearchClient {
self.auth = Auth::Header(format!("Bearer {}", token.into()));
self
}
pub fn timeout(mut self, timeout: Duration) -> SearchClient {
self.http = self.http.timeout(timeout);
self
}
pub fn retries(mut self, retries: u32) -> SearchClient {
self.retries = retries;
self
}
pub fn faking(mut self, fake: rustlavel_client::Fake) -> SearchClient {
self.http = self.http.faking(fake);
self
}
pub fn fake(&self) -> Option<&std::sync::Arc<rustlavel_client::Fake>> {
self.http.fake()
}
pub fn base_url(&self) -> &str {
&self.base_url
}
pub fn url(&self, path: &str) -> String {
format!("{}/{}", self.base_url, path.trim_start_matches('/'))
}
pub async fn health(&self) -> Result<ClusterHealth> {
let body = self.json(Method::Get, "_cluster/health", "_cluster/health", None).await?;
ClusterHealth::parse(&body)
}
pub async fn info(&self) -> Result<ClusterInfo> {
let body = self.json(Method::Get, "", "/", None).await?;
Ok(ClusterInfo {
name: string_at(&body, "name"),
cluster_name: string_at(&body, "cluster_name"),
version: string_at(&body, "version.number"),
distribution: body
.get("version.distribution")
.and_then(Json::as_str)
.unwrap_or("elasticsearch")
.to_string(),
})
}
pub async fn search(&self, index: &str, search: &Search) -> Result<SearchResults> {
let path = format!("{}/_search", encode(index));
let body = self.json(Method::Post, &path, index, Some(search.body())).await?;
SearchResults::parse(&body)
}
pub async fn count(&self, index: &str, search: &Search) -> Result<u64> {
let path = format!("{}/_count", encode(index));
let body = self.json(Method::Post, &path, index, Some(search.count_body())).await?;
body.get("count")
.and_then(Json::as_i64)
.map(|count| count.max(0) as u64)
.ok_or_else(|| SearchError::Malformed {
context: path,
message: "the reply has no `count`".to_string(),
})
}
pub(crate) async fn json(
&self,
method: Method,
path: &str,
context: &str,
body: Option<Json>,
) -> Result<Json> {
let raw = self.send(method, path, context, body).await?;
if !raw.is_success() {
return Err(raw.error(context));
}
raw.into_json(path)
}
pub(crate) async fn ndjson(&self, path: &str, context: &str, body: String) -> Result<Json> {
let raw = self
.send_payload(Method::Post, path, context, Some(Payload::Ndjson(body)))
.await?;
if !raw.is_success() {
return Err(raw.error(context));
}
raw.into_json(path)
}
pub(crate) async fn exists_at(&self, path: &str, context: &str) -> Result<bool> {
let raw = self.send(Method::Get, path, context, None).await?;
match raw.status {
_ if raw.is_success() => Ok(true),
404 => Ok(false),
_ => Err(raw.error(context)),
}
}
pub(crate) async fn send(
&self,
method: Method,
path: &str,
context: &str,
body: Option<Json>,
) -> Result<Raw> {
self.send_payload(method, path, context, body.map(Payload::Json)).await
}
async fn send_payload(
&self,
method: Method,
path: &str,
context: &str,
body: Option<Payload>,
) -> Result<Raw> {
let mut attempt = 0;
loop {
let (failed, error) = match self.send_once(method, path, body.clone()).await {
Ok(raw) if raw.is_success() => return Ok(raw),
Ok(raw) => {
let error = raw.error(context);
(Some(raw), error)
}
Err(error) => (None, error),
};
let safe = matches!(method, Method::Get | Method::Head);
let worth_retrying = match &error {
SearchError::Transport(_) => safe,
other => other.is_retryable(),
};
if attempt >= self.retries || !worth_retrying {
return match failed {
Some(raw) => Ok(raw),
None => Err(error),
};
}
tokio::time::sleep(Duration::from_millis(200 * 2u64.pow(attempt))).await;
attempt += 1;
}
}
async fn send_once(&self, method: Method, path: &str, body: Option<Payload>) -> Result<Raw> {
let mut request = self.http.request(method, self.url(path));
if let Auth::Header(value) = &self.auth {
request = request.header("authorization", value.clone());
}
match body {
Some(Payload::Json(value)) => request = request.json(value),
Some(Payload::Ndjson(text)) => {
request = request.header("content-type", "application/x-ndjson").body(text)
}
None => {}
}
let response =
request.send().await.map_err(|e| SearchError::Transport(e.to_string()))?;
Ok(Raw { status: response.status.code(), body: response.text() })
}
}
impl std::fmt::Debug for SearchClient {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SearchClient")
.field("base_url", &self.base_url)
.field(
"auth",
&match &self.auth {
Auth::None => "<none>",
Auth::Header(_) => "<redacted>",
},
)
.field("retries", &self.retries)
.finish()
}
}
#[derive(Clone)]
enum Payload {
Json(Json),
Ndjson(String),
}
pub(crate) struct Raw {
pub(crate) status: u16,
pub(crate) body: String,
}
impl Raw {
pub(crate) fn is_success(&self) -> bool {
(200..300).contains(&self.status)
}
pub(crate) fn error(&self, context: &str) -> SearchError {
SearchError::from_response(self.status, context, &self.body)
}
pub(crate) fn into_json(self, path: &str) -> Result<Json> {
if self.body.trim().is_empty() {
return Ok(Json::Null);
}
Json::parse(&self.body).map_err(|e| SearchError::Malformed {
context: path.to_string(),
message: e.to_string(),
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClusterHealth {
pub cluster_name: String,
pub status: HealthStatus,
pub number_of_nodes: u32,
pub active_shards: u32,
pub unassigned_shards: u32,
pub timed_out: bool,
}
impl ClusterHealth {
fn parse(body: &Json) -> Result<ClusterHealth> {
let status = body.get("status").and_then(Json::as_str).ok_or_else(|| {
SearchError::Malformed {
context: "_cluster/health".to_string(),
message: "the reply has no `status`".to_string(),
}
})?;
Ok(ClusterHealth {
cluster_name: string_at(body, "cluster_name"),
status: HealthStatus::parse(status),
number_of_nodes: number_at(body, "number_of_nodes"),
active_shards: number_at(body, "active_shards"),
unassigned_shards: number_at(body, "unassigned_shards"),
timed_out: body.get("timed_out").and_then(Json::as_bool).unwrap_or(false),
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HealthStatus {
Green,
Yellow,
Red,
Unknown,
}
impl HealthStatus {
fn parse(value: &str) -> HealthStatus {
match value {
"green" => HealthStatus::Green,
"yellow" => HealthStatus::Yellow,
"red" => HealthStatus::Red,
_ => HealthStatus::Unknown,
}
}
pub fn is_usable(self) -> bool {
matches!(self, HealthStatus::Green | HealthStatus::Yellow)
}
}
impl std::fmt::Display for HealthStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(match self {
HealthStatus::Green => "green",
HealthStatus::Yellow => "yellow",
HealthStatus::Red => "red",
HealthStatus::Unknown => "unknown",
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClusterInfo {
pub name: String,
pub cluster_name: String,
pub version: String,
pub distribution: String,
}
impl ClusterInfo {
pub fn is_opensearch(&self) -> bool {
self.distribution == "opensearch"
}
}
fn string_at(body: &Json, path: &str) -> String {
body.get(path).and_then(Json::as_str).unwrap_or_default().to_string()
}
fn number_at(body: &Json, path: &str) -> u32 {
body.get(path).and_then(Json::as_i64).unwrap_or(0).max(0) as u32
}
pub(crate) fn encode(segment: &str) -> String {
let mut out = String::with_capacity(segment.len());
for byte in segment.bytes() {
match byte {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' | b',' | b'*' => {
out.push(byte as char)
}
_ => out.push_str(&format!("%{byte:02X}")),
}
}
out
}
fn base64(input: &[u8]) -> String {
const ALPHABET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
let mut out = String::with_capacity(input.len().div_ceil(3) * 4);
for chunk in input.chunks(3) {
let b = [chunk[0], *chunk.get(1).unwrap_or(&0), *chunk.get(2).unwrap_or(&0)];
let bits = ((b[0] as u32) << 16) | ((b[1] as u32) << 8) | b[2] as u32;
out.push(ALPHABET[(bits >> 18) as usize & 63] as char);
out.push(ALPHABET[(bits >> 12) as usize & 63] as char);
out.push(if chunk.len() > 1 { ALPHABET[(bits >> 6) as usize & 63] as char } else { '=' });
out.push(if chunk.len() > 2 { ALPHABET[bits as usize & 63] as char } else { '=' });
}
out
}
#[cfg(test)]
mod tests {
use super::*;
use rustlavel_client::{Fake, FakeResponse};
#[test]
fn from_config_reads_the_address_and_the_first_credential_it_finds() {
let config = Config::new();
config.set("search.url", "https://es.internal:9200/");
config.set("search.username", "elastic");
config.set("search.password", "s3cret");
config.set("search.api_key", "ignored-because-basic-wins");
config.set("search.retries", Json::from(2_i64));
let client = SearchClient::from_config(&config);
assert_eq!(client.base_url, "https://es.internal:9200", "trailing slash trimmed");
assert!(matches!(&client.auth, Auth::Header(h) if h.starts_with("Basic ")));
assert_eq!(client.retries, 2);
}
#[test]
fn from_config_with_nothing_set_points_at_localhost_unauthenticated() {
let client = SearchClient::from_config(&Config::new());
assert_eq!(client.base_url, "http://localhost:9200");
assert!(matches!(client.auth, Auth::None));
}
fn json(body: &str) -> Json {
Json::parse(body).expect("valid JSON in a test")
}
#[test]
fn builds_urls_whatever_the_slashes() {
let client = SearchClient::new("http://localhost:9200/");
assert_eq!(client.url("posts/_search"), "http://localhost:9200/posts/_search");
assert_eq!(client.url("/posts/_search"), "http://localhost:9200/posts/_search");
}
#[test]
fn debug_never_prints_the_password() {
let client = SearchClient::new("http://localhost:9200").basic_auth("elastic", "hunter2");
let printed = format!("{client:?}");
assert!(!printed.contains("hunter2"), "the password reached a log: {printed}");
assert!(!printed.contains("ZWxhc3RpYzpodW50ZXIy"), "got {printed}");
assert!(printed.contains("redacted"));
}
#[test]
fn debug_distinguishes_no_credentials_from_hidden_ones() {
let printed = format!("{:?}", SearchClient::new("http://localhost:9200"));
assert!(printed.contains("none"), "got {printed}");
}
#[test]
fn base64_matches_the_worked_example_in_rfc_4648() {
assert_eq!(base64(b""), "");
assert_eq!(base64(b"f"), "Zg==");
assert_eq!(base64(b"fo"), "Zm8=");
assert_eq!(base64(b"foo"), "Zm9v");
assert_eq!(base64(b"foob"), "Zm9vYg==");
assert_eq!(base64(b"fooba"), "Zm9vYmE=");
assert_eq!(base64(b"foobar"), "Zm9vYmFy");
assert_eq!(base64(b"elastic:changeme"), "ZWxhc3RpYzpjaGFuZ2VtZQ==");
}
#[test]
fn an_id_cannot_change_which_endpoint_is_called() {
assert_eq!(encode("a/b"), "a%2Fb");
assert_eq!(encode("with space"), "with%20space");
assert_eq!(encode("q?x=1"), "q%3Fx%3D1");
assert_eq!(encode("plain-id_1.2~3"), "plain-id_1.2~3");
assert_eq!(encode("logs-*,metrics-*"), "logs-*,metrics-*");
}
#[tokio::test]
async fn sends_the_authorization_header_basic_auth_builds() {
let client = SearchClient::new("http://localhost:9200")
.basic_auth("elastic", "changeme")
.faking(Fake::new().on("localhost:9200", FakeResponse::json(json(r#"{"status":"green"}"#))));
client.health().await.unwrap();
let sent = client.fake().unwrap().recorded();
assert_eq!(sent[0].headers.get("authorization"), Some("Basic ZWxhc3RpYzpjaGFuZ2VtZQ=="));
}
#[tokio::test]
async fn an_api_key_is_passed_through_rather_than_encoded_again() {
let client = SearchClient::new("http://localhost:9200")
.api_key("VnVhQ2ZHY0JDZGJrUW0tZTVhT3g6dWkybHAyYXhUTm1zeWFrdzl0dk5udw==")
.faking(Fake::new().fallback(FakeResponse::json(json(r#"{"status":"green"}"#))));
client.health().await.unwrap();
assert_eq!(
client.fake().unwrap().recorded()[0].headers.get("authorization"),
Some("ApiKey VnVhQ2ZHY0JDZGJrUW0tZTVhT3g6dWkybHAyYXhUTm1zeWFrdzl0dk5udw==")
);
}
#[tokio::test]
async fn no_credentials_means_no_header_at_all() {
let client = SearchClient::new("http://localhost:9200")
.faking(Fake::new().fallback(FakeResponse::json(json(r#"{"status":"green"}"#))));
client.health().await.unwrap();
assert_eq!(client.fake().unwrap().recorded()[0].headers.get("authorization"), None);
}
#[tokio::test]
async fn reads_the_health_a_single_node_cluster_actually_reports() {
let body = r#"{"cluster_name":"docker-cluster","status":"yellow","timed_out":false,"number_of_nodes":1,"number_of_data_nodes":1,"active_primary_shards":1,"active_shards":1,"relocating_shards":0,"initializing_shards":0,"unassigned_shards":1,"delayed_unassigned_shards":0,"number_of_pending_tasks":0,"number_of_in_flight_fetch":0,"task_max_waiting_in_queue_millis":0,"active_shards_percent_as_number":50.0}"#;
let client = SearchClient::new("http://localhost:9200")
.faking(Fake::new().on("_cluster/health", FakeResponse::json(json(body))));
let health = client.health().await.unwrap();
assert_eq!(health.status, HealthStatus::Yellow);
assert_eq!(health.cluster_name, "docker-cluster");
assert_eq!(health.unassigned_shards, 1);
assert!(health.status.is_usable(), "yellow is the normal state of a one-node cluster");
assert!(!HealthStatus::Red.is_usable());
}
#[tokio::test]
async fn tells_elasticsearch_and_opensearch_apart_without_preferring_either() {
let elastic = r#"{"name":"e9f1c933ed49","cluster_name":"docker-cluster","cluster_uuid":"n3-0v9hLR4mBTvOqPRVFmA","version":{"number":"8.15.0","build_flavor":"default","build_type":"docker","build_hash":"1a77947f34deddb41af25e6f0ddb8e830159c179","build_date":"2024-08-05T10:05:34.233336849Z","build_snapshot":false,"lucene_version":"9.11.1","minimum_wire_compatibility_version":"7.17.0","minimum_index_compatibility_version":"7.0.0"},"tagline":"You Know, for Search"}"#;
let opensearch = r#"{"name":"opensearch-node1","cluster_name":"docker-cluster","cluster_uuid":"C1v1zpcYRW6cLtLBEIyOZg","version":{"distribution":"opensearch","number":"2.11.0","build_type":"tar","build_hash":"4dcad6dd1fd45b6bd91f041a041829c8687278fa","build_date":"2023-10-13T02:57:04.940191475Z","build_snapshot":false,"lucene_version":"9.7.0","minimum_wire_compatibility_version":"7.10.0","minimum_index_compatibility_version":"7.0.0"},"tagline":"The OpenSearch Project: https://opensearch.org/"}"#;
for (body, distribution, version) in [
(elastic, "elasticsearch", "8.15.0"),
(opensearch, "opensearch", "2.11.0"),
] {
let client = SearchClient::new("http://localhost:9200")
.faking(Fake::new().fallback(FakeResponse::json(json(body))));
let info = client.info().await.unwrap();
assert_eq!(info.distribution, distribution);
assert_eq!(info.version, version);
assert_eq!(info.is_opensearch(), distribution == "opensearch");
}
}
#[tokio::test]
async fn a_rejected_request_is_retried_and_a_mapping_error_is_not() {
let rejected = r#"{"error":{"root_cause":[{"type":"circuit_breaking_exception","reason":"[parent] Data too large","bytes_wanted":1,"bytes_limit":0,"durability":"TRANSIENT"}],"type":"circuit_breaking_exception","reason":"[parent] Data too large","bytes_wanted":1,"bytes_limit":0,"durability":"TRANSIENT"},"status":429}"#;
let broken = r#"{"error":{"root_cause":[{"type":"parsing_exception","reason":"unknown query [matchh]","line":1,"col":30}],"type":"parsing_exception","reason":"unknown query [matchh]","line":1,"col":30},"status":400}"#;
let client = SearchClient::new("http://localhost:9200")
.retries(2)
.faking(Fake::new().fallback(FakeResponse::text(rejected).status(429)));
assert!(client.search("posts", &Search::new()).await.is_err());
assert_eq!(client.fake().unwrap().count(), 3, "the first attempt plus two retries");
let client = SearchClient::new("http://localhost:9200")
.retries(5)
.faking(Fake::new().fallback(FakeResponse::text(broken).status(400)));
assert!(client.search("posts", &Search::new()).await.is_err());
assert_eq!(client.fake().unwrap().count(), 1, "a broken query cannot be fixed by repeating it");
}
#[tokio::test]
async fn a_search_against_a_missing_index_says_so_by_name() {
let body = r#"{"error":{"root_cause":[{"type":"index_not_found_exception","reason":"no such index [pots]","resource.type":"index_or_alias","resource.id":"pots","index_uuid":"_na_","index":"pots"}],"type":"index_not_found_exception","reason":"no such index [pots]","resource.type":"index_or_alias","resource.id":"pots","index_uuid":"_na_","index":"pots"},"status":404}"#;
let client = SearchClient::new("http://localhost:9200")
.faking(Fake::new().fallback(FakeResponse::text(body).status(404)));
let error = client.search("pots", &Search::new()).await.unwrap_err();
assert!(error.is_index_not_found());
assert!(error.to_string().contains("pots"), "got {error}");
assert!(!error.is_retryable(), "the index will not appear by itself");
}
#[tokio::test]
async fn counting_does_not_send_the_fields_count_refuses() {
let client = SearchClient::new("http://localhost:9200")
.faking(Fake::new().fallback(FakeResponse::json(json(r#"{"count":42,"_shards":{"total":1,"successful":1,"skipped":0,"failed":0}}"#))));
let search = Search::new().size(10).from(20).sort_desc("created_at");
assert_eq!(client.count("posts", &search).await.unwrap(), 42);
let sent = client.fake().unwrap().recorded()[0].json().unwrap();
assert!(sent.get("size").is_none(), "size would make _count answer 400");
assert!(sent.get("sort").is_none(), "sort would make _count answer 400");
assert!(sent.get("query").is_some());
}
}