frostbow_cli/
command.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Command within CLI
19
20use crate::cli_context::CliSessionContext;
21use crate::exec::{exec_and_print, exec_from_lines};
22use crate::functions::{display_all_functions, Function};
23use crate::print_format::PrintFormat;
24use crate::print_options::PrintOptions;
25use clap::ValueEnum;
26use datafusion::arrow::array::{ArrayRef, StringArray};
27use datafusion::arrow::datatypes::{DataType, Field, Schema};
28use datafusion::arrow::record_batch::RecordBatch;
29use datafusion::common::exec_err;
30use datafusion::common::instant::Instant;
31use datafusion::error::{DataFusionError, Result};
32use std::fs::File;
33use std::io::BufReader;
34use std::str::FromStr;
35use std::sync::Arc;
36
37/// Command
38#[derive(Debug)]
39pub enum Command {
40    Quit,
41    Help,
42    ListTables,
43    DescribeTableStmt(String),
44    ListFunctions,
45    Include(Option<String>),
46    SearchFunctions(String),
47    QuietMode(Option<bool>),
48    OutputFormat(Option<String>),
49}
50
51pub enum OutputFormat {
52    ChangeFormat(String),
53}
54
55impl Command {
56    pub async fn execute(
57        &self,
58        ctx: &dyn CliSessionContext,
59        print_options: &mut PrintOptions,
60    ) -> Result<()> {
61        match self {
62            Self::Help => {
63                let now = Instant::now();
64                let command_batch = all_commands_info();
65                print_options.print_batches(command_batch.schema(), &[command_batch], now)
66            }
67            Self::ListTables => {
68                exec_and_print(ctx, print_options, "SHOW TABLES".into()).await
69            }
70            Self::DescribeTableStmt(name) => {
71                exec_and_print(ctx, print_options, format!("SHOW COLUMNS FROM {}", name))
72                    .await
73            }
74            Self::Include(filename) => {
75                if let Some(filename) = filename {
76                    let file = File::open(filename).map_err(|e| {
77                        DataFusionError::Execution(format!(
78                            "Error opening {:?} {}",
79                            filename, e
80                        ))
81                    })?;
82                    exec_from_lines(ctx, &mut BufReader::new(file), print_options)
83                        .await?;
84                    Ok(())
85                } else {
86                    exec_err!("Required filename argument is missing")
87                }
88            }
89            Self::QuietMode(quiet) => {
90                if let Some(quiet) = quiet {
91                    print_options.quiet = *quiet;
92                    println!(
93                        "Quiet mode set to {}",
94                        if print_options.quiet { "true" } else { "false" }
95                    );
96                } else {
97                    println!(
98                        "Quiet mode is {}",
99                        if print_options.quiet { "true" } else { "false" }
100                    );
101                }
102                Ok(())
103            }
104            Self::Quit => exec_err!("Unexpected quit, this should be handled outside"),
105            Self::ListFunctions => display_all_functions(),
106            Self::SearchFunctions(function) => {
107                if let Ok(func) = function.parse::<Function>() {
108                    let details = func.function_details()?;
109                    println!("{}", details);
110                    Ok(())
111                } else {
112                    exec_err!("{function} is not a supported function")
113                }
114            }
115            Self::OutputFormat(_) => exec_err!(
116                "Unexpected change output format, this should be handled outside"
117            ),
118        }
119    }
120
121    fn get_name_and_description(&self) -> (&'static str, &'static str) {
122        match self {
123            Self::Quit => ("\\q", "quit datafusion-cli"),
124            Self::ListTables => ("\\d", "list tables"),
125            Self::DescribeTableStmt(_) => ("\\d name", "describe table"),
126            Self::Help => ("\\?", "help"),
127            Self::Include(_) => {
128                ("\\i filename", "reads input from the specified filename")
129            }
130            Self::ListFunctions => ("\\h", "function list"),
131            Self::SearchFunctions(_) => ("\\h function", "search function"),
132            Self::QuietMode(_) => ("\\quiet (true|false)?", "print or set quiet mode"),
133            Self::OutputFormat(_) => {
134                ("\\pset [NAME [VALUE]]", "set table output option\n(format)")
135            }
136        }
137    }
138}
139
140const ALL_COMMANDS: [Command; 9] = [
141    Command::ListTables,
142    Command::DescribeTableStmt(String::new()),
143    Command::Quit,
144    Command::Help,
145    Command::Include(Some(String::new())),
146    Command::ListFunctions,
147    Command::SearchFunctions(String::new()),
148    Command::QuietMode(None),
149    Command::OutputFormat(None),
150];
151
152fn all_commands_info() -> RecordBatch {
153    let schema = Arc::new(Schema::new(vec![
154        Field::new("Command", DataType::Utf8, false),
155        Field::new("Description", DataType::Utf8, false),
156    ]));
157    let (names, description): (Vec<&str>, Vec<&str>) = ALL_COMMANDS
158        .into_iter()
159        .map(|c| c.get_name_and_description())
160        .unzip();
161    RecordBatch::try_new(
162        schema,
163        [names, description]
164            .into_iter()
165            .map(|i| Arc::new(StringArray::from(i)) as ArrayRef)
166            .collect::<Vec<_>>(),
167    )
168    .expect("This should not fail")
169}
170
171impl FromStr for Command {
172    type Err = ();
173
174    fn from_str(s: &str) -> Result<Self, Self::Err> {
175        let (c, arg) = if let Some((a, b)) = s.split_once(' ') {
176            (a, Some(b))
177        } else {
178            (s, None)
179        };
180        Ok(match (c, arg) {
181            ("q", None) => Self::Quit,
182            ("d", None) => Self::ListTables,
183            ("d", Some(name)) => Self::DescribeTableStmt(name.into()),
184            ("?", None) => Self::Help,
185            ("h", None) => Self::ListFunctions,
186            ("h", Some(function)) => Self::SearchFunctions(function.into()),
187            ("i", None) => Self::Include(None),
188            ("i", Some(filename)) => Self::Include(Some(filename.to_owned())),
189            ("quiet", Some("true" | "t" | "yes" | "y" | "on")) => {
190                Self::QuietMode(Some(true))
191            }
192            ("quiet", Some("false" | "f" | "no" | "n" | "off")) => {
193                Self::QuietMode(Some(false))
194            }
195            ("quiet", None) => Self::QuietMode(None),
196            ("pset", Some(subcommand)) => {
197                Self::OutputFormat(Some(subcommand.to_string()))
198            }
199            ("pset", None) => Self::OutputFormat(None),
200            _ => return Err(()),
201        })
202    }
203}
204
205impl FromStr for OutputFormat {
206    type Err = ();
207
208    fn from_str(s: &str) -> Result<Self, Self::Err> {
209        let (c, arg) = if let Some((a, b)) = s.split_once(' ') {
210            (a, Some(b))
211        } else {
212            (s, None)
213        };
214        Ok(match (c, arg) {
215            ("format", Some(format)) => Self::ChangeFormat(format.to_string()),
216            _ => return Err(()),
217        })
218    }
219}
220
221impl OutputFormat {
222    pub async fn execute(&self, print_options: &mut PrintOptions) -> Result<()> {
223        match self {
224            Self::ChangeFormat(format) => {
225                if let Ok(format) = format.parse::<PrintFormat>() {
226                    print_options.format = format;
227                    println!("Output format is {:?}.", print_options.format);
228                    Ok(())
229                } else {
230                    exec_err!(
231                        "{:?} is not a valid format type [possible values: {:?}]",
232                        format,
233                        PrintFormat::value_variants()
234                    )
235                }
236            }
237        }
238    }
239}
frostbow_cli/command.rs

frostbow_cli/
command.rs
