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