datafusion_opensearch/
client.rs

1//! The HTTP client the provider talks to OpenSearch through.
2//!
3//! OpenSearch's REST API is stateless, so there is no connection pool to manage: a
4//! [`reqwest::Client`] (which pools connections internally) plus the cluster's base URL is the
5//! whole "connection". Add credentials with [`with_basic_auth`](OpenSearchClient::with_basic_auth)
6//! or [`with_bearer_token`](OpenSearchClient::with_bearer_token), or bring your own client via
7//! [`with_client`](OpenSearchClient::with_client) for timeouts, TLS roots and custom headers.
8
9use serde_json::{json, Value};
10
11use crate::{DecodeSnafu, Error, HttpSnafu, IndexNotFoundSnafu, Result};
12use snafu::ResultExt;
13
14#[derive(Clone, Debug)]
15enum Auth {
16    None,
17    Basic { user: String, password: String },
18    Bearer(String),
19}
20
21/// A handle to one OpenSearch cluster.
22#[derive(Clone, Debug)]
23pub struct OpenSearchClient {
24    base_url: String,
25    http: reqwest::Client,
26    auth: Auth,
27}
28
29/// One page of search hits: the `_source` documents plus the scroll cursor, if one is open.
30#[derive(Clone, Debug, Default)]
31pub struct SearchPage {
32    /// Each hit's `_source`.
33    pub hits: Vec<Value>,
34    /// The `_scroll_id` to continue with, when the search opened a scroll.
35    pub scroll_id: Option<String>,
36}
37
38impl OpenSearchClient {
39    /// `base_url` e.g. `http://localhost:9200` (a trailing slash is tolerated).
40    pub fn new(base_url: impl Into<String>) -> Self {
41        Self {
42            base_url: base_url.into(),
43            http: reqwest::Client::new(),
44            auth: Auth::None,
45        }
46    }
47
48    /// Use a preconfigured [`reqwest::Client`] (timeouts, connection limits, TLS roots, headers).
49    pub fn with_client(mut self, http: reqwest::Client) -> Self {
50        self.http = http;
51        self
52    }
53
54    /// Send HTTP basic credentials with every request (the OpenSearch security plugin's default).
55    pub fn with_basic_auth(mut self, user: impl Into<String>, password: impl Into<String>) -> Self {
56        self.auth = Auth::Basic {
57            user: user.into(),
58            password: password.into(),
59        };
60        self
61    }
62
63    /// Send a bearer token with every request.
64    pub fn with_bearer_token(mut self, token: impl Into<String>) -> Self {
65        self.auth = Auth::Bearer(token.into());
66        self
67    }
68
69    /// The cluster base URL, without a trailing slash.
70    pub fn base_url(&self) -> &str {
71        self.base_url.trim_end_matches('/')
72    }
73
74    fn get(&self, url: &str) -> reqwest::RequestBuilder {
75        self.authed(self.http.get(url))
76    }
77
78    fn post(&self, url: &str) -> reqwest::RequestBuilder {
79        self.authed(self.http.post(url))
80    }
81
82    fn delete(&self, url: &str) -> reqwest::RequestBuilder {
83        self.authed(self.http.delete(url))
84    }
85
86    fn authed(&self, rb: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
87        match &self.auth {
88            Auth::None => rb,
89            Auth::Basic { user, password } => rb.basic_auth(user, Some(password)),
90            Auth::Bearer(t) => rb.bearer_auth(t),
91        }
92    }
93
94    /// Run one `_search` against `index` and return each hit's `_source`. A missing index (404)
95    /// yields an empty result rather than an error, so a table over an index that has not been
96    /// created yet simply reads as empty.
97    pub async fn search(&self, index: &str, body: &Value) -> Result<Vec<Value>> {
98        Ok(self.search_page(index, body, None).await?.hits)
99    }
100
101    /// Run a `_search`, optionally opening a scroll cursor (`keep_alive` e.g. `"1m"`) so the rest
102    /// of the result can be paged with [`scroll`](Self::scroll). A missing index reads as empty.
103    pub async fn search_page(&self, index: &str, body: &Value, keep_alive: Option<&str>) -> Result<SearchPage> {
104        let url = match keep_alive {
105            Some(k) => format!("{}/{index}/_search?scroll={k}", self.base_url()),
106            None => format!("{}/{index}/_search", self.base_url()),
107        };
108        let resp = self
109            .post(&url)
110            .json(body)
111            .send()
112            .await
113            .context(HttpSnafu { url: url.clone() })?;
114        if resp.status().as_u16() == 404 {
115            return Ok(SearchPage::default());
116        }
117        let resp = check_status(resp, &url).await?;
118        let doc: Value = resp.json().await.context(DecodeSnafu { url })?;
119        Ok(page_from(doc))
120    }
121
122    /// Fetch the next page of an open scroll.
123    pub async fn scroll(&self, scroll_id: &str, keep_alive: &str) -> Result<SearchPage> {
124        let url = format!("{}/_search/scroll", self.base_url());
125        let resp = self
126            .post(&url)
127            .json(&json!({ "scroll": keep_alive, "scroll_id": scroll_id }))
128            .send()
129            .await
130            .context(HttpSnafu { url: url.clone() })?;
131        let resp = check_status(resp, &url).await?;
132        let doc: Value = resp.json().await.context(DecodeSnafu { url })?;
133        Ok(page_from(doc))
134    }
135
136    /// Release a scroll cursor early (best effort — the server expires it anyway).
137    pub async fn clear_scroll(&self, scroll_id: &str) {
138        let url = format!("{}/_search/scroll", self.base_url());
139        let _ = self
140            .delete(&url)
141            .json(&json!({ "scroll_id": [scroll_id] }))
142            .send()
143            .await;
144    }
145
146    /// Fetch the index `_mapping` (the raw response, keyed by index name). A missing index is
147    /// [`Error::IndexNotFound`] — unlike [`search`](Self::search), a caller asking for the schema of
148    /// an index that does not exist has nothing sensible to fall back to.
149    pub async fn mapping(&self, index: &str) -> Result<Value> {
150        let url = format!("{}/{index}/_mapping", self.base_url());
151        let resp = self.get(&url).send().await.context(HttpSnafu { url: url.clone() })?;
152        if resp.status().as_u16() == 404 {
153            return IndexNotFoundSnafu { index }.fail();
154        }
155        let resp = check_status(resp, &url).await?;
156        resp.json().await.context(DecodeSnafu { url })
157    }
158
159    /// The names of the cluster's indices, sorted, excluding system indices (names starting with
160    /// `.`). Aliases are not included.
161    pub async fn indices(&self) -> Result<Vec<String>> {
162        let url = format!("{}/_cat/indices?format=json&h=index", self.base_url());
163        let resp = self.get(&url).send().await.context(HttpSnafu { url: url.clone() })?;
164        let resp = check_status(resp, &url).await?;
165        let doc: Value = resp.json().await.context(DecodeSnafu { url })?;
166        let mut names: Vec<String> = doc
167            .as_array()
168            .map(|rows| {
169                rows.iter()
170                    .filter_map(|r| r["index"].as_str())
171                    .filter(|n| !n.starts_with('.'))
172                    .map(str::to_string)
173                    .collect()
174            })
175            .unwrap_or_default();
176        names.sort();
177        Ok(names)
178    }
179}
180
181fn page_from(doc: Value) -> SearchPage {
182    let hits = doc["hits"]["hits"]
183        .as_array()
184        .map(|hits| hits.iter().map(|h| h["_source"].clone()).collect())
185        .unwrap_or_default();
186    let scroll_id = doc["_scroll_id"].as_str().map(str::to_string);
187    SearchPage { hits, scroll_id }
188}
189
190/// Turn a non-2xx response into [`Error::Status`], carrying the body for diagnosis.
191async fn check_status(resp: reqwest::Response, url: &str) -> Result<reqwest::Response> {
192    if resp.status().is_success() {
193        return Ok(resp);
194    }
195    let status = resp.status().as_u16();
196    let body = resp.text().await.unwrap_or_default();
197    Err(Error::Status {
198        status,
199        url: url.to_string(),
200        body,
201    })
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207
208    #[test]
209    fn base_url_drops_a_trailing_slash() {
210        assert_eq!(OpenSearchClient::new("http://x:9200/").base_url(), "http://x:9200");
211        assert_eq!(OpenSearchClient::new("http://x:9200").base_url(), "http://x:9200");
212    }
213
214    #[test]
215    fn auth_is_recorded() {
216        let c = OpenSearchClient::new("http://x").with_basic_auth("admin", "pw");
217        assert!(matches!(c.auth, Auth::Basic { .. }));
218        let c = OpenSearchClient::new("http://x").with_bearer_token("t");
219        assert!(matches!(c.auth, Auth::Bearer(_)));
220    }
221
222    #[test]
223    fn page_parses_hits_and_scroll_id() {
224        let p = page_from(json!({ "_scroll_id": "abc", "hits": { "hits": [ { "_source": { "a": 1 } } ] } }));
225        assert_eq!(p.hits, vec![json!({ "a": 1 })]);
226        assert_eq!(p.scroll_id.as_deref(), Some("abc"));
227        let empty = page_from(json!({ "hits": { "hits": [] } }));
228        assert!(empty.hits.is_empty() && empty.scroll_id.is_none());
229    }
230}
datafusion_opensearch/client.rs

datafusion_opensearch/
client.rs
