nu_plugin_from_sqlite/
from_sqlite.rs

1use bigdecimal::FromPrimitive;
2use nu_errors::ShellError;
3use nu_protocol::{Primitive, ReturnSuccess, ReturnValue, TaggedDictBuilder, UntaggedValue, Value};
4use nu_source::Tag;
5use rusqlite::{types::ValueRef, Connection, Row};
6use std::io::Write;
7use std::path::Path;
8
9#[derive(Default)]
10pub struct FromSqlite {
11    pub state: Vec<u8>,
12    pub name_tag: Tag,
13    pub tables: Vec<String>,
14}
15
16impl FromSqlite {
17    pub fn new() -> FromSqlite {
18        FromSqlite {
19            state: vec![],
20            name_tag: Tag::unknown(),
21            tables: vec![],
22        }
23    }
24}
25
26pub fn convert_sqlite_file_to_nu_value(
27    path: &Path,
28    tag: impl Into<Tag> + Clone,
29    tables: Vec<String>,
30) -> Result<Value, rusqlite::Error> {
31    let conn = Connection::open(path)?;
32
33    let mut meta_out = Vec::new();
34    let mut meta_stmt = conn.prepare("select name from sqlite_master where type='table'")?;
35    let mut meta_rows = meta_stmt.query([])?;
36
37    while let Some(meta_row) = meta_rows.next()? {
38        let table_name: String = meta_row.get(0)?;
39        if tables.is_empty() || tables.contains(&table_name) {
40            let mut meta_dict = TaggedDictBuilder::new(tag.clone());
41            let mut out = Vec::new();
42            let mut table_stmt = conn.prepare(&format!("select * from [{}]", table_name))?;
43            let mut table_rows = table_stmt.query([])?;
44            while let Some(table_row) = table_rows.next()? {
45                out.push(convert_sqlite_row_to_nu_value(table_row, tag.clone()))
46            }
47            meta_dict.insert_value(
48                "table_name".to_string(),
49                UntaggedValue::Primitive(Primitive::String(table_name)).into_value(tag.clone()),
50            );
51            meta_dict.insert_value(
52                "table_values",
53                UntaggedValue::Table(out).into_value(tag.clone()),
54            );
55            meta_out.push(meta_dict.into_value());
56        }
57    }
58    let tag = tag.into();
59    Ok(UntaggedValue::Table(meta_out).into_value(tag))
60}
61
62fn convert_sqlite_row_to_nu_value(row: &Row, tag: impl Into<Tag> + Clone) -> Value {
63    let mut collected = TaggedDictBuilder::new(tag.clone());
64    for (i, c) in row.as_ref().column_names().iter().enumerate() {
65        collected.insert_value(
66            c.to_string(),
67            convert_sqlite_value_to_nu_value(row.get_ref_unwrap(i), tag.clone()),
68        );
69    }
70    collected.into_value()
71}
72
73fn convert_sqlite_value_to_nu_value(value: ValueRef, tag: impl Into<Tag> + Clone) -> Value {
74    match value {
75        ValueRef::Null => {
76            UntaggedValue::Primitive(Primitive::String(String::from(""))).into_value(tag)
77        }
78        ValueRef::Integer(i) => UntaggedValue::int(i).into_value(tag),
79        ValueRef::Real(f) => {
80            let f = bigdecimal::BigDecimal::from_f64(f);
81            let tag = tag.into();
82            let span = tag.span;
83            match f {
84                Some(d) => UntaggedValue::decimal(d).into_value(tag),
85                None => UntaggedValue::Error(ShellError::labeled_error(
86                    "Can not convert f64 to big decimal",
87                    "can not convert to decimal",
88                    span,
89                ))
90                .into_value(tag),
91            }
92        }
93        ValueRef::Text(s) => {
94            // this unwrap is safe because we know the ValueRef is Text.
95            UntaggedValue::Primitive(Primitive::String(String::from_utf8_lossy(s).to_string()))
96                .into_value(tag)
97        }
98        ValueRef::Blob(u) => UntaggedValue::binary(u.to_owned()).into_value(tag),
99    }
100}
101
102pub fn from_sqlite_bytes_to_value(
103    mut bytes: Vec<u8>,
104    tag: impl Into<Tag> + Clone,
105    tables: Vec<String>,
106) -> Result<Value, std::io::Error> {
107    // FIXME: should probably write a sqlite virtual filesystem
108    // that will allow us to use bytes as a file to avoid this
109    // write out, but this will require C code. Might be
110    // best done as a PR to rusqlite.
111    let mut tempfile = tempfile::NamedTempFile::new()?;
112    tempfile.write_all(bytes.as_mut_slice())?;
113    match convert_sqlite_file_to_nu_value(tempfile.path(), tag, tables) {
114        Ok(value) => Ok(value),
115        Err(e) => Err(std::io::Error::new(std::io::ErrorKind::Other, e)),
116    }
117}
118
119pub fn from_sqlite(
120    bytes: Vec<u8>,
121    name_tag: Tag,
122    tables: Vec<String>,
123) -> Result<Vec<ReturnValue>, ShellError> {
124    match from_sqlite_bytes_to_value(bytes, name_tag.clone(), tables) {
125        Ok(x) => match x {
126            Value {
127                value: UntaggedValue::Table(list),
128                ..
129            } => Ok(list.into_iter().map(ReturnSuccess::value).collect()),
130            _ => Ok(vec![ReturnSuccess::value(x)]),
131        },
132        Err(_) => Err(ShellError::labeled_error(
133            "Could not parse as SQLite",
134            "input cannot be parsed as SQLite",
135            &name_tag,
136        )),
137    }
138}
nu_plugin_from_sqlite/from_sqlite.rs

nu_plugin_from_sqlite/
from_sqlite.rs
