elasticctl_api/search/
esql.rs

1use crate::compose::require_path_segment;
2use elasticctl_core::{Error, ErrorKind, Result, Transport, urlencode};
3use serde_json::{Map, Value};
4use std::time::Duration;
5
6#[derive(Debug, Clone, PartialEq)]
7pub struct EsqlColumn {
8    pub name: String,
9    pub r#type: String,
10}
11
12#[derive(Debug, Clone, PartialEq)]
13pub struct EsqlResponse {
14    pub columns: Vec<EsqlColumn>,
15    pub values: Vec<Vec<Value>>,
16    pub is_partial: bool,
17}
18
19fn parse_columns(obj: &Map<String, Value>) -> Result<Vec<EsqlColumn>> {
20    let columns = obj
21        .get("columns")
22        .and_then(Value::as_array)
23        .ok_or_else(|| Error::new(ErrorKind::Http, "decoding esql response field `columns`"))?;
24    let mut out = Vec::with_capacity(columns.len());
25    for (i, col) in columns.iter().enumerate() {
26        let name = col.get("name").and_then(Value::as_str).ok_or_else(|| {
27            Error::new(
28                ErrorKind::Http,
29                format!("decoding esql response column {i} field `name`"),
30            )
31        })?;
32        let ty = col.get("type").and_then(Value::as_str).ok_or_else(|| {
33            Error::new(
34                ErrorKind::Http,
35                format!("decoding esql response column {i} field `type`"),
36            )
37        })?;
38        out.push(EsqlColumn {
39            name: name.to_string(),
40            r#type: ty.to_string(),
41        });
42    }
43    Ok(out)
44}
45
46fn is_partial(obj: &Map<String, Value>) -> bool {
47    obj.get("is_partial")
48        .and_then(Value::as_bool)
49        .unwrap_or(false)
50}
51
52pub fn decode(value: &Value) -> Result<EsqlResponse> {
53    let obj = value.as_object().ok_or_else(|| {
54        Error::new(
55            ErrorKind::Http,
56            "decoding esql response: expected an object",
57        )
58    })?;
59    let columns = parse_columns(obj)?;
60
61    let values = obj
62        .get("values")
63        .and_then(Value::as_array)
64        .ok_or_else(|| Error::new(ErrorKind::Http, "decoding esql response field `values`"))?
65        .iter()
66        .map(|row| {
67            let cells = row.as_array().ok_or_else(|| {
68                Error::new(
69                    ErrorKind::Http,
70                    "decoding esql response: `values` rows must be arrays",
71                )
72            })?;
73            if cells.len() != columns.len() {
74                return Err(Error::new(
75                    ErrorKind::Http,
76                    "decoding esql response: `values` row width does not match `columns`",
77                ));
78            }
79            Ok(cells.clone())
80        })
81        .collect::<Result<Vec<_>>>()?;
82
83    Ok(EsqlResponse {
84        columns,
85        values,
86        is_partial: is_partial(obj),
87    })
88}
89
90pub fn decode_columnar(value: &Value) -> Result<EsqlResponse> {
91    let obj = value.as_object().ok_or_else(|| {
92        Error::new(
93            ErrorKind::Http,
94            "decoding esql response: expected an object",
95        )
96    })?;
97    let columns = parse_columns(obj)?;
98
99    let cols = obj
100        .get("values")
101        .and_then(Value::as_array)
102        .ok_or_else(|| Error::new(ErrorKind::Http, "decoding esql response field `values`"))?;
103    if cols.is_empty() {
104        return Ok(EsqlResponse {
105            columns,
106            values: Vec::new(),
107            is_partial: is_partial(obj),
108        });
109    }
110    if cols.len() != columns.len() {
111        return Err(Error::new(
112            ErrorKind::Http,
113            "decoding esql response: `values` column count does not match `columns`",
114        ));
115    }
116    let arrays = cols
117        .iter()
118        .enumerate()
119        .map(|(i, col)| {
120            col.as_array().ok_or_else(|| {
121                Error::new(
122                    ErrorKind::Http,
123                    format!("decoding esql response: `values` column {i} must be an array"),
124                )
125            })
126        })
127        .collect::<Result<Vec<_>>>()?;
128    let row_count = arrays.first().map(|col| col.len()).unwrap_or(0);
129    if arrays.iter().any(|col| col.len() != row_count) {
130        return Err(Error::new(
131            ErrorKind::Http,
132            "decoding esql response: `values` columns have unequal lengths",
133        ));
134    }
135
136    let mut values = Vec::with_capacity(row_count);
137    for r in 0..row_count {
138        let mut row = Vec::with_capacity(columns.len());
139        for col in &arrays {
140            row.push(col[r].clone());
141        }
142        values.push(row);
143    }
144
145    Ok(EsqlResponse {
146        columns,
147        values,
148        is_partial: is_partial(obj),
149    })
150}
151
152pub async fn run_sync(t: &Transport, query: &str) -> Result<EsqlResponse> {
153    let body = serde_json::json!({ "query": query });
154    let response = t.post_absolute_es("/_query", &body).await?;
155    decode(&response)
156}
157
158pub async fn run_async(t: &Transport, query: &str) -> Result<EsqlResponse> {
159    let start = t
160        .post_absolute_es(
161            "/_query/async",
162            &serde_json::json!({ "query": query, "wait_for_completion_timeout": "1ms", "columnar": true }),
163        )
164        .await?;
165    let id = match start.get("id").and_then(Value::as_str) {
166        Some(id) => id.to_string(),
167        None => return decode_columnar(&start),
168    };
169    if start.get("is_running").and_then(Value::as_bool) == Some(false) {
170        let _ = t.delete_absolute_es(&async_path(&id)?).await;
171        return decode_columnar(&start);
172    }
173
174    poll_until_complete(t, &id, 300, Duration::from_secs(1)).await
175}
176
177fn async_path(id: &str) -> Result<String> {
178    Ok(format!(
179        "/_query/async/{}",
180        urlencode(require_path_segment(id)?)
181    ))
182}
183
184pub async fn poll_until_complete(
185    t: &Transport,
186    id: &str,
187    max_polls: usize,
188    interval: Duration,
189) -> Result<EsqlResponse> {
190    let path = async_path(id)?;
191    for _ in 0..max_polls {
192        let resp = match t.get_absolute_es(&path).await {
193            Ok(resp) => resp,
194            Err(err) => {
195                let _ = t.delete_absolute_es(&path).await;
196                return Err(err);
197            }
198        };
199        if resp.get("is_running").and_then(Value::as_bool) == Some(false) {
200            let _ = t.delete_absolute_es(&path).await;
201            return decode_columnar(&resp);
202        }
203        tokio::time::sleep(interval).await;
204    }
205    let _ = t.delete_absolute_es(&path).await;
206    Err(Error::new(
207        ErrorKind::Timeout,
208        format!("async query {id} still running after {max_polls} polls"),
209    ))
210}
elasticctl_api/search/esql.rs

elasticctl_api/search/
esql.rs
