nushell/crates/nu_plugin_to_sqlite/src/to_sqlite.rs
2021-09-10 10:44:22 +12:00

173 lines
5.7 KiB
Rust

use hex::encode;
use nu_errors::ShellError;
use nu_protocol::{Dictionary, Primitive, ReturnSuccess, ReturnValue, UntaggedValue, Value};
use nu_source::Tag;
use rusqlite::Connection;
use std::io::Read;
#[derive(Default)]
pub struct ToSqlite {
pub state: Vec<Value>,
}
impl ToSqlite {
pub fn new() -> ToSqlite {
ToSqlite { state: vec![] }
}
}
fn comma_concat(acc: String, current: String) -> String {
if acc.is_empty() {
current
} else {
format!("{}, {}", acc, current)
}
}
fn get_columns(rows: &[Value]) -> Result<String, std::io::Error> {
match &rows[0].value {
UntaggedValue::Row(d) => Ok(d
.entries
.iter()
.map(|(k, _v)| k.clone())
.fold("".to_string(), comma_concat)),
_ => Err(std::io::Error::new(
std::io::ErrorKind::Other,
"Could not find table column names",
)),
}
}
fn nu_value_to_sqlite_string(v: Value) -> String {
match &v.value {
UntaggedValue::Primitive(p) => match p {
Primitive::Nothing => "NULL".into(),
Primitive::BigInt(i) => i.to_string(),
Primitive::Int(i) => i.to_string(),
Primitive::Duration(i) => i.to_string(),
Primitive::Decimal(f) => f.to_string(),
Primitive::Filesize(u) => u.to_string(),
Primitive::GlobPattern(s) => format!("'{}'", s.replace("'", "''")),
Primitive::String(s) => format!("'{}'", s.replace("'", "''")),
Primitive::Boolean(true) => "1".into(),
Primitive::Boolean(_) => "0".into(),
Primitive::Date(d) => format!("'{}'", d),
Primitive::FilePath(p) => format!("'{}'", p.display().to_string().replace("'", "''")),
Primitive::Binary(u) => format!("x'{}'", encode(u)),
Primitive::BeginningOfStream
| Primitive::EndOfStream
| Primitive::ColumnPath(_)
| Primitive::Range(_) => "NULL".into(),
},
_ => "NULL".into(),
}
}
fn get_insert_values(rows: Vec<Value>) -> Result<String, std::io::Error> {
let values: Result<Vec<_>, _> = rows
.into_iter()
.map(|value| match value.value {
UntaggedValue::Row(d) => Ok(format!(
"({})",
d.entries
.iter()
.map(|(_k, v)| nu_value_to_sqlite_string(v.clone()))
.fold("".to_string(), comma_concat)
)),
_ => Err(std::io::Error::new(
std::io::ErrorKind::Other,
"Could not find table column names",
)),
})
.collect();
let values = values?;
Ok(values.join(", "))
}
fn generate_statements(table: Dictionary) -> Result<(String, String), std::io::Error> {
let table_name = match table.entries.get("table_name") {
Some(Value {
value: UntaggedValue::Primitive(Primitive::String(table_name)),
..
}) => table_name,
_ => {
return Err(std::io::Error::new(
std::io::ErrorKind::Other,
"Could not find table name",
))
}
};
let (columns, insert_values) = match table.entries.get("table_values") {
Some(Value {
value: UntaggedValue::Table(l),
..
}) => {
if l.is_empty() {
return Ok((String::new(), String::new()));
}
(get_columns(l), get_insert_values(l.to_vec()))
}
_ => {
return Err(std::io::Error::new(
std::io::ErrorKind::Other,
"Could not find table values",
))
}
};
let create = format!("create table {}({})", table_name, columns?);
let insert = format!("insert into {} values {}", table_name, insert_values?);
Ok((create, insert))
}
fn sqlite_input_stream_to_bytes(values: Vec<Value>) -> 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
// write out, but this will require C code. Might be
// best done as a PR to rusqlite.
let mut tempfile = tempfile::NamedTempFile::new()?;
let conn = match Connection::open(tempfile.path()) {
Ok(conn) => conn,
Err(e) => return Err(std::io::Error::new(std::io::ErrorKind::Other, e)),
};
let tag = values[0].tag.clone();
for value in values {
match &value.value {
UntaggedValue::Row(d) => {
let (create, insert) = generate_statements(d.to_owned())?;
if create.is_empty() {
continue;
}
match conn
.execute(&create, [])
.and_then(|_| conn.execute(&insert, []))
{
Ok(_) => (),
Err(e) => {
return Err(std::io::Error::new(std::io::ErrorKind::Other, e));
}
}
}
other => {
return Err(std::io::Error::new(
std::io::ErrorKind::Other,
format!("Expected row, found {:?}", other),
))
}
}
}
let mut out = Vec::new();
tempfile.read_to_end(&mut out)?;
Ok(UntaggedValue::binary(out).into_value(tag))
}
pub fn to_sqlite(input: Vec<Value>, name_tag: Tag) -> Result<Vec<ReturnValue>, ShellError> {
match sqlite_input_stream_to_bytes(input) {
Ok(out) => Ok(vec![ReturnSuccess::value(out)]),
_ => Err(ShellError::labeled_error(
"Expected a table with SQLite-compatible structure from pipeline",
"requires SQLite-compatible input",
name_tag,
)),
}
}