add param validation for datafusion-cli (#284)

* add param validation for datafusion-cli * Update datafusion-cli/src/main.rs Co-authored-by: Andrew Lamb <[email protected]> * Update datafusion-cli/src/main.rs Co-authored-by: Andrew Lamb <[email protected]> Co-authored-by: Andrew Lamb <[email protected]>
apache · May 8, 2021 · a947f11 · a947f11
1 parent b8805d4
commit a947f11
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/datafusion-cli/src/main.rs b/datafusion-cli/src/main.rs
@@ -40,13 +40,15 @@ pub async fn main() {
                 .help("Path to your data, default to current directory")
                 .short("p")
                 .long("data-path")
+                .validator(is_valid_data_dir)
                 .takes_value(true),
         )
         .arg(
             Arg::with_name("batch-size")
                 .help("The batch size of each query, or use DataFusion default")
                 .short("c")
                 .long("batch-size")
+                .validator(is_valid_batch_size)
                 .takes_value(true),
         )
         .get_matches();
@@ -100,6 +102,21 @@ pub async fn main() {
     rl.save_history(".history").ok();
 }
 
+fn is_valid_data_dir(dir: String) -> std::result::Result<(), String> {
+    if Path::new(&dir).is_dir() {
+        Ok(())
+    } else {
+        Err(format!("Invalid data directory '{}'", dir))
+    }
+}
+
+fn is_valid_batch_size(size: String) -> std::result::Result<(), String> {
+    match size.parse::<usize>() {
+        Ok(size) if size > 0 => Ok(()),
+        _ => Err(format!("Invalid batch size '{}'", size)),
+    }
+}
+
 fn is_exit_command(line: &str) -> bool {
     let line = line.trim_end().to_lowercase();
     line == "quit" || line == "exit"