from sqlite: Add parameter --tables (#3529)

* from sqlite: Add parameter '--tables'

* from sqlite: Enhance documentation

crates/nu_plugin_from_sqlite/src/from_sqlite.rs

@@ -10,6 +10,7 @@ use std::path::Path;
 pub struct FromSqlite {
     pub state: Vec<u8>,
     pub name_tag: Tag,
+    pub tables: Vec<String>,
 }
 
 impl FromSqlite {
@@ -17,6 +18,7 @@ impl FromSqlite {
         FromSqlite {
             state: vec![],
             name_tag: Tag::unknown(),
+            tables: vec![],
         }
     }
 }
@@ -24,6 +26,7 @@ impl FromSqlite {
 pub fn convert_sqlite_file_to_nu_value(
     path: &Path,
     tag: impl Into<Tag> + Clone,
+    tables: Vec<String>,
 ) -> Result<Value, rusqlite::Error> {
     let conn = Connection::open(path)?;
 
@@ -33,22 +36,24 @@ pub fn convert_sqlite_file_to_nu_value(
 
     while let Some(meta_row) = meta_rows.next()? {
         let table_name: String = meta_row.get(0)?;
-        let mut meta_dict = TaggedDictBuilder::new(tag.clone());
-        let mut out = Vec::new();
-        let mut table_stmt = conn.prepare(&format!("select * from [{}]", table_name))?;
-        let mut table_rows = table_stmt.query([])?;
-        while let Some(table_row) = table_rows.next()? {
-            out.push(convert_sqlite_row_to_nu_value(table_row, tag.clone()))
+        if tables.is_empty() || tables.contains(&table_name) {
+            let mut meta_dict = TaggedDictBuilder::new(tag.clone());
+            let mut out = Vec::new();
+            let mut table_stmt = conn.prepare(&format!("select * from [{}]", table_name))?;
+            let mut table_rows = table_stmt.query([])?;
+            while let Some(table_row) = table_rows.next()? {
+                out.push(convert_sqlite_row_to_nu_value(table_row, tag.clone()))
+            }
+            meta_dict.insert_value(
+                "table_name".to_string(),
+                UntaggedValue::Primitive(Primitive::String(table_name)).into_value(tag.clone()),
+            );
+            meta_dict.insert_value(
+                "table_values",
+                UntaggedValue::Table(out).into_value(tag.clone()),
+            );
+            meta_out.push(meta_dict.into_value());
         }
-        meta_dict.insert_value(
-            "table_name".to_string(),
-            UntaggedValue::Primitive(Primitive::String(table_name)).into_value(tag.clone()),
-        );
-        meta_dict.insert_value(
-            "table_values",
-            UntaggedValue::Table(out).into_value(tag.clone()),
-        );
-        meta_out.push(meta_dict.into_value());
     }
     let tag = tag.into();
     Ok(UntaggedValue::Table(meta_out).into_value(tag))
@@ -97,6 +102,7 @@ fn convert_sqlite_value_to_nu_value(value: ValueRef, tag: impl Into<Tag> + Clone
 pub fn from_sqlite_bytes_to_value(
     mut bytes: Vec<u8>,
     tag: impl Into<Tag> + Clone,
+    tables: Vec<String>,
 ) -> Result<Value, std::io::Error> {
     // FIXME: should probably write a sqlite virtual filesystem
     // that will allow us to use bytes as a file to avoid this
@@ -104,14 +110,18 @@ pub fn from_sqlite_bytes_to_value(
     // best done as a PR to rusqlite.
     let mut tempfile = tempfile::NamedTempFile::new()?;
     tempfile.write_all(bytes.as_mut_slice())?;
-    match convert_sqlite_file_to_nu_value(tempfile.path(), tag) {
+    match convert_sqlite_file_to_nu_value(tempfile.path(), tag, tables) {
         Ok(value) => Ok(value),
         Err(e) => Err(std::io::Error::new(std::io::ErrorKind::Other, e)),
     }
 }
 
-pub fn from_sqlite(bytes: Vec<u8>, name_tag: Tag) -> Result<Vec<ReturnValue>, ShellError> {
-    match from_sqlite_bytes_to_value(bytes, name_tag.clone()) {
+pub fn from_sqlite(
+    bytes: Vec<u8>,
+    name_tag: Tag,
+    tables: Vec<String>,
+) -> Result<Vec<ReturnValue>, ShellError> {
+    match from_sqlite_bytes_to_value(bytes, name_tag.clone(), tables) {
         Ok(x) => match x {
             Value {
                 value: UntaggedValue::Table(list),

crates/nu_plugin_from_sqlite/src/nu/mod.rs

@@ -4,18 +4,47 @@ mod tests;
 use crate::FromSqlite;
 use nu_errors::ShellError;
 use nu_plugin::Plugin;
-use nu_protocol::{CallInfo, Primitive, ReturnValue, Signature, UntaggedValue, Value};
+use nu_protocol::{CallInfo, Primitive, ReturnValue, Signature, SyntaxShape, UntaggedValue, Value};
 use nu_source::Tag;
 
+// Adapted from crates/nu-command/src/commands/dataframe/utils.rs
+fn convert_columns(columns: &[Value]) -> Result<Vec<String>, ShellError> {
+    let res = columns
+        .iter()
+        .map(|value| match &value.value {
+            UntaggedValue::Primitive(Primitive::String(s)) => Ok(s.clone()),
+            _ => Err(ShellError::labeled_error(
+                "Incorrect column format",
+                "Only string as column name",
+                &value.tag,
+            )),
+        })
+        .collect::<Result<Vec<String>, _>>()?;
+
+    Ok(res)
+}
+
 impl Plugin for FromSqlite {
     fn config(&mut self) -> Result<Signature, ShellError> {
         Ok(Signature::build("from sqlite")
+            .named(
+                "tables",
+                SyntaxShape::Table,
+                "Only convert specified tables",
+                Some('t'),
+            )
             .desc("Convert from sqlite binary into table")
             .filter())
     }
 
     fn begin_filter(&mut self, call_info: CallInfo) -> Result<Vec<ReturnValue>, ShellError> {
         self.name_tag = call_info.name_tag;
+
+        if let Some(t) = call_info.args.get("tables") {
+            if let UntaggedValue::Table(columns) = t.value.clone() {
+                self.tables = convert_columns(columns.as_slice())?;
+            }
+        }
         Ok(vec![])
     }
 
@@ -41,6 +70,6 @@ impl Plugin for FromSqlite {
     }
 
     fn end_filter(&mut self) -> Result<Vec<ReturnValue>, ShellError> {
-        crate::from_sqlite::from_sqlite(self.state.clone(), Tag::unknown())
+        crate::from_sqlite::from_sqlite(self.state.clone(), Tag::unknown(), self.tables.clone())
     }
 }

docs/commands/from-sqlite.md

@@ -8,4 +8,4 @@ Convert from sqlite binary into table
 
 ## Flags
 * -h, --help: Display this help message
-
+* -t, --tables \[\<table_name_1> \<table_name_2> ... \<table_name_N>]: Only convert specified tables