elasticctl_api/search/
esql.rs

1//! ES|QL responses and the sync/async `/_query` runners.
2
3use crate::compose::require_path_segment;
4use elasticctl_core::{Error, ErrorKind, Result, Transport, urlencode};
5use serde_json::{Map, Value};
6use std::time::Duration;
7
8#[derive(Debug, Clone, PartialEq)]
9pub struct EsqlColumn {
10    pub name: String,
11    pub r#type: String,
12}
13
14/// Row-major `values`: one array of cells per row. `columnar` responses are
15/// transposed into this shape by `decode_columnar` so callers always see rows.
16#[derive(Debug, Clone, PartialEq)]
17pub struct EsqlResponse {
18    pub columns: Vec<EsqlColumn>,
19    pub values: Vec<Vec<Value>>,
20    pub is_partial: bool,
21}
22
23fn parse_columns(obj: &Map<String, Value>) -> Result<Vec<EsqlColumn>> {
24    let columns = obj
25        .get("columns")
26        .and_then(Value::as_array)
27        .ok_or_else(|| Error::new(ErrorKind::Http, "decoding esql response field `columns`"))?;
28    let mut out = Vec::with_capacity(columns.len());
29    for (i, col) in columns.iter().enumerate() {
30        let name = col.get("name").and_then(Value::as_str).ok_or_else(|| {
31            Error::new(
32                ErrorKind::Http,
33                format!("decoding esql response column {i} field `name`"),
34            )
35        })?;
36        let ty = col.get("type").and_then(Value::as_str).ok_or_else(|| {
37            Error::new(
38                ErrorKind::Http,
39                format!("decoding esql response column {i} field `type`"),
40            )
41        })?;
42        out.push(EsqlColumn {
43            name: name.to_string(),
44            r#type: ty.to_string(),
45        });
46    }
47    Ok(out)
48}
49
50fn is_partial(obj: &Map<String, Value>) -> bool {
51    obj.get("is_partial")
52        .and_then(Value::as_bool)
53        .unwrap_or(false)
54}
55
56/// Strict decode of a row-major `POST /_query` response. Unknown fields are
57/// accepted; a missing or mistyped `columns`/`values` is an error, never an
58/// empty result.
59pub fn decode(value: &Value) -> Result<EsqlResponse> {
60    let obj = value.as_object().ok_or_else(|| {
61        Error::new(
62            ErrorKind::Http,
63            "decoding esql response: expected an object",
64        )
65    })?;
66    let columns = parse_columns(obj)?;
67
68    let values = obj
69        .get("values")
70        .and_then(Value::as_array)
71        .ok_or_else(|| Error::new(ErrorKind::Http, "decoding esql response field `values`"))?
72        .iter()
73        .map(|row| {
74            let cells = row.as_array().ok_or_else(|| {
75                Error::new(
76                    ErrorKind::Http,
77                    "decoding esql response: `values` rows must be arrays",
78                )
79            })?;
80            if cells.len() != columns.len() {
81                return Err(Error::new(
82                    ErrorKind::Http,
83                    "decoding esql response: `values` row width does not match `columns`",
84                ));
85            }
86            Ok(cells.clone())
87        })
88        .collect::<Result<Vec<_>>>()?;
89
90    Ok(EsqlResponse {
91        columns,
92        values,
93        is_partial: is_partial(obj),
94    })
95}
96
97/// Strict decode of a column-major (`columnar: true`) response. `values` holds
98/// one array per column, each the length of the row count; the result is
99/// transposed into the row-major `EsqlResponse` shape.
100pub fn decode_columnar(value: &Value) -> Result<EsqlResponse> {
101    let obj = value.as_object().ok_or_else(|| {
102        Error::new(
103            ErrorKind::Http,
104            "decoding esql response: expected an object",
105        )
106    })?;
107    let columns = parse_columns(obj)?;
108
109    let cols = obj
110        .get("values")
111        .and_then(Value::as_array)
112        .ok_or_else(|| Error::new(ErrorKind::Http, "decoding esql response field `values`"))?;
113    // A zero-row columnar result is `values: []`, not one empty array per
114    // column (measured against Serverless 9.6.0). Treat it as an empty result
115    // rather than a column-count mismatch.
116    if cols.is_empty() {
117        return Ok(EsqlResponse {
118            columns,
119            values: Vec::new(),
120            is_partial: is_partial(obj),
121        });
122    }
123    if cols.len() != columns.len() {
124        return Err(Error::new(
125            ErrorKind::Http,
126            "decoding esql response: `values` column count does not match `columns`",
127        ));
128    }
129    let arrays = cols
130        .iter()
131        .enumerate()
132        .map(|(i, col)| {
133            col.as_array().ok_or_else(|| {
134                Error::new(
135                    ErrorKind::Http,
136                    format!("decoding esql response: `values` column {i} must be an array"),
137                )
138            })
139        })
140        .collect::<Result<Vec<_>>>()?;
141    let row_count = arrays.first().map(|col| col.len()).unwrap_or(0);
142    if arrays.iter().any(|col| col.len() != row_count) {
143        return Err(Error::new(
144            ErrorKind::Http,
145            "decoding esql response: `values` columns have unequal lengths",
146        ));
147    }
148
149    let mut values = Vec::with_capacity(row_count);
150    for r in 0..row_count {
151        let mut row = Vec::with_capacity(columns.len());
152        for col in &arrays {
153            row.push(col[r].clone());
154        }
155        values.push(row);
156    }
157
158    Ok(EsqlResponse {
159        columns,
160        values,
161        is_partial: is_partial(obj),
162    })
163}
164
165/// Run a synchronous ES|QL query. `query` carries its own `FROM` and `LIMIT`.
166pub async fn run_sync(t: &Transport, query: &str) -> Result<EsqlResponse> {
167    let body = serde_json::json!({ "query": query });
168    let response = t.post_absolute_es("/_query", &body).await?;
169    decode(&response)
170}
171
172/// Run a query through the async API and poll until complete. ES|QL has no
173/// page-by-page cursor; the full result returns in one response. `columnar:
174/// true` keeps the payload and memory footprint low.
175pub async fn run_async(t: &Transport, query: &str) -> Result<EsqlResponse> {
176    let start = t
177        .post_absolute_es(
178            "/_query/async",
179            &serde_json::json!({ "query": query, "wait_for_completion_timeout": "1ms", "columnar": true }),
180        )
181        .await?;
182    // A query finishing within wait_for_completion_timeout returns the inline
183    // result with is_running: false and no `id`; decode it directly.
184    let id = match start.get("id").and_then(Value::as_str) {
185        Some(id) => id.to_string(),
186        None => return decode_columnar(&start),
187    };
188    // The start response can also carry both an `id` and an inline result.
189    // Clean the id up before decoding either way.
190    if start.get("is_running").and_then(Value::as_bool) == Some(false) {
191        let _ = t.delete_absolute_es(&async_path(&id)?).await;
192        return decode_columnar(&start);
193    }
194
195    // Poll once a second for up to five minutes: a bulk export may legitimately
196    // run that long, but a query that never finishes must still fail closed.
197    poll_until_complete(t, &id, 300, Duration::from_secs(1)).await
198}
199
200/// The async query route for an id the server returned. Encoding keeps the
201/// id one path segment, and `.` or `..` is refused.
202fn async_path(id: &str) -> Result<String> {
203    Ok(format!(
204        "/_query/async/{}",
205        urlencode(require_path_segment(id)?)
206    ))
207}
208
209/// Poll `GET /_query/async/{id}` until it reports completion, cleaning up on
210/// every exit path. `max_polls * interval` bounds the wait; `run_async` passes
211/// a bulk-export-sized budget, while tests pass a tiny one.
212pub async fn poll_until_complete(
213    t: &Transport,
214    id: &str,
215    max_polls: usize,
216    interval: Duration,
217) -> Result<EsqlResponse> {
218    let path = async_path(id)?;
219    for _ in 0..max_polls {
220        let resp = match t.get_absolute_es(&path).await {
221            Ok(resp) => resp,
222            Err(err) => {
223                let _ = t.delete_absolute_es(&path).await;
224                return Err(err);
225            }
226        };
227        if resp.get("is_running").and_then(Value::as_bool) == Some(false) {
228            let _ = t.delete_absolute_es(&path).await;
229            return decode_columnar(&resp);
230        }
231        tokio::time::sleep(interval).await;
232    }
233    let _ = t.delete_absolute_es(&path).await;
234    Err(Error::new(
235        ErrorKind::Timeout,
236        format!("async query {id} still running after {max_polls} polls"),
237    ))
238}
elasticctl_api/search/esql.rs

elasticctl_api/search/
esql.rs
