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

* from sqlite: Add parameter '--tables'

* from sqlite: Enhance documentation
This commit is contained in:
Christian Menges 2021-06-02 00:06:32 +02:00 committed by GitHub
parent 7bf10b980c
commit 2486492c4d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 60 additions and 21 deletions

View File

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

View File

@ -4,18 +4,47 @@ mod tests;
use crate::FromSqlite; use crate::FromSqlite;
use nu_errors::ShellError; use nu_errors::ShellError;
use nu_plugin::Plugin; 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; 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 { impl Plugin for FromSqlite {
fn config(&mut self) -> Result<Signature, ShellError> { fn config(&mut self) -> Result<Signature, ShellError> {
Ok(Signature::build("from sqlite") Ok(Signature::build("from sqlite")
.named(
"tables",
SyntaxShape::Table,
"Only convert specified tables",
Some('t'),
)
.desc("Convert from sqlite binary into table") .desc("Convert from sqlite binary into table")
.filter()) .filter())
} }
fn begin_filter(&mut self, call_info: CallInfo) -> Result<Vec<ReturnValue>, ShellError> { fn begin_filter(&mut self, call_info: CallInfo) -> Result<Vec<ReturnValue>, ShellError> {
self.name_tag = call_info.name_tag; 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![]) Ok(vec![])
} }
@ -41,6 +70,6 @@ impl Plugin for FromSqlite {
} }
fn end_filter(&mut self) -> Result<Vec<ReturnValue>, ShellError> { 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())
} }
} }

View File

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