datafusion_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                let schema = command_batch.schema();
66                let num_rows = command_batch.num_rows();
67                let task_ctx = ctx.task_ctx();
68                let config = &task_ctx.session_config().options().format;
69                print_options.print_batches(
70                    schema,
71                    &[command_batch],
72                    now,
73                    num_rows,
74                    config,
75                )
76            }
77            Self::ListTables => {
78                exec_and_print(ctx, print_options, "SHOW TABLES".into()).await
79            }
80            Self::DescribeTableStmt(name) => {
81                exec_and_print(ctx, print_options, format!("SHOW COLUMNS FROM {name}"))
82                    .await
83            }
84            Self::Include(filename) => {
85                if let Some(filename) = filename {
86                    let file = File::open(filename).map_err(|e| {
87                        DataFusionError::Execution(format!(
88                            "Error opening {filename:?} {e}"
89                        ))
90                    })?;
91                    exec_from_lines(ctx, &mut BufReader::new(file), print_options)
92                        .await?;
93                    Ok(())
94                } else {
95                    exec_err!("Required filename argument is missing")
96                }
97            }
98            Self::QuietMode(quiet) => {
99                if let Some(quiet) = quiet {
100                    print_options.quiet = *quiet;
101                    println!(
102                        "Quiet mode set to {}",
103                        if print_options.quiet { "true" } else { "false" }
104                    );
105                } else {
106                    println!(
107                        "Quiet mode is {}",
108                        if print_options.quiet { "true" } else { "false" }
109                    );
110                }
111                Ok(())
112            }
113            Self::Quit => exec_err!("Unexpected quit, this should be handled outside"),
114            Self::ListFunctions => display_all_functions(),
115            Self::SearchFunctions(function) => {
116                if let Ok(func) = function.parse::<Function>() {
117                    let details = func.function_details()?;
118                    println!("{details}");
119                    Ok(())
120                } else {
121                    exec_err!("{function} is not a supported function")
122                }
123            }
124            Self::OutputFormat(_) => exec_err!(
125                "Unexpected change output format, this should be handled outside"
126            ),
127        }
128    }
129
130    fn get_name_and_description(&self) -> (&'static str, &'static str) {
131        match self {
132            Self::Quit => ("\\q", "quit datafusion-cli"),
133            Self::ListTables => ("\\d", "list tables"),
134            Self::DescribeTableStmt(_) => ("\\d name", "describe table"),
135            Self::Help => ("\\?", "help"),
136            Self::Include(_) => {
137                ("\\i filename", "reads input from the specified filename")
138            }
139            Self::ListFunctions => ("\\h", "function list"),
140            Self::SearchFunctions(_) => ("\\h function", "search function"),
141            Self::QuietMode(_) => ("\\quiet (true|false)?", "print or set quiet mode"),
142            Self::OutputFormat(_) => {
143                ("\\pset [NAME [VALUE]]", "set table output option\n(format)")
144            }
145        }
146    }
147}
148
149const ALL_COMMANDS: [Command; 9] = [
150    Command::ListTables,
151    Command::DescribeTableStmt(String::new()),
152    Command::Quit,
153    Command::Help,
154    Command::Include(Some(String::new())),
155    Command::ListFunctions,
156    Command::SearchFunctions(String::new()),
157    Command::QuietMode(None),
158    Command::OutputFormat(None),
159];
160
161fn all_commands_info() -> RecordBatch {
162    let schema = Arc::new(Schema::new(vec![
163        Field::new("Command", DataType::Utf8, false),
164        Field::new("Description", DataType::Utf8, false),
165    ]));
166    let (names, description): (Vec<&str>, Vec<&str>) = ALL_COMMANDS
167        .into_iter()
168        .map(|c| c.get_name_and_description())
169        .unzip();
170    RecordBatch::try_new(
171        schema,
172        [names, description]
173            .into_iter()
174            .map(|i| Arc::new(StringArray::from(i)) as ArrayRef)
175            .collect::<Vec<_>>(),
176    )
177    .expect("This should not fail")
178}
179
180impl FromStr for Command {
181    type Err = ();
182
183    fn from_str(s: &str) -> Result<Self, Self::Err> {
184        let (c, arg) = if let Some((a, b)) = s.split_once(' ') {
185            (a, Some(b))
186        } else {
187            (s, None)
188        };
189        Ok(match (c, arg) {
190            ("q", None) => Self::Quit,
191            ("d", None) => Self::ListTables,
192            ("d", Some(name)) => Self::DescribeTableStmt(name.into()),
193            ("?", None) => Self::Help,
194            ("h", None) => Self::ListFunctions,
195            ("h", Some(function)) => Self::SearchFunctions(function.into()),
196            ("i", None) => Self::Include(None),
197            ("i", Some(filename)) => Self::Include(Some(filename.to_owned())),
198            ("quiet", Some("true" | "t" | "yes" | "y" | "on")) => {
199                Self::QuietMode(Some(true))
200            }
201            ("quiet", Some("false" | "f" | "no" | "n" | "off")) => {
202                Self::QuietMode(Some(false))
203            }
204            ("quiet", None) => Self::QuietMode(None),
205            ("pset", Some(subcommand)) => {
206                Self::OutputFormat(Some(subcommand.to_string()))
207            }
208            ("pset", None) => Self::OutputFormat(None),
209            _ => return Err(()),
210        })
211    }
212}
213
214impl FromStr for OutputFormat {
215    type Err = ();
216
217    fn from_str(s: &str) -> Result<Self, Self::Err> {
218        let (c, arg) = if let Some((a, b)) = s.split_once(' ') {
219            (a, Some(b))
220        } else {
221            (s, None)
222        };
223        Ok(match (c, arg) {
224            ("format", Some(format)) => Self::ChangeFormat(format.to_string()),
225            _ => return Err(()),
226        })
227    }
228}
229
230impl OutputFormat {
231    pub async fn execute(&self, print_options: &mut PrintOptions) -> Result<()> {
232        match self {
233            Self::ChangeFormat(format) => {
234                if let Ok(format) = format.parse::<PrintFormat>() {
235                    print_options.format = format;
236                    println!("Output format is {:?}.", print_options.format);
237                    Ok(())
238                } else {
239                    exec_err!(
240                        "{:?} is not a valid format type [possible values: {:?}]",
241                        format,
242                        PrintFormat::value_variants()
243                    )
244                }
245            }
246        }
247    }
248}
datafusion_cli/command.rs

datafusion_cli/
command.rs
