Skip to main content

elasticctl_api/search/
dsl.rs

1use elasticctl_core::{Error, ErrorKind, Result, Transport};
2use serde_json::{Value, json};
3
4#[derive(Debug, Clone, PartialEq)]
5pub struct DslHit {
6    pub source: Value,
7    pub sort: Option<Vec<Value>>,
8    pub id: Option<String>,
9    pub index: Option<String>,
10    pub score: Option<f64>,
11}
12
13#[derive(Debug, Clone, PartialEq)]
14pub struct DslPage {
15    pub hits: Vec<DslHit>,
16    pub total: Option<u64>,
17}
18
19pub fn decode(value: &Value) -> Result<DslPage> {
20    let hits = value
21        .pointer("/hits/hits")
22        .and_then(Value::as_array)
23        .ok_or_else(|| {
24            Error::new(
25                ErrorKind::Http,
26                "decoding search response field `hits.hits`",
27            )
28        })?;
29    let out = hits
30        .iter()
31        .map(|h| {
32            Ok(DslHit {
33                source: h.get("_source").cloned().unwrap_or(Value::Null),
34                sort: h.get("sort").and_then(Value::as_array).cloned(),
35                id: h.get("_id").and_then(Value::as_str).map(str::to_owned),
36                index: h.get("_index").and_then(Value::as_str).map(str::to_owned),
37                score: h.get("_score").and_then(Value::as_f64),
38            })
39        })
40        .collect::<Result<Vec<_>>>()?;
41    let total = value.pointer("/hits/total/value").and_then(Value::as_u64);
42    Ok(DslPage { hits: out, total })
43}
44
45pub async fn run_sync(t: &Transport, index: &str, body: &Value) -> Result<DslPage> {
46    let response = t
47        .post_absolute_es(&format!("/{index}/_search"), body)
48        .await?;
49    decode(&response)
50}
51
52fn total_sort(sort: &Value) -> Value {
53    let mut entries = match sort {
54        Value::Array(items) => items.clone(),
55        other => vec![other.clone()],
56    };
57    let has_tiebreaker = entries.iter().any(|entry| {
58        entry
59            .as_object()
60            .is_some_and(|obj| obj.contains_key("_shard_doc"))
61    });
62    if !has_tiebreaker {
63        entries.push(json!({"_shard_doc": "asc"}));
64    }
65    Value::Array(entries)
66}
67
68pub async fn run_stream(
69    t: &Transport,
70    index: &str,
71    query: &Value,
72    sort: &Value,
73    limit: Option<usize>,
74) -> Result<Vec<DslHit>> {
75    let open: Value = t
76        .post_absolute_es(&format!("/{index}/_pit?keep_alive=1m"), &json!({}))
77        .await?;
78    let pit_id = open
79        .get("id")
80        .and_then(Value::as_str)
81        .ok_or_else(|| Error::new(ErrorKind::Http, "decoding _pit open response field `id`"))?
82        .to_string();
83
84    let sort = total_sort(sort);
85    let result = page_loop(t, &pit_id, query, &sort, limit).await;
86
87    let _ = t
88        .delete_absolute_es_json("/_pit", &json!({ "id": pit_id }))
89        .await;
90    result
91}
92
93async fn page_loop(
94    t: &Transport,
95    pit_id: &str,
96    query: &Value,
97    sort: &Value,
98    limit: Option<usize>,
99) -> Result<Vec<DslHit>> {
100    let mut all = Vec::new();
101    let mut search_after: Option<Vec<Value>> = None;
102    let size = limit.map_or(1000, |n| n.min(1000));
103    loop {
104        let mut body = json!({
105            "size": size,
106            "sort": sort,
107            "pit": { "id": pit_id, "keep_alive": "1m" },
108            "query": query
109        });
110        if let Some(sa) = &search_after {
111            body["search_after"] = json!(sa);
112        }
113        let page = decode(&t.post_absolute_es("/_search", &body).await?)?;
114        if page.hits.is_empty() {
115            return Ok(all);
116        }
117        let last_sort = page.hits.last().and_then(|h| h.sort.clone());
118        all.extend(page.hits);
119        if let Some(limit) = limit
120            && all.len() >= limit
121        {
122            all.truncate(limit);
123            return Ok(all);
124        }
125        match last_sort {
126            Some(sa) => search_after = Some(sa),
127            None => return Ok(all),
128        }
129    }
130}