diff --git a/crates/nu-command/src/database/mod.rs b/crates/nu-command/src/database/mod.rs new file mode 100644 index 0000000000..d4f755050d --- /dev/null +++ b/crates/nu-command/src/database/mod.rs @@ -0,0 +1,3 @@ +mod sqlite; + +pub use sqlite::SQLiteDatabase; diff --git a/crates/nu-command/src/database/sqlite.rs b/crates/nu-command/src/database/sqlite.rs new file mode 100644 index 0000000000..550661e6be --- /dev/null +++ b/crates/nu-command/src/database/sqlite.rs @@ -0,0 +1,323 @@ +use std::path::{Path, PathBuf}; + +use nu_protocol::{CustomValue, ShellError, Span, Spanned, Value}; +use rusqlite::{types::ValueRef, Connection, Row}; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Serialize, Deserialize)] +pub struct SQLiteDatabase { + // I considered storing a SQLite connection here, but decided against it because + // 1) YAGNI, 2) it's not obvious how cloning a connection could work, 3) state + // management gets tricky quick. Revisit this approach if we find a compelling use case. + path: PathBuf, +} + +impl SQLiteDatabase { + pub fn new(path: &Path) -> SQLiteDatabase { + SQLiteDatabase { + path: PathBuf::from(path), + } + } + + pub fn query(&self, sql: &Spanned, call_span: Span) -> Result { + let db = open_sqlite_db(&self.path, call_span)?; + to_shell_error( + run_sql_query(db, sql), + "Failed to query SQLite database", + sql.span, + ) + } + + pub fn describe(&self) -> String { + format!("A SQLite database at {:?}", self.path) + } +} + +impl CustomValue for SQLiteDatabase { + fn clone_value(&self, span: Span) -> Value { + let cloned = SQLiteDatabase { + path: self.path.clone(), + }; + + Value::CustomValue { + val: Box::new(cloned), + span, + } + } + + fn value_string(&self) -> String { + self.typetag_name().to_string() + } + + fn to_base_value(&self, span: Span) -> Result { + let db = open_sqlite_db(&self.path, span)?; + to_shell_error( + read_entire_sqlite_db(db, span), + "Failed to read from SQLite database", + span, + ) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn follow_path_int(&self, _count: usize, span: Span) -> Result { + // In theory we could support this, but tables don't have an especially well-defined order + Err(ShellError::IncompatiblePathAccess("SQLite databases do not support integer-indexed access. Try specifying a table name instead".into(), span)) + } + + fn follow_path_string(&self, _column_name: String, span: Span) -> Result { + let db = open_sqlite_db(&self.path, span)?; + + to_shell_error( + read_single_table(db, _column_name, span), + "Failed to read from SQLite database", + span, + ) + } + + fn typetag_name(&self) -> &'static str { + "SQLiteDatabase" + } + + fn typetag_deserialize(&self) { + unimplemented!("typetag_deserialize") + } +} + +// TODO: is there a more elegant way to map rusqlite errors to ShellErrors? +fn to_shell_error( + result: Result, + message: &str, + span: Span, +) -> Result { + match result { + Ok(val) => Ok(val), + Err(err) => Err(ShellError::GenericError( + message.to_string(), + err.to_string(), + Some(span), + None, + Vec::new(), + )), + } +} + +fn open_sqlite_db(path: &Path, call_span: Span) -> Result { + let path = path.to_string_lossy().to_string(); + to_shell_error( + Connection::open(path), + "Failed to open SQLite database", + call_span, + ) +} + +fn read_entire_sqlite_db(conn: Connection, call_span: Span) -> Result { + let mut table_names: Vec = Vec::new(); + let mut tables: Vec = Vec::new(); + + let mut get_table_names = + conn.prepare("SELECT name FROM sqlite_master WHERE type = 'table'")?; + let rows = get_table_names.query_map([], |row| row.get(0))?; + + for row in rows { + let table_name: String = row?; + table_names.push(table_name.clone()); + + let mut rows = 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()? { + rows.push(convert_sqlite_row_to_nu_value(table_row, call_span)) + } + + let table_record = Value::List { + vals: rows, + span: call_span, + }; + + tables.push(table_record); + } + + Ok(Value::Record { + cols: table_names, + vals: tables, + span: call_span, + }) +} + +fn run_sql_query(conn: Connection, sql: &Spanned) -> Result { + let mut stmt = conn.prepare(&sql.item)?; + let mut results = stmt.query([])?; + let mut nu_records = Vec::new(); + + while let Some(table_row) = results.next()? { + nu_records.push(convert_sqlite_row_to_nu_value(table_row, sql.span)) + } + + Ok(Value::List { + vals: nu_records, + span: sql.span, + }) +} + +fn read_single_table( + conn: Connection, + table_name: String, + call_span: Span, +) -> Result { + let mut stmt = conn.prepare(&format!("SELECT * FROM {}", table_name))?; + let mut results = stmt.query([])?; + let mut nu_records = Vec::new(); + + while let Some(table_row) = results.next()? { + nu_records.push(convert_sqlite_row_to_nu_value(table_row, call_span)) + } + + Ok(Value::List { + vals: nu_records, + span: call_span, + }) +} + +fn convert_sqlite_row_to_nu_value(row: &Row, span: Span) -> Value { + let mut vals = Vec::new(); + let colnamestr = row.as_ref().column_names().to_vec(); + let colnames = colnamestr.iter().map(|s| s.to_string()).collect(); + + for (i, c) in row.as_ref().column_names().iter().enumerate() { + let _column = c.to_string(); + let val = convert_sqlite_value_to_nu_value(row.get_ref_unwrap(i), span); + vals.push(val); + } + + Value::Record { + cols: colnames, + vals, + span, + } +} + +fn convert_sqlite_value_to_nu_value(value: ValueRef, span: Span) -> Value { + match value { + ValueRef::Null => Value::Nothing { span }, + ValueRef::Integer(i) => Value::Int { val: i, span }, + ValueRef::Real(f) => Value::Float { val: f, span }, + ValueRef::Text(buf) => { + let s = match std::str::from_utf8(buf) { + Ok(v) => v, + Err(_) => { + return Value::Error { + error: ShellError::NonUtf8(span), + } + } + }; + Value::String { + val: s.to_string(), + span, + } + } + ValueRef::Blob(u) => Value::Binary { + val: u.to_vec(), + span, + }, + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn can_read_empty_db() { + let db = Connection::open_in_memory().unwrap(); + let converted_db = read_entire_sqlite_db(db, Span::test_data()).unwrap(); + + let expected = Value::Record { + cols: vec![], + vals: vec![], + span: Span::test_data(), + }; + + assert_eq!(converted_db, expected); + } + + #[test] + fn can_read_empty_table() { + let db = Connection::open_in_memory().unwrap(); + + db.execute( + "CREATE TABLE person ( + id INTEGER PRIMARY KEY, + name TEXT NOT NULL, + data BLOB + )", + [], + ) + .unwrap(); + let converted_db = read_entire_sqlite_db(db, Span::test_data()).unwrap(); + + let expected = Value::Record { + cols: vec!["person".to_string()], + vals: vec![Value::List { + vals: vec![], + span: Span::test_data(), + }], + span: Span::test_data(), + }; + + assert_eq!(converted_db, expected); + } + + #[test] + fn can_read_null_and_non_null_data() { + let span = Span::test_data(); + let db = Connection::open_in_memory().unwrap(); + + db.execute( + "CREATE TABLE item ( + id INTEGER PRIMARY KEY, + name TEXT + )", + [], + ) + .unwrap(); + + db.execute("INSERT INTO item (id, name) VALUES (123, NULL)", []) + .unwrap(); + + db.execute("INSERT INTO item (id, name) VALUES (456, 'foo bar')", []) + .unwrap(); + + let converted_db = read_entire_sqlite_db(db, span).unwrap(); + + let expected = Value::Record { + cols: vec!["item".to_string()], + vals: vec![Value::List { + vals: vec![ + Value::Record { + cols: vec!["id".to_string(), "name".to_string()], + vals: vec![Value::Int { val: 123, span }, Value::Nothing { span }], + span, + }, + Value::Record { + cols: vec!["id".to_string(), "name".to_string()], + vals: vec![ + Value::Int { val: 456, span }, + Value::String { + val: "foo bar".to_string(), + span, + }, + ], + span, + }, + ], + span, + }], + span, + }; + + assert_eq!(converted_db, expected); + } +} diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index 783a95ec6b..a47236767e 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -361,6 +361,11 @@ pub fn create_default_context(cwd: impl AsRef) -> EngineState { ViewSource, }; + // Database-related + bind_command! { + QueryDb + }; + // Deprecated bind_command! { PivotDeprecated, diff --git a/crates/nu-command/src/filesystem/open.rs b/crates/nu-command/src/filesystem/open.rs index e35ea5d367..874c5d2f76 100644 --- a/crates/nu-command/src/filesystem/open.rs +++ b/crates/nu-command/src/filesystem/open.rs @@ -1,13 +1,12 @@ use crate::filesystem::util::BufferedReader; +use crate::SQLiteDatabase; use nu_engine::{eval_block, get_full_help, CallExt}; use nu_protocol::ast::Call; use nu_protocol::engine::{Command, EngineState, Stack}; use nu_protocol::{ - Category, Example, IntoPipelineData, PipelineData, RawStream, ShellError, Signature, Span, - Spanned, SyntaxShape, Value, + Category, Example, IntoPipelineData, PipelineData, RawStream, ShellError, Signature, Spanned, + SyntaxShape, Value, }; -use rusqlite::types::ValueRef; -use rusqlite::{Connection, Row}; use std::io::{BufReader, Read, Seek}; #[cfg(unix)] @@ -121,8 +120,12 @@ impl Command for Open { let mut buf: [u8; 16] = [0; 16]; if file.read_exact(&mut buf).is_ok() && buf == sqlite_magic_bytes { - return open_and_read_sqlite_db(path, call_span) - .map(|val| PipelineData::Value(val, None)); + let custom_val = Value::CustomValue { + val: Box::new(SQLiteDatabase::new(path)), + span: call.head, + }; + + return Ok(custom_val.into_pipeline_data()); } if file.rewind().is_err() { @@ -196,208 +199,9 @@ impl Command for Open { } } -fn open_and_read_sqlite_db(path: &Path, call_span: Span) -> Result { - let path = path.to_string_lossy().to_string(); - - match Connection::open(path) { - Ok(conn) => match read_sqlite_db(conn, call_span) { - Ok(data) => Ok(data), - Err(err) => Err(ShellError::GenericError( - "Failed to read from SQLite database".into(), - err.to_string(), - Some(call_span), - None, - Vec::new(), - )), - }, - Err(err) => Err(ShellError::GenericError( - "Failed to open SQLite database".into(), - err.to_string(), - Some(call_span), - None, - Vec::new(), - )), - } -} - -fn read_sqlite_db(conn: Connection, call_span: Span) -> Result { - let mut table_names: Vec = Vec::new(); - let mut tables: Vec = Vec::new(); - - let mut get_table_names = - conn.prepare("SELECT name from sqlite_master where type = 'table'")?; - let rows = get_table_names.query_map([], |row| row.get(0))?; - - for row in rows { - let table_name: String = row?; - table_names.push(table_name.clone()); - - let mut rows = 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()? { - rows.push(convert_sqlite_row_to_nu_value(table_row, call_span)) - } - - let table_record = Value::List { - vals: rows, - span: call_span, - }; - - tables.push(table_record); - } - - Ok(Value::Record { - cols: table_names, - vals: tables, - span: call_span, - }) -} - -fn convert_sqlite_row_to_nu_value(row: &Row, span: Span) -> Value { - let mut vals = Vec::new(); - let colnamestr = row.as_ref().column_names().to_vec(); - let colnames = colnamestr.iter().map(|s| s.to_string()).collect(); - - for (i, c) in row.as_ref().column_names().iter().enumerate() { - let _column = c.to_string(); - let val = convert_sqlite_value_to_nu_value(row.get_ref_unwrap(i), span); - vals.push(val); - } - - Value::Record { - cols: colnames, - vals, - span, - } -} - -fn convert_sqlite_value_to_nu_value(value: ValueRef, span: Span) -> Value { - match value { - ValueRef::Null => Value::Nothing { span }, - ValueRef::Integer(i) => Value::Int { val: i, span }, - ValueRef::Real(f) => Value::Float { val: f, span }, - ValueRef::Text(buf) => { - let s = match std::str::from_utf8(buf) { - Ok(v) => v, - Err(_) => { - return Value::Error { - error: ShellError::NonUtf8(span), - } - } - }; - Value::String { - val: s.to_string(), - span, - } - } - ValueRef::Blob(u) => Value::Binary { - val: u.to_vec(), - span, - }, - } -} - fn permission_denied(dir: impl AsRef) -> bool { match dir.as_ref().read_dir() { Err(e) => matches!(e.kind(), std::io::ErrorKind::PermissionDenied), Ok(_) => false, } } - -#[cfg(test)] -mod test { - use super::*; - - #[test] - fn can_read_empty_db() { - let db = Connection::open_in_memory().unwrap(); - let converted_db = read_sqlite_db(db, Span::test_data()).unwrap(); - - let expected = Value::Record { - cols: vec![], - vals: vec![], - span: Span::test_data(), - }; - - assert_eq!(converted_db, expected); - } - - #[test] - fn can_read_empty_table() { - let db = Connection::open_in_memory().unwrap(); - - db.execute( - "CREATE TABLE person ( - id INTEGER PRIMARY KEY, - name TEXT NOT NULL, - data BLOB - )", - [], - ) - .unwrap(); - let converted_db = read_sqlite_db(db, Span::test_data()).unwrap(); - - let expected = Value::Record { - cols: vec!["person".to_string()], - vals: vec![Value::List { - vals: vec![], - span: Span::test_data(), - }], - span: Span::test_data(), - }; - - assert_eq!(converted_db, expected); - } - - #[test] - fn can_read_null_and_non_null_data() { - let span = Span::test_data(); - let db = Connection::open_in_memory().unwrap(); - - db.execute( - "CREATE TABLE item ( - id INTEGER PRIMARY KEY, - name TEXT - )", - [], - ) - .unwrap(); - - db.execute("INSERT INTO item (id, name) VALUES (123, NULL)", []) - .unwrap(); - - db.execute("INSERT INTO item (id, name) VALUES (456, 'foo bar')", []) - .unwrap(); - - let converted_db = read_sqlite_db(db, span).unwrap(); - - let expected = Value::Record { - cols: vec!["item".to_string()], - vals: vec![Value::List { - vals: vec![ - Value::Record { - cols: vec!["id".to_string(), "name".to_string()], - vals: vec![Value::Int { val: 123, span }, Value::Nothing { span }], - span, - }, - Value::Record { - cols: vec!["id".to_string(), "name".to_string()], - vals: vec![ - Value::Int { val: 456, span }, - Value::String { - val: "foo bar".to_string(), - span, - }, - ], - span, - }, - ], - span, - }], - span, - }; - - assert_eq!(converted_db, expected); - } -} diff --git a/crates/nu-command/src/filters/columns.rs b/crates/nu-command/src/filters/columns.rs index 6d3e45f780..2ad1b58cdf 100644 --- a/crates/nu-command/src/filters/columns.rs +++ b/crates/nu-command/src/filters/columns.rs @@ -73,6 +73,16 @@ fn getcol( .map(move |x| Value::String { val: x, span }) .into_pipeline_data(engine_state.ctrlc.clone())) } + PipelineData::Value(Value::CustomValue { val, span }, ..) => { + // TODO: should we get CustomValue to expose columns in a more efficient way? + // Would be nice to be able to get columns without generating the whole value + let input_as_base_value = val.to_base_value(span)?; + let input_cols = get_columns(&[input_as_base_value]); + Ok(input_cols + .into_iter() + .map(move |x| Value::String { val: x, span }) + .into_pipeline_data(engine_state.ctrlc.clone())) + } PipelineData::ListStream(stream, ..) => { let v: Vec<_> = stream.into_iter().collect(); let input_cols = get_columns(&v); diff --git a/crates/nu-command/src/lib.rs b/crates/nu-command/src/lib.rs index 29863951fb..2d4f8a6a23 100644 --- a/crates/nu-command/src/lib.rs +++ b/crates/nu-command/src/lib.rs @@ -1,5 +1,6 @@ mod conversions; mod core_commands; +mod database; mod date; mod default_context; mod deprecated; @@ -15,6 +16,7 @@ mod math; mod network; mod path; mod platform; +mod query; mod random; mod shells; mod strings; @@ -23,6 +25,7 @@ mod viewers; pub use conversions::*; pub use core_commands::*; +pub use database::*; pub use date::*; pub use default_context::*; pub use deprecated::*; @@ -39,6 +42,7 @@ pub use math::*; pub use network::*; pub use path::*; pub use platform::*; +pub use query::*; pub use random::*; pub use shells::*; pub use strings::*; diff --git a/crates/nu-command/src/query/db.rs b/crates/nu-command/src/query/db.rs new file mode 100644 index 0000000000..4f3e7afc8f --- /dev/null +++ b/crates/nu-command/src/query/db.rs @@ -0,0 +1,96 @@ +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, PipelineData, ShellError, Signature, Span, Spanned, SyntaxShape, Value, +}; + +use crate::database::SQLiteDatabase; + +#[derive(Clone)] +pub struct SubCommand; + +impl Command for SubCommand { + fn name(&self) -> &str { + "query db" + } + + fn signature(&self) -> Signature { + Signature::build("query db") + .required( + "query", + SyntaxShape::String, + "SQL to execute against the database", + ) + .category(Category::Date) // TODO: change category + } + + fn usage(&self) -> &str { + "Query a database using SQL." + } + + fn search_terms(&self) -> Vec<&str> { + vec!["database", "SQLite"] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + let sql: Spanned = call.req(engine_state, stack, 0)?; + let call_head = call.head; + + input.map( + move |value| query_input(value, call_head, &sql), + engine_state.ctrlc.clone(), + ) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Get 1 table out of a SQLite database", + example: r#"open foo.db | query db "SELECT * FROM Bar""#, + result: None, + }] + } +} + +fn query_input(input: Value, call_span: Span, sql: &Spanned) -> Value { + // this fn has slightly awkward error handling because it needs to jam errors into Value instead of returning a Result + if let Value::CustomValue { + val, + span: input_span, + } = input + { + let sqlite = val.as_any().downcast_ref::(); + + if let Some(db) = sqlite { + return match db.query(sql, call_span) { + Ok(val) => val, + Err(error) => Value::Error { error }, + }; + } + + return Value::Error { + error: ShellError::PipelineMismatch( + "a SQLite database".to_string(), + call_span, + input_span, + ), + }; + } + + match input.span() { + Ok(input_span) => Value::Error { + error: ShellError::PipelineMismatch( + "a SQLite database".to_string(), + call_span, + input_span, + ), + }, + Err(err) => Value::Error { error: err }, + } +} diff --git a/crates/nu-command/src/query/mod.rs b/crates/nu-command/src/query/mod.rs new file mode 100644 index 0000000000..26a4247dea --- /dev/null +++ b/crates/nu-command/src/query/mod.rs @@ -0,0 +1,3 @@ +mod db; + +pub use db::SubCommand as QueryDb; diff --git a/crates/nu-command/tests/commands/mod.rs b/crates/nu-command/tests/commands/mod.rs index ba76af6239..e7ff873a73 100644 --- a/crates/nu-command/tests/commands/mod.rs +++ b/crates/nu-command/tests/commands/mod.rs @@ -39,6 +39,7 @@ mod open; mod parse; mod path; mod prepend; +mod query; mod random; mod range; mod reduce; diff --git a/crates/nu-command/tests/commands/open.rs b/crates/nu-command/tests/commands/open.rs index d796addf89..cbb6ebe3ba 100644 --- a/crates/nu-command/tests/commands/open.rs +++ b/crates/nu-command/tests/commands/open.rs @@ -109,7 +109,22 @@ fn parses_more_bson_complexity() { // ╰───┴──────╯ #[test] + fn parses_sqlite() { + let actual = nu!( + cwd: "tests/fixtures/formats", pipeline( + r#" + open sample.db + | columns + | length + "# + )); + + assert_eq!(actual.out, "3"); +} + +#[test] +fn parses_sqlite_get_column_name() { let actual = nu!( cwd: "tests/fixtures/formats", pipeline( r#" diff --git a/crates/nu-command/tests/commands/query/db.rs b/crates/nu-command/tests/commands/query/db.rs new file mode 100644 index 0000000000..a7a854ca13 --- /dev/null +++ b/crates/nu-command/tests/commands/query/db.rs @@ -0,0 +1,47 @@ +use nu_test_support::{nu, pipeline}; + +#[test] + +fn can_query_single_table() { + let actual = nu!( + cwd: "tests/fixtures/formats", pipeline( + r#" + open sample.db + | query db "select * from strings" + | where x =~ ell + | length + "# + )); + + assert_eq!(actual.out, "4"); +} + +#[test] +fn invalid_sql_fails() { + let actual = nu!( + cwd: "tests/fixtures/formats", pipeline( + r#" + open sample.db + | query db "select *asdfasdf" + "# + )); + + assert!(actual.err.contains("syntax error")); +} + +#[test] +fn invalid_input_fails() { + let actual = nu!( + cwd: "tests/fixtures/formats", pipeline( + r#" + "foo" | query db "select * from asdf" + "# + )); + + assert!(actual.err.contains("pipeline_mismatch")); +} + +#[test] +fn scratch() { + assert!(true); +} diff --git a/crates/nu-command/tests/commands/query/mod.rs b/crates/nu-command/tests/commands/query/mod.rs new file mode 100644 index 0000000000..b91801835d --- /dev/null +++ b/crates/nu-command/tests/commands/query/mod.rs @@ -0,0 +1 @@ +mod db; diff --git a/crates/nu-protocol/src/value/custom_value.rs b/crates/nu-protocol/src/value/custom_value.rs index 35b4d34f84..25470dd900 100644 --- a/crates/nu-protocol/src/value/custom_value.rs +++ b/crates/nu-protocol/src/value/custom_value.rs @@ -28,14 +28,14 @@ pub trait CustomValue: fmt::Debug + Send + Sync { // Follow cell path functions fn follow_path_int(&self, _count: usize, span: Span) -> Result { Err(ShellError::IncompatiblePathAccess( - format!("{} does't support path access", self.value_string()), + format!("{} doesn't support path access", self.value_string()), span, )) } fn follow_path_string(&self, _column_name: String, span: Span) -> Result { Err(ShellError::IncompatiblePathAccess( - format!("{} does't support path access", self.value_string()), + format!("{} doesn't support path access", self.value_string()), span, )) }