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

datafusion_cli/
command.rs
