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

kapot_cli/
command.rs
