Skip to main content

haematite_cli/client/
query.rs

1//! `haematite query`: a query's rows, every page of them, one per line.
2//!
3//! Each page is sized from the service's hello, as the Studio sizes it, and
4//! every page is followed until `next` is null, at the moment the first page
5//! was read. An answer that stops part-way says so, with how many rows were
6//! already printed, so a complete answer is never mistaken for a cut one.
7
8use std::io::{self, Read, Write};
9use std::process::ExitCode;
10use std::time::Duration;
11
12use serde::{Deserialize, Serialize};
13use serde_json::value::RawValue;
14
15use super::{Client, Common, Ended, Refusal, Reply, Target, exit_code, socket_of};
16
17const USAGE: &str = "\
18usage: haematite query (--socket <path> | --config <serve.json>) [--timeout-ms N] --branch <name>
19                  <text | ->
20<text> is the query as typed: jq, or a pipeline; - reads it from standard input. It runs at
21the branch's head, and every page is read at the moment the first page was read.
22Each row is printed on its own line as the service sent it, value beside origin.
23--timeout-ms bounds connecting, writing and reading; the default is 30000.
24Exit 0: every row printed. Exit 1: could not ask; nothing was sent. Exit 2: refused by the
25service. Exit 3: the outcome is not known. On 2 or 3 after rows were printed, standard error
26says how many, and that the answer is not complete.";
27
28struct Query {
29    target: Target,
30    timeout: Duration,
31    branch: String,
32    text: Option<String>,
33}
34
35#[derive(Serialize)]
36struct Run<'a> {
37    selection: Selection<'a>,
38    query: &'a str,
39    page: PageSize,
40    #[serde(skip_serializing_if = "Option::is_none")]
41    continuation: Option<&'a str>,
42}
43
44#[derive(Serialize)]
45struct Selection<'a> {
46    branch: &'a str,
47    #[serde(skip_serializing_if = "Option::is_none")]
48    version: Option<&'a RawValue>,
49}
50
51#[derive(Clone, Copy, Serialize)]
52struct PageSize {
53    rows: u32,
54    bytes: u32,
55}
56
57#[derive(Deserialize)]
58struct Page {
59    selection: Moment,
60    rows: Vec<Box<RawValue>>,
61    /// Null or text, never absent: an answer without `next` lost its ending
62    /// and is not read as complete.
63    #[serde(deserialize_with = "Option::deserialize")]
64    next: Option<String>,
65}
66
67#[derive(Deserialize)]
68struct Moment {
69    version: Box<RawValue>,
70}
71
72/// How a query stopped short of its last page.
73enum Stop {
74    Refused(Refusal),
75    Ended(Ended),
76}
77
78/// `haematite query`, ending in its exit code.
79#[must_use]
80pub fn run(args: impl Iterator<Item = String>) -> ExitCode {
81    let mut printed = 0_u64;
82    let outcome = {
83        let mut output = io::stdout().lock();
84        query(args, &mut output, &mut printed)
85    };
86    let code = match outcome {
87        Ok(()) => return ExitCode::SUCCESS,
88        Err(Stop::Refused(refusal)) => {
89            if let Err(error) = super::print_refusal(&refusal) {
90                eprintln!("haematite: receipt_output_failed: {error}");
91            }
92            2
93        }
94        Err(Stop::Ended(ended)) => {
95            if let Err(error) = super::print_ended(&ended) {
96                eprintln!("haematite: receipt_output_failed: {error}");
97            }
98            exit_code(&ended)
99        }
100    };
101    if printed > 0 {
102        eprintln!(
103            "haematite: stopped part-way: {printed} rows were printed before this; the answer is not \
104             complete"
105        );
106    }
107    ExitCode::from(code)
108}
109
110fn query(
111    args: impl Iterator<Item = String>,
112    output: &mut impl Write,
113    printed: &mut u64,
114) -> Result<(), Stop> {
115    let asked = parse(args).map_err(|why| Stop::Ended(Ended::Usage(format!("{why}\n{USAGE}"))))?;
116    let text = text_of(asked.text).map_err(Stop::Ended)?;
117    let socket = socket_of(&asked.target).map_err(Stop::Ended)?;
118    let mut client = Client::connect(&socket, asked.timeout).map_err(Stop::Ended)?;
119    let page = PageSize {
120        rows: client.max_page().get(),
121        bytes: client.max_frame().get(),
122    };
123    let unknown = |why: String| Stop::Ended(Ended::Unknown(why));
124    let mut moment: Option<Box<RawValue>> = None;
125    let mut continuation: Option<String> = None;
126    loop {
127        let request = Run {
128            selection: Selection {
129                branch: &asked.branch,
130                version: moment.as_deref(),
131            },
132            query: &text,
133            page,
134            continuation: continuation.as_deref(),
135        };
136        let params = serde_json::value::to_raw_value(&request).map_err(|error| {
137            Stop::Ended(Ended::NotAsked(format!(
138                "the request could not be written: {error}"
139            )))
140        })?;
141        let answer = match client
142            .call("query.run", Some(&params))
143            .map_err(Stop::Ended)?
144        {
145            Reply::Answered(answer) => answer,
146            Reply::Refused(refusal) => return Err(Stop::Refused(refusal)),
147        };
148        let read: Page = serde_json::from_str(answer.get())
149            .map_err(|error| unknown(format!("the answer could not be read: {error}")))?;
150        if read.next.is_some() && read.next == continuation {
151            return Err(unknown(format!(
152                "the service repeated the continuation {:?}, so the answer would never end",
153                continuation.unwrap_or_default()
154            )));
155        }
156        for row in &read.rows {
157            writeln!(output, "{}", row.get())
158                .map_err(|error| unknown(format!("a row could not be written: {error}")))?;
159            *printed += 1;
160        }
161        let Some(next) = read.next else {
162            break;
163        };
164        moment = Some(read.selection.version);
165        continuation = Some(next);
166    }
167    output
168        .flush()
169        .map_err(|error| unknown(format!("the rows could not be written: {error}")))
170}
171
172fn parse(mut args: impl Iterator<Item = String>) -> Result<Query, String> {
173    let mut common = Common::default();
174    let mut branch = None;
175    let mut rest = Vec::new();
176    while let Some(arg) = args.next() {
177        if common.take(&arg, &mut args)? {
178            continue;
179        }
180        match arg.as_str() {
181            "--branch" if branch.is_none() => {
182                branch = Some(args.next().ok_or("--branch needs a name")?);
183            }
184            flag if flag.starts_with("--") => {
185                return Err(format!("unknown or repeated flag {flag:?}"));
186            }
187            _ => rest.push(arg),
188        }
189    }
190    let (target, timeout) = common.finish()?;
191    let branch = branch.ok_or("name the branch with --branch")?;
192    let mut rest = rest.into_iter();
193    let text = rest.next().ok_or("give the query text, or -")?;
194    if let Some(extra) = rest.next() {
195        return Err(format!("unexpected argument {extra:?}"));
196    }
197    Ok(Query {
198        target,
199        timeout,
200        branch,
201        text: (text != "-").then_some(text),
202    })
203}
204
205/// The query text as given, or read from standard input when `None`.
206fn text_of(given: Option<String>) -> Result<String, Ended> {
207    given.map_or_else(
208        || {
209            let mut text = String::new();
210            io::stdin().read_to_string(&mut text).map_err(|error| {
211                Ended::NotAsked(format!(
212                    "the query could not be read from standard input: {error}"
213                ))
214            })?;
215            Ok(text)
216        },
217        Ok,
218    )
219}