diff --git a/crates/nu_plugin_from_sqlite/src/from_sqlite.rs b/crates/nu_plugin_from_sqlite/src/from_sqlite.rs index fbd112df7d..b837a75b5e 100644 --- a/crates/nu_plugin_from_sqlite/src/from_sqlite.rs +++ b/crates/nu_plugin_from_sqlite/src/from_sqlite.rs @@ -10,6 +10,7 @@ use std::path::Path; pub struct FromSqlite { pub state: Vec, pub name_tag: Tag, + pub tables: Vec, } 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 + Clone, + tables: Vec, ) -> Result { 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 + Clone pub fn from_sqlite_bytes_to_value( mut bytes: Vec, tag: impl Into + Clone, + tables: Vec, ) -> Result { // 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, name_tag: Tag) -> Result, ShellError> { - match from_sqlite_bytes_to_value(bytes, name_tag.clone()) { +pub fn from_sqlite( + bytes: Vec, + name_tag: Tag, + tables: Vec, +) -> Result, ShellError> { + match from_sqlite_bytes_to_value(bytes, name_tag.clone(), tables) { Ok(x) => match x { Value { value: UntaggedValue::Table(list), diff --git a/crates/nu_plugin_from_sqlite/src/nu/mod.rs b/crates/nu_plugin_from_sqlite/src/nu/mod.rs index deac5dafcd..f39463386e 100644 --- a/crates/nu_plugin_from_sqlite/src/nu/mod.rs +++ b/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, 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::, _>>()?; + + Ok(res) +} + impl Plugin for FromSqlite { fn config(&mut self) -> Result { 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, 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, ShellError> { - crate::from_sqlite::from_sqlite(self.state.clone(), Tag::unknown()) + crate::from_sqlite::from_sqlite(self.state.clone(), Tag::unknown(), self.tables.clone()) } } diff --git a/docs/commands/from-sqlite.md b/docs/commands/from-sqlite.md index 168681c505..0ac84120bf 100644 --- a/docs/commands/from-sqlite.md +++ b/docs/commands/from-sqlite.md @@ -8,4 +8,4 @@ Convert from sqlite binary into table ## Flags * -h, --help: Display this help message - +* -t, --tables \[\ \ ... \]: Only convert specified tables