ballista_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 std::str::FromStr;
21use std::sync::Arc;
22use std::time::Instant;
23
24use datafusion::arrow::array::{ArrayRef, StringArray};
25use datafusion::arrow::datatypes::{DataType, Field, Schema};
26use datafusion::arrow::record_batch::RecordBatch;
27use datafusion::common::Result;
28use datafusion::error::DataFusionError;
29use datafusion::prelude::SessionContext;
30
31use crate::functions::{display_all_functions, Function};
32use crate::print_format::PrintFormat;
33use crate::print_options::PrintOptions;
34
35/// Command
36#[derive(Debug)]
37pub enum Command {
38    Quit,
39    Help,
40    ListTables,
41    DescribeTable(String),
42    ListFunctions,
43    SearchFunctions(String),
44    QuietMode(Option<bool>),
45    OutputFormat(Option<String>),
46}
47
48pub enum OutputFormat {
49    ChangeFormat(String),
50}
51
52impl Command {
53    pub async fn execute(
54        &self,
55        ctx: &SessionContext,
56        print_options: &mut PrintOptions,
57    ) -> Result<()> {
58        let now = Instant::now();
59        let max_rows = match print_options.maxrows {
60            datafusion_cli::print_options::MaxRows::Unlimited => usize::MAX,
61            datafusion_cli::print_options::MaxRows::Limited(max_rows) => max_rows,
62        };
63        match self {
64            Self::Help =>
65            // TODO need to provide valid schema
66            {
67                print_options.print_batches(
68                    Arc::new(Schema::empty()),
69                    &[all_commands_info()],
70                    now,
71                    max_rows,
72                    &Default::default(),
73                )
74            }
75            Self::ListTables => {
76                let df = ctx.sql("SHOW TABLES").await?;
77                let schema = Arc::new(df.schema().as_arrow().clone());
78                let batches = df.collect().await?;
79                print_options.print_batches(
80                    schema,
81                    &batches,
82                    now,
83                    max_rows,
84                    &Default::default(),
85                )
86            }
87            Self::DescribeTable(name) => {
88                let df = ctx.sql(&format!("SHOW COLUMNS FROM {name}")).await?;
89                let schema = Arc::new(df.schema().as_arrow().clone());
90                let batches = df.collect().await?;
91                print_options.print_batches(
92                    schema,
93                    &batches,
94                    now,
95                    max_rows,
96                    &Default::default(),
97                )
98            }
99            Self::QuietMode(quiet) => {
100                if let Some(quiet) = quiet {
101                    print_options.quiet = *quiet;
102                    println!(
103                        "Quiet mode set to {}",
104                        if print_options.quiet { "true" } else { "false" }
105                    );
106                } else {
107                    println!(
108                        "Quiet mode is {}",
109                        if print_options.quiet { "true" } else { "false" }
110                    );
111                }
112                Ok(())
113            }
114            Self::Quit => Err(DataFusionError::Internal(
115                "Unexpected quit, this should be handled outside".to_string(),
116            )),
117            Self::ListFunctions => display_all_functions(),
118            Self::SearchFunctions(function) => {
119                if let Ok(func) = function.parse::<Function>() {
120                    let details = func.function_details()?;
121                    println!("{details}");
122                    Ok(())
123                } else {
124                    let msg = format!("{function} is not a supported function");
125                    Err(DataFusionError::NotImplemented(msg))
126                }
127            }
128            Self::OutputFormat(_) => Err(DataFusionError::Internal(
129                "Unexpected change output format, this should be handled outside"
130                    .to_string(),
131            )),
132        }
133    }
134
135    fn get_name_and_description(&self) -> (&'static str, &'static str) {
136        match self {
137            Self::Quit => ("\\q", "quit ballista-cli"),
138            Self::ListTables => ("\\d", "list tables"),
139            Self::DescribeTable(_) => ("\\d name", "describe table"),
140            Self::Help => ("\\?", "help"),
141            Self::ListFunctions => ("\\h", "function list"),
142            Self::SearchFunctions(_) => ("\\h function", "search function"),
143            Self::QuietMode(_) => ("\\quiet (true|false)?", "print or set quiet mode"),
144            Self::OutputFormat(_) => {
145                ("\\pset [NAME [VALUE]]", "set table output option\n(format)")
146            }
147        }
148    }
149}
150
151const ALL_COMMANDS: [Command; 8] = [
152    Command::ListTables,
153    Command::DescribeTable(String::new()),
154    Command::Quit,
155    Command::Help,
156    Command::ListFunctions,
157    Command::SearchFunctions(String::new()),
158    Command::QuietMode(None),
159    Command::OutputFormat(None),
160];
161
162fn all_commands_info() -> RecordBatch {
163    let schema = Arc::new(Schema::new(vec![
164        Field::new("Command", DataType::Utf8, false),
165        Field::new("Description", DataType::Utf8, false),
166    ]));
167    let (names, description): (Vec<&str>, Vec<&str>) = ALL_COMMANDS
168        .into_iter()
169        .map(|c| c.get_name_and_description())
170        .unzip();
171    RecordBatch::try_new(
172        schema,
173        [names, description]
174            .into_iter()
175            .map(|i| Arc::new(StringArray::from(i)) as ArrayRef)
176            .collect::<Vec<_>>(),
177    )
178    .expect("This should not fail")
179}
180
181impl FromStr for Command {
182    type Err = ();
183
184    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
185        let (c, arg) = if let Some((a, b)) = s.split_once(' ') {
186            (a, Some(b))
187        } else {
188            (s, None)
189        };
190        Ok(match (c, arg) {
191            ("q", None) => Self::Quit,
192            ("d", None) => Self::ListTables,
193            ("d", Some(name)) => Self::DescribeTable(name.into()),
194            ("?", None) => Self::Help,
195            ("h", None) => Self::ListFunctions,
196            ("h", Some(function)) => Self::SearchFunctions(function.into()),
197            ("quiet", Some("true" | "t" | "yes" | "y" | "on")) => {
198                Self::QuietMode(Some(true))
199            }
200            ("quiet", Some("false" | "f" | "no" | "n" | "off")) => {
201                Self::QuietMode(Some(false))
202            }
203            ("quiet", None) => Self::QuietMode(None),
204            ("pset", Some(subcommand)) => {
205                Self::OutputFormat(Some(subcommand.to_string()))
206            }
207            ("pset", None) => Self::OutputFormat(None),
208            _ => return Err(()),
209        })
210    }
211}
212
213impl FromStr for OutputFormat {
214    type Err = ();
215
216    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
217        let (c, arg) = if let Some((a, b)) = s.split_once(' ') {
218            (a, Some(b))
219        } else {
220            (s, None)
221        };
222        Ok(match (c, arg) {
223            ("format", Some(format)) => Self::ChangeFormat(format.to_string()),
224            _ => return Err(()),
225        })
226    }
227}
228
229impl OutputFormat {
230    pub async fn execute(&self, print_options: &mut PrintOptions) -> Result<()> {
231        match self {
232            Self::ChangeFormat(format) => {
233                if let Ok(format) = format.parse::<PrintFormat>() {
234                    print_options.format = format;
235                    println!("Output format is {:?}.", print_options.format);
236                    Ok(())
237                } else {
238                    Err(DataFusionError::Execution(format!(
239                        "{:?} is not a valid format type [possible values: {:?}]",
240                        format,
241                        "TO BE FIXED", //PrintFormat::value_variants()
242                    )))
243                }
244            }
245        }
246    }
247}
ballista_cli/command.rs

ballista_cli/
command.rs
