haematite_cli/client/

query.rs

//! `haematite query`: a query's rows, every page of them, one per line.
//!
//! Each page is sized from the service's hello, as the Studio sizes it, and
//! every page is followed until `next` is null, at the moment the first page
//! was read. An answer that stops part-way says so, with how many rows were
//! already printed, so a complete answer is never mistaken for a cut one.

use std::io::{self, Read, Write};
use std::process::ExitCode;
use std::time::Duration;

use serde::{Deserialize, Serialize};
use serde_json::value::RawValue;

use super::{Client, Common, Ended, Refusal, Reply, Target, exit_code, socket_of};

const USAGE: &str = "\
usage: haematite query (--socket <path> | --config <serve.json>) [--timeout-ms N] --branch <name>
                  <text | ->
<text> is the query as typed: jq, or a pipeline; - reads it from standard input. It runs at
the branch's head, and every page is read at the moment the first page was read.
Each row is printed on its own line as the service sent it, value beside origin.
--timeout-ms bounds connecting, writing and reading; the default is 30000.
Exit 0: every row printed. Exit 1: could not ask; nothing was sent. Exit 2: refused by the
service. Exit 3: the outcome is not known. On 2 or 3 after rows were printed, standard error
says how many, and that the answer is not complete.";

struct Query {
    target: Target,
    timeout: Duration,
    branch: String,
    text: Option<String>,
}

#[derive(Serialize)]
struct Run<'a> {
    selection: Selection<'a>,
    query: &'a str,
    page: PageSize,
    #[serde(skip_serializing_if = "Option::is_none")]
    continuation: Option<&'a str>,
}

#[derive(Serialize)]
struct Selection<'a> {
    branch: &'a str,
    #[serde(skip_serializing_if = "Option::is_none")]
    version: Option<&'a RawValue>,
}

#[derive(Clone, Copy, Serialize)]
struct PageSize {
    rows: u32,
    bytes: u32,
}

#[derive(Deserialize)]
struct Page {
    selection: Moment,
    rows: Vec<Box<RawValue>>,
    /// Null or text, never absent: an answer without `next` lost its ending
    /// and is not read as complete.
    #[serde(deserialize_with = "Option::deserialize")]
    next: Option<String>,
}

#[derive(Deserialize)]
struct Moment {
    version: Box<RawValue>,
}

/// How a query stopped short of its last page.
enum Stop {
    Refused(Refusal),
    Ended(Ended),
}

/// `haematite query`, ending in its exit code.
#[must_use]
pub fn run(args: impl Iterator<Item = String>) -> ExitCode {
    let mut printed = 0_u64;
    let outcome = {
        let mut output = io::stdout().lock();
        query(args, &mut output, &mut printed)
    };
    let code = match outcome {
        Ok(()) => return ExitCode::SUCCESS,
        Err(Stop::Refused(refusal)) => {
            if let Err(error) = super::print_refusal(&refusal) {
                eprintln!("haematite: receipt_output_failed: {error}");
            }
            2
        }
        Err(Stop::Ended(ended)) => {
            if let Err(error) = super::print_ended(&ended) {
                eprintln!("haematite: receipt_output_failed: {error}");
            }
            exit_code(&ended)
        }
    };
    if printed > 0 {
        eprintln!(
            "haematite: stopped part-way: {printed} rows were printed before this; the answer is not \
             complete"
        );
    }
    ExitCode::from(code)
}

fn query(
    args: impl Iterator<Item = String>,
    output: &mut impl Write,
    printed: &mut u64,
) -> Result<(), Stop> {
    let asked = parse(args).map_err(|why| Stop::Ended(Ended::Usage(format!("{why}\n{USAGE}"))))?;
    let text = text_of(asked.text).map_err(Stop::Ended)?;
    let socket = socket_of(&asked.target).map_err(Stop::Ended)?;
    let mut client = Client::connect(&socket, asked.timeout).map_err(Stop::Ended)?;
    let page = PageSize {
        rows: client.max_page().get(),
        bytes: client.max_frame().get(),
    };
    let unknown = |why: String| Stop::Ended(Ended::Unknown(why));
    let mut moment: Option<Box<RawValue>> = None;
    let mut continuation: Option<String> = None;
    loop {
        let request = Run {
            selection: Selection {
                branch: &asked.branch,
                version: moment.as_deref(),
            },
            query: &text,
            page,
            continuation: continuation.as_deref(),
        };
        let params = serde_json::value::to_raw_value(&request).map_err(|error| {
            Stop::Ended(Ended::NotAsked(format!(
                "the request could not be written: {error}"
            )))
        })?;
        let answer = match client
            .call("query.run", Some(&params))
            .map_err(Stop::Ended)?
        {
            Reply::Answered(answer) => answer,
            Reply::Refused(refusal) => return Err(Stop::Refused(refusal)),
        };
        let read: Page = serde_json::from_str(answer.get())
            .map_err(|error| unknown(format!("the answer could not be read: {error}")))?;
        if read.next.is_some() && read.next == continuation {
            return Err(unknown(format!(
                "the service repeated the continuation {:?}, so the answer would never end",
                continuation.unwrap_or_default()
            )));
        }
        for row in &read.rows {
            writeln!(output, "{}", row.get())
                .map_err(|error| unknown(format!("a row could not be written: {error}")))?;
            *printed += 1;
        }
        let Some(next) = read.next else {
            break;
        };
        moment = Some(read.selection.version);
        continuation = Some(next);
    }
    output
        .flush()
        .map_err(|error| unknown(format!("the rows could not be written: {error}")))
}

fn parse(mut args: impl Iterator<Item = String>) -> Result<Query, String> {
    let mut common = Common::default();
    let mut branch = None;
    let mut rest = Vec::new();
    while let Some(arg) = args.next() {
        if common.take(&arg, &mut args)? {
            continue;
        }
        match arg.as_str() {
            "--branch" if branch.is_none() => {
                branch = Some(args.next().ok_or("--branch needs a name")?);
            }
            flag if flag.starts_with("--") => {
                return Err(format!("unknown or repeated flag {flag:?}"));
            }
            _ => rest.push(arg),
        }
    }
    let (target, timeout) = common.finish()?;
    let branch = branch.ok_or("name the branch with --branch")?;
    let mut rest = rest.into_iter();
    let text = rest.next().ok_or("give the query text, or -")?;
    if let Some(extra) = rest.next() {
        return Err(format!("unexpected argument {extra:?}"));
    }
    Ok(Query {
        target,
        timeout,
        branch,
        text: (text != "-").then_some(text),
    })
}

/// The query text as given, or read from standard input when `None`.
fn text_of(given: Option<String>) -> Result<String, Ended> {
    given.map_or_else(
        || {
            let mut text = String::new();
            io::stdin().read_to_string(&mut text).map_err(|error| {
                Ended::NotAsked(format!(
                    "the query could not be read from standard input: {error}"
                ))
            })?;
            Ok(text)
        },
        Ok,
    )
}