Skip to main content

rudb_cli/
shell.rs

1//! The shell itself: read a line, decide whether it is SQL or a dot command, run it, print it.
2
3use std::fmt::Write as _;
4use std::fs::File;
5use std::io::{self, BufWriter, Write};
6use std::path::{Path, PathBuf};
7use std::time::Instant;
8
9use rudb::{Database, QueryResult};
10use rudb_common::{Error, Span};
11
12use crate::args::{Command, Options};
13use crate::format::{Format, Settings, escaped, render};
14
15/// Where printed results go.
16///
17/// `.output FILE` and `.output` back again is the reason this is a type rather than a
18/// `Box<dyn Write>` handed in once. A shell that can only write to the stream it was started with
19/// cannot be used to produce a file, which is most of what the CSV and JSON modes are for.
20enum Sink {
21    /// The stream the shell was started with.
22    Given(Box<dyn Write>),
23    /// A file opened by `.output`.
24    File(BufWriter<File>, PathBuf),
25}
26
27impl Write for Sink {
28    fn write(&mut self, buffer: &[u8]) -> io::Result<usize> {
29        match self {
30            Self::Given(out) => out.write(buffer),
31            Self::File(out, _) => out.write(buffer),
32        }
33    }
34
35    fn flush(&mut self) -> io::Result<()> {
36        match self {
37            Self::Given(out) => out.flush(),
38            Self::File(out, _) => out.flush(),
39        }
40    }
41}
42
43/// Why the shell stopped.
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub enum Stop {
46    /// `.quit`, `.exit`, or the end of the input.
47    Done,
48    /// An error, and `-bail` was on.
49    Failed,
50}
51
52/// A running shell.
53pub struct Shell {
54    database: Database,
55    settings: Settings,
56    out: Sink,
57    err: Box<dyn Write>,
58    filename: String,
59    given: Option<Box<dyn Write>>,
60    timer: bool,
61    echo: bool,
62    bail: bool,
63    failed: bool,
64}
65
66impl std::fmt::Debug for Shell {
67    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
68        f.debug_struct("Shell")
69            .field("settings", &self.settings)
70            .field("failed", &self.failed)
71            .finish_non_exhaustive()
72    }
73}
74
75impl Shell {
76    /// A shell over a fresh database, writing results to `out` and errors to `err`.
77    pub fn new(options: &Options, out: Box<dyn Write>, err: Box<dyn Write>) -> Self {
78        Self {
79            database: Database::new(),
80            settings: options.settings.clone(),
81            out: Sink::Given(out),
82            err,
83            filename: options.database.clone(),
84            given: None,
85            timer: false,
86            echo: options.echo,
87            bail: options.bail,
88            failed: false,
89        }
90    }
91
92    /// Whether anything has failed since the shell started, which is what the exit code is.
93    pub fn failed(&self) -> bool {
94        self.failed
95    }
96
97    /// Runs everything the command line asked for, in order.
98    pub fn run_commands(&mut self, commands: &[Command]) -> Stop {
99        for command in commands {
100            let stop = match command {
101                Command::Sql(sql) => self.run_input(sql),
102                Command::File(path) => self.run_file(path),
103            };
104            if stop == Stop::Failed {
105                return Stop::Failed;
106            }
107        }
108        Stop::Done
109    }
110
111    /// Runs a file of SQL and dot commands.
112    pub fn run_file(&mut self, path: &Path) -> Stop {
113        match std::fs::read_to_string(path) {
114            Ok(text) => self.run_input(&text),
115            Err(problem) => {
116                let why = format!("Cannot open file \"{}\": {problem}", path.display());
117                self.report(&Error::io(why), "");
118                self.after_error()
119            }
120        }
121    }
122
123    /// The two lines a terminal gets before the first prompt.
124    pub fn greet(&mut self) {
125        let _ = writeln!(self.out, "rudb {}", crate::VERSION);
126        let _ = writeln!(self.out, "Enter \".help\" for usage hints.");
127        let _ = self.out.flush();
128    }
129
130    /// Reads and runs lines from a terminal until the user stops.
131    ///
132    /// The continuation marker is what says the statement is not finished, and it is the reason
133    /// [`rudb::is_complete`] exists rather than the shell guessing from a trailing semicolon.
134    pub fn prompt(&mut self, stdin: &io::Stdin) -> Stop {
135        let mut pending = String::new();
136        loop {
137            let marker = if pending.is_empty() { "D " } else { "ยท " };
138            let _ = write!(self.out, "{marker}");
139            let _ = self.out.flush();
140            let mut line = String::new();
141            match stdin.read_line(&mut line) {
142                Ok(0) => {
143                    let _ = writeln!(self.out);
144                    return Stop::Done;
145                }
146                Ok(_) => {}
147                Err(_) => return Stop::Done,
148            }
149            let line = line.trim_end_matches(['\n', '\r']);
150            if pending.is_empty() && line.trim_start().starts_with('.') {
151                if self.run_dot(line.trim()) == Stop::Failed {
152                    return Stop::Done;
153                }
154                continue;
155            }
156            if !pending.is_empty() {
157                pending.push('\n');
158            }
159            pending.push_str(line);
160            if rudb::is_complete(&pending) {
161                let statement = std::mem::take(&mut pending);
162                if self.run_sql(&statement) == Stop::Failed {
163                    return Stop::Done;
164                }
165            }
166        }
167    }
168
169    /// Runs a block of input, which may hold any mixture of dot commands and statements.
170    ///
171    /// Line oriented rather than statement oriented, because a dot command is a line and SQL is
172    /// not. Lines accumulate into a statement until the tokenizer says the statement is finished,
173    /// which is how a multi line `CREATE TABLE` works at a prompt and in a file alike.
174    pub fn run_input(&mut self, text: &str) -> Stop {
175        let mut pending = String::new();
176        for line in text.lines() {
177            if pending.trim().is_empty() && line.trim_start().starts_with('.') {
178                pending.clear();
179                if self.run_dot(line.trim()) == Stop::Failed {
180                    return Stop::Failed;
181                }
182                continue;
183            }
184            if !pending.is_empty() {
185                pending.push('\n');
186            }
187            pending.push_str(line);
188            if rudb::is_complete(&pending) {
189                let statement = std::mem::take(&mut pending);
190                if self.run_sql(&statement) == Stop::Failed {
191                    return Stop::Failed;
192                }
193            }
194        }
195        if pending.trim().is_empty() {
196            return Stop::Done;
197        }
198        self.run_sql(&pending)
199    }
200
201    /// Runs whatever statements are in one piece of text.
202    fn run_sql(&mut self, text: &str) -> Stop {
203        let found = match rudb::statements(text) {
204            Ok(found) => found,
205            Err(problem) => {
206                self.report(&problem, text);
207                return self.after_error();
208            }
209        };
210        for statement in found {
211            if self.echo {
212                let _ = writeln!(self.out, "{}", statement.sql());
213            }
214            let started = Instant::now();
215            match self.database.execute(statement.sql()) {
216                Ok(result) => {
217                    self.print(&result);
218                    if self.timer {
219                        let _ = writeln!(
220                            self.err,
221                            "Run Time (s): real {:.3}",
222                            started.elapsed().as_secs_f64()
223                        );
224                    }
225                }
226                Err(problem) => {
227                    self.report(&problem, statement.sql());
228                    return self.after_error();
229                }
230            }
231        }
232        Stop::Done
233    }
234
235    /// Prints a result, unless it is the empty one a writing statement hands back.
236    fn print(&mut self, result: &QueryResult) {
237        let text = render(result, &self.settings);
238        if !text.is_empty() {
239            let _ = write!(self.out, "{text}");
240            let _ = self.out.flush();
241        }
242    }
243
244    /// What an error does to the run, which depends on `-bail`.
245    fn after_error(&mut self) -> Stop {
246        self.failed = true;
247        if self.bail { Stop::Failed } else { Stop::Done }
248    }
249
250    /// Prints an error the way DuckDB prints one: the message, then the line it is about with a
251    /// caret under the offending token.
252    fn report(&mut self, problem: &Error, sql: &str) {
253        let _ = writeln!(self.err, "{problem}");
254        if let Some(span) = problem.span() {
255            if let Some(text) = pointer(sql, span) {
256                let _ = writeln!(self.err);
257                let _ = write!(self.err, "{text}");
258            }
259        }
260        let _ = self.err.flush();
261    }
262
263    /// Runs one dot command.
264    fn run_dot(&mut self, line: &str) -> Stop {
265        let mut words = split(line);
266        if words.is_empty() {
267            return Stop::Done;
268        }
269        let name = words.remove(0);
270        let argument = |at: usize| words.get(at).cloned().unwrap_or_default();
271        match name.as_str() {
272            ".quit" | ".exit" => return Stop::Failed,
273            ".help" => {
274                let _ = write!(self.out, "{}", crate::help::DOT_COMMANDS);
275            }
276            ".mode" => {
277                if words.is_empty() {
278                    let _ =
279                        writeln!(self.out, "current output mode: {}", self.settings.format.name());
280                } else if let Some(format) = Format::from_name(&argument(0)) {
281                    self.settings.set_format(format);
282                    if let Some(table) = words.get(1) {
283                        self.settings.table = table.clone();
284                    }
285                } else {
286                    return self.complain(&format!(
287                        "Error: mode should be one of: {}",
288                        crate::help::MODES
289                    ));
290                }
291            }
292            ".headers" | ".header" => self.settings.header = on(&argument(0)),
293            ".separator" => {
294                self.settings.separator = argument(0);
295                if let Some(newline) = words.get(1) {
296                    self.settings.newline = newline.clone();
297                }
298            }
299            ".nullvalue" | ".nullValue" => self.settings.nullvalue = argument(0),
300            ".timer" => self.timer = on(&argument(0)),
301            ".echo" => self.echo = on(&argument(0)),
302            ".bail" => self.bail = on(&argument(0)),
303            ".print" => {
304                let _ = writeln!(self.out, "{}", words.join(" "));
305            }
306            ".read" => return self.run_file(Path::new(&argument(0))),
307            ".output" => return self.redirect(words.first().map(String::as_str)),
308            ".tables" => self.tables(words.first().map(String::as_str)),
309            ".schema" => self.schema(words.first().map(String::as_str)),
310            ".databases" => {
311                let _ = writeln!(self.out, "memory:");
312            }
313            ".show" => self.show(),
314            ".open" => {
315                return self.complain(
316                    "Error: .open needs a storage format, which is not built yet. See https://github.com/tamnd/rudb/issues/103",
317                );
318            }
319            other => {
320                return self
321                    .complain(&format!("Error: unknown command or invalid arguments:  \"{}\". Enter \".help\" for help", other.trim_start_matches('.')));
322            }
323        }
324        let _ = self.out.flush();
325        Stop::Done
326    }
327
328    /// Prints a complaint about a dot command, which is an error like any other.
329    fn complain(&mut self, message: &str) -> Stop {
330        let _ = writeln!(self.err, "{message}");
331        let _ = self.err.flush();
332        self.after_error()
333    }
334
335    /// `.output`, both directions.
336    fn redirect(&mut self, path: Option<&str>) -> Stop {
337        let _ = self.out.flush();
338        match path {
339            None | Some("stdout") => {
340                if let Some(given) = self.given.take() {
341                    self.out = Sink::Given(given);
342                }
343            }
344            Some(path) => {
345                let path = PathBuf::from(path);
346                match File::create(&path) {
347                    Ok(file) => {
348                        let opened = Sink::File(BufWriter::new(file), path);
349                        if let Sink::Given(given) = std::mem::replace(&mut self.out, opened) {
350                            self.given = Some(given);
351                        }
352                    }
353                    Err(problem) => {
354                        return self.complain(&format!(
355                            "Error: cannot open \"{}\": {problem}",
356                            path.display()
357                        ));
358                    }
359                }
360            }
361        }
362        Stop::Done
363    }
364
365    /// `.tables`, one name per line.
366    fn tables(&mut self, pattern: Option<&str>) {
367        let mut names = self.database.table_names();
368        names.sort();
369        for name in names {
370            if pattern.is_none_or(|pattern| matches(&name, pattern)) {
371                let _ = writeln!(self.out, "{name}");
372            }
373        }
374    }
375
376    /// `.schema`, the `CREATE TABLE` for every table or for one of them.
377    fn schema(&mut self, wanted: Option<&str>) {
378        let mut names = self.database.table_names();
379        names.sort();
380        for name in names {
381            if wanted.is_some_and(|wanted| !matches(&name, wanted)) {
382                continue;
383            }
384            if let Ok(sql) = self.database.table_sql(&name) {
385                let _ = writeln!(self.out, "{sql}");
386            }
387        }
388    }
389
390    /// `.show`, in the order and the spacing DuckDB prints it.
391    ///
392    /// `width` is blank because there is no column width setting yet, and it is listed anyway so
393    /// that a script reading this output finds the line where it expects it.
394    fn show(&mut self) {
395        let mut out = String::new();
396        let _ = writeln!(out, "        echo: {}", off_on(self.echo));
397        let _ = writeln!(out, "     headers: {}", off_on(self.settings.header));
398        let _ = writeln!(out, "        mode: {}", self.settings.format.name());
399        let _ = writeln!(out, "   nullvalue: \"{}\"", self.settings.nullvalue);
400        let _ = writeln!(out, "      output: {}", self.output_name());
401        let _ = writeln!(out, "colseparator: \"{}\"", escaped(&self.settings.separator));
402        let _ = writeln!(out, "rowseparator: \"{}\"", escaped(&self.settings.newline));
403        let _ = writeln!(out, "       width: ");
404        let _ = writeln!(out, "    filename: {}", self.filename);
405        let _ = write!(self.out, "{out}");
406    }
407
408    /// What `.show` calls the place output is going.
409    fn output_name(&self) -> String {
410        match &self.out {
411            Sink::Given(_) => "stdout".to_string(),
412            Sink::File(_, path) => path.display().to_string(),
413        }
414    }
415}
416
417/// Whether a name matches a `.tables` or `.schema` pattern, where `%` stands for any run.
418fn matches(name: &str, pattern: &str) -> bool {
419    let pattern = pattern.trim_matches('\'');
420    if let Some(prefix) = pattern.strip_suffix('%') {
421        name.starts_with(prefix)
422    } else {
423        name.eq_ignore_ascii_case(pattern)
424    }
425}
426
427/// How a dot command's arguments are split: on whitespace, with quoted runs kept together.
428fn split(line: &str) -> Vec<String> {
429    let mut words = Vec::new();
430    let mut current = String::new();
431    let mut quote = None;
432    let mut started = false;
433    for character in line.chars() {
434        match quote {
435            Some(open) if character == open => quote = None,
436            Some(_) => current.push(character),
437            None if character == '\'' || character == '"' => {
438                quote = Some(character);
439                started = true;
440            }
441            None if character.is_whitespace() => {
442                if started || !current.is_empty() {
443                    words.push(std::mem::take(&mut current));
444                    started = false;
445                }
446            }
447            None => current.push(character),
448        }
449    }
450    if started || !current.is_empty() {
451        words.push(current);
452    }
453    words
454}
455
456/// How a dot command spells a boolean, where anything that is not a recognized "off" is "on".
457fn on(word: &str) -> bool {
458    !matches!(word, "off" | "0" | "false" | "no")
459}
460
461/// How `.show` spells one back.
462fn off_on(flag: bool) -> &'static str {
463    if flag { "on" } else { "off" }
464}
465
466/// The `LINE n:` and caret that go under an error message.
467///
468/// `None` when the span does not point into the text, which happens for an error raised about a
469/// statement the caller did not hand us, and printing a caret under the wrong thing is worse than
470/// printing none.
471fn pointer(sql: &str, span: Span) -> Option<String> {
472    let start = span.start as usize;
473    if start > sql.len() || !sql.is_char_boundary(start) {
474        return None;
475    }
476    let before = &sql[..start];
477    let number = before.matches('\n').count() + 1;
478    let line_start = before.rfind('\n').map_or(0, |at| at + 1);
479    let line_end = sql[line_start..].find('\n').map_or(sql.len(), |at| line_start + at);
480    let line = &sql[line_start..line_end];
481    let prefix = format!("LINE {number}: ");
482    let column = sql[line_start..start].chars().count();
483    Some(format!("{prefix}{line}\n{}^\n", " ".repeat(prefix.chars().count() + column)))
484}
485
486#[cfg(test)]
487mod tests {
488    use super::{matches, on, pointer, split};
489    use rudb_common::Span;
490
491    #[test]
492    fn a_dot_command_splits_on_whitespace() {
493        assert_eq!(split(".mode csv"), vec![".mode", "csv"]);
494        assert_eq!(split("  .timer   on  "), vec![".timer", "on"]);
495    }
496
497    #[test]
498    fn a_quoted_argument_keeps_its_spaces() {
499        assert_eq!(split(".separator ' | '"), vec![".separator", " | "]);
500        assert_eq!(split(".nullvalue \"\""), vec![".nullvalue", ""]);
501    }
502
503    #[test]
504    fn off_is_the_only_way_to_turn_something_off() {
505        assert!(on("on"));
506        assert!(on(""));
507        assert!(!on("off"));
508        assert!(!on("0"));
509    }
510
511    #[test]
512    fn a_pattern_ending_in_a_percent_is_a_prefix() {
513        assert!(matches("orders", "orders"));
514        assert!(matches("orders", "ORDERS"));
515        assert!(matches("orders", "ord%"));
516        assert!(!matches("orders", "lineitem"));
517    }
518
519    #[test]
520    fn the_caret_lands_under_the_span() {
521        let sql = "SELECT nosuch";
522        let text = pointer(sql, Span::new(7, 13)).expect("a pointer");
523        assert_eq!(text, "LINE 1: SELECT nosuch\n               ^\n");
524    }
525
526    #[test]
527    fn the_caret_counts_lines() {
528        let sql = "SELECT\n  nosuch";
529        let text = pointer(sql, Span::new(9, 15)).expect("a pointer");
530        assert_eq!(text, "LINE 2:   nosuch\n          ^\n");
531    }
532
533    #[test]
534    fn a_span_past_the_end_gets_no_pointer() {
535        assert!(pointer("SELECT 1", Span::new(100, 101)).is_none());
536    }
537}