From 1cb449b2d14bcd77305024716db13a7dbecaede7 Mon Sep 17 00:00:00 2001 From: Fernando Herrera Date: Sat, 7 May 2022 13:33:33 +0100 Subject: [PATCH] Database commands (#5466) * change query to statement * internal functions and over definitions * cargo fmt --- .../{expressions => commands}/alias.rs | 72 +++++++----- .../nu-command/src/database/commands/and.rs | 55 +++++---- .../database/{expressions => commands}/col.rs | 2 +- .../nu-command/src/database/commands/from.rs | 52 ++++++--- .../src/database/commands/function.rs | 85 ++++++++++++++ .../src/database/commands/group_by.rs | 102 ++++++++++++++++ .../nu-command/src/database/commands/limit.rs | 19 ++- .../nu-command/src/database/commands/mod.rs | 17 ++- crates/nu-command/src/database/commands/or.rs | 55 +++++---- .../src/database/commands/order_by.rs | 110 ++++++++++++++---- .../nu-command/src/database/commands/over.rs | 80 +++++++++++++ .../src/database/commands/select.rs | 51 +++++--- .../src/database/commands/where_.rs | 36 +++--- .../src/database/expressions/mod.rs | 21 ---- crates/nu-command/src/database/mod.rs | 10 +- .../src/database/values/dsl/expression.rs | 43 ++++++- .../nu-command/src/database/values/sqlite.rs | 18 +-- 17 files changed, 630 insertions(+), 198 deletions(-) rename crates/nu-command/src/database/{expressions => commands}/alias.rs (62%) rename crates/nu-command/src/database/{expressions => commands}/col.rs (97%) create mode 100644 crates/nu-command/src/database/commands/function.rs create mode 100644 crates/nu-command/src/database/commands/group_by.rs create mode 100644 crates/nu-command/src/database/commands/over.rs delete mode 100644 crates/nu-command/src/database/expressions/mod.rs diff --git a/crates/nu-command/src/database/expressions/alias.rs b/crates/nu-command/src/database/commands/alias.rs similarity index 62% rename from crates/nu-command/src/database/expressions/alias.rs rename to crates/nu-command/src/database/commands/alias.rs index 9cda68a0ba..407a0341fb 100644 --- a/crates/nu-command/src/database/expressions/alias.rs +++ b/crates/nu-command/src/database/commands/alias.rs @@ -8,7 +8,7 @@ use nu_protocol::{ engine::{Command, EngineState, Stack}, Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, SyntaxShape, }; -use sqlparser::ast::{Ident, SelectItem, SetExpr, TableAlias, TableFactor}; +use sqlparser::ast::{Ident, SelectItem, SetExpr, Statement, TableAlias, TableFactor}; #[derive(Clone)] pub struct AliasExpr; @@ -110,44 +110,56 @@ fn alias_db( new_alias: String, call: &Call, ) -> Result { - match db.query { + match db.statement.as_mut() { None => Err(ShellError::GenericError( "Error creating alias".into(), - "there is no query defined yet".into(), + "there is no statement defined yet".into(), Some(call.head), None, Vec::new(), )), - Some(ref mut query) => match &mut query.body { - SetExpr::Select(ref mut select) => { - select.as_mut().from.iter_mut().for_each(|table| { - let new_alias = Some(TableAlias { - name: Ident { - value: new_alias.clone(), - quote_style: None, - }, - columns: Vec::new(), + Some(statement) => match statement { + Statement::Query(query) => match &mut query.body { + SetExpr::Select(select) => { + select.as_mut().from.iter_mut().for_each(|table| { + let new_alias = Some(TableAlias { + name: Ident { + value: new_alias.clone(), + quote_style: None, + }, + columns: Vec::new(), + }); + + if let TableFactor::Table { ref mut alias, .. } = table.relation { + *alias = new_alias; + } else if let TableFactor::Derived { ref mut alias, .. } = table.relation { + *alias = new_alias; + } else if let TableFactor::TableFunction { ref mut alias, .. } = + table.relation + { + *alias = new_alias; + } }); - if let TableFactor::Table { ref mut alias, .. } = table.relation { - *alias = new_alias; - } else if let TableFactor::Derived { ref mut alias, .. } = table.relation { - *alias = new_alias; - } else if let TableFactor::TableFunction { ref mut alias, .. } = table.relation - { - *alias = new_alias; - } - }); - - Ok(db.into_value(call.head).into_pipeline_data()) + Ok(db.into_value(call.head).into_pipeline_data()) + } + _ => Err(ShellError::GenericError( + "Error creating alias".into(), + "Query has no select from defined".into(), + Some(call.head), + None, + Vec::new(), + )), + }, + s => { + return Err(ShellError::GenericError( + "Connection doesnt define a query".into(), + format!("Expected a connection with query. Got {}", s), + Some(call.head), + None, + Vec::new(), + )) } - _ => Err(ShellError::GenericError( - "Error creating alias".into(), - "Query has no select from defined".into(), - Some(call.head), - None, - Vec::new(), - )), }, } } diff --git a/crates/nu-command/src/database/commands/and.rs b/crates/nu-command/src/database/commands/and.rs index 45cdbc6094..8e4d6dc607 100644 --- a/crates/nu-command/src/database/commands/and.rs +++ b/crates/nu-command/src/database/commands/and.rs @@ -8,7 +8,7 @@ use nu_protocol::{ Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, Span, SyntaxShape, Value, }; -use sqlparser::ast::{BinaryOperator, Expr, Query, Select, SetExpr}; +use sqlparser::ast::{BinaryOperator, Expr, Query, Select, SetExpr, Statement}; #[derive(Clone)] pub struct AndDb; @@ -78,12 +78,23 @@ impl Command for AndDb { Ok(expression.into_value(call.head).into_pipeline_data()) } else if let Ok(mut db) = SQLiteDatabase::try_from_value(value.clone()) { - db.query = match db.query { - Some(query) => Some(modify_query(query, expr, call.head)?), + match db.statement.as_mut() { + Some(statement) => match statement { + Statement::Query(query) => modify_query(query, expr, call.head)?, + s => { + return Err(ShellError::GenericError( + "Connection doesnt define a query".into(), + format!("Expected a connection with query. Got {}", s), + Some(call.head), + None, + Vec::new(), + )) + } + }, None => { return Err(ShellError::GenericError( - "Connection without query".into(), - "Missing query in the connection".into(), + "Connection without statement".into(), + "The connection needs a statement defined".into(), Some(call.head), None, Vec::new(), @@ -103,26 +114,24 @@ impl Command for AndDb { } } -fn modify_query(mut query: Query, expression: Expr, span: Span) -> Result { - query.body = match query.body { - SetExpr::Select(select) => Ok(SetExpr::Select(modify_select(select, expression, span)?)), - _ => Err(ShellError::GenericError( - "Query without a select".into(), - "Missing a WHERE clause before an AND clause".into(), - Some(span), - None, - Vec::new(), - )), - }?; +fn modify_query(query: &mut Box, expression: Expr, span: Span) -> Result<(), ShellError> { + match query.body { + SetExpr::Select(ref mut select) => modify_select(select, expression, span)?, + _ => { + return Err(ShellError::GenericError( + "Query without a select".into(), + "Missing a WHERE clause before an AND clause".into(), + Some(span), + None, + Vec::new(), + )) + } + }; - Ok(query) + Ok(()) } -fn modify_select( - mut select: Box, expression: Expr, span: Span) -> Result<(), ShellError> { let new_expression = match &select.selection { Some(expr) => Ok(Expr::BinaryOp { left: Box::new(expr.clone()), @@ -139,5 +148,5 @@ fn modify_select( }?; select.as_mut().selection = Some(new_expression); - Ok(select) + Ok(()) } diff --git a/crates/nu-command/src/database/expressions/col.rs b/crates/nu-command/src/database/commands/col.rs similarity index 97% rename from crates/nu-command/src/database/expressions/col.rs rename to crates/nu-command/src/database/commands/col.rs index f604d26b69..c94c294b93 100644 --- a/crates/nu-command/src/database/expressions/col.rs +++ b/crates/nu-command/src/database/commands/col.rs @@ -27,7 +27,7 @@ impl Command for ColExpr { fn examples(&self) -> Vec { vec![Example { description: "Creates a named column expression", - example: "col name_1", + example: "db col name_1", result: None, }] } diff --git a/crates/nu-command/src/database/commands/from.rs b/crates/nu-command/src/database/commands/from.rs index 4c9398eaa4..fd7692ae41 100644 --- a/crates/nu-command/src/database/commands/from.rs +++ b/crates/nu-command/src/database/commands/from.rs @@ -3,9 +3,11 @@ use nu_engine::CallExt; use nu_protocol::{ ast::Call, engine::{Command, EngineState, Stack}, - Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, SyntaxShape, + Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, Span, SyntaxShape, +}; +use sqlparser::ast::{ + Ident, ObjectName, Query, Select, SetExpr, Statement, TableFactor, TableWithJoins, }; -use sqlparser::ast::{Ident, ObjectName, Query, Select, SetExpr, TableFactor, TableWithJoins}; #[derive(Clone)] pub struct FromDb; @@ -51,17 +53,17 @@ impl Command for FromDb { let table: String = call.req(engine_state, stack, 0)?; let mut db = SQLiteDatabase::try_from_pipeline(input, call.head)?; - db.query = match db.query { - None => Some(create_query(table)), - Some(query) => Some(modify_query(query, table)), + db.statement = match db.statement { + None => Some(create_statement(table)), + Some(statement) => Some(modify_statement(statement, table, call.head)?), }; Ok(db.into_value(call.head).into_pipeline_data()) } } -fn create_query(table: String) -> Query { - Query { +fn create_statement(table: String) -> Statement { + let query = Query { with: None, body: SetExpr::Select(Box::new(create_select(table))), order_by: Vec::new(), @@ -69,21 +71,35 @@ fn create_query(table: String) -> Query { offset: None, fetch: None, lock: None, - } -} - -fn modify_query(mut query: Query, table: String) -> Query { - query.body = match query.body { - SetExpr::Select(select) => SetExpr::Select(modify_select(select, table)), - _ => SetExpr::Select(Box::new(create_select(table))), }; - query + Statement::Query(Box::new(query)) } -fn modify_select(mut select: Box { - select.as_mut().from = create_from(table); - select +fn modify_statement( + mut statement: Statement, + table: String, + span: Span, +) -> Result { + match statement { + Statement::Query(ref mut query) => { + match query.body { + SetExpr::Select(ref mut select) => select.as_mut().from = create_from(table), + _ => { + query.as_mut().body = SetExpr::Select(Box::new(create_select(table))); + } + }; + + Ok(statement) + } + s => Err(ShellError::GenericError( + "Connection doesnt define a statement".into(), + format!("Expected a connection with query. Got {}", s), + Some(span), + None, + Vec::new(), + )), + } } fn create_select(table: String) -> Select { diff --git a/crates/nu-command/src/database/commands/function.rs b/crates/nu-command/src/database/commands/function.rs new file mode 100644 index 0000000000..50e36c0191 --- /dev/null +++ b/crates/nu-command/src/database/commands/function.rs @@ -0,0 +1,85 @@ +use crate::database::values::dsl::ExprDb; +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, SyntaxShape, Value, +}; +use sqlparser::ast::{Expr, Function, FunctionArg, FunctionArgExpr, Ident, ObjectName}; + +#[derive(Clone)] +pub struct FunctionExpr; + +impl Command for FunctionExpr { + fn name(&self) -> &str { + "db fn" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required("name", SyntaxShape::String, "function name") + .switch("distinct", "distict values", Some('d')) + .rest("arguments", SyntaxShape::Any, "function arguments") + .category(Category::Custom("database".into())) + } + + fn usage(&self) -> &str { + "Creates function expression for a select operation" + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Creates a function expression", + example: "db fn count name_1", + result: None, + }] + } + + fn search_terms(&self) -> Vec<&str> { + vec!["database", "function", "expression"] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + _input: PipelineData, + ) -> Result { + let name: String = call.req(engine_state, stack, 0)?; + let vals: Vec = call.rest(engine_state, stack, 1)?; + let value = Value::List { + vals, + span: call.head, + }; + let expressions = ExprDb::extract_exprs(value)?; + + let name: Vec = name + .split('.') + .map(|part| Ident { + value: part.to_string(), + quote_style: None, + }) + .collect(); + let name = ObjectName(name); + + let args: Vec = expressions + .into_iter() + .map(|expr| { + let arg = FunctionArgExpr::Expr(expr); + + FunctionArg::Unnamed(arg) + }) + .collect(); + + let expression: ExprDb = Expr::Function(Function { + name, + args, + over: None, + distinct: call.has_flag("distinct"), + }) + .into(); + + Ok(expression.into_value(call.head).into_pipeline_data()) + } +} diff --git a/crates/nu-command/src/database/commands/group_by.rs b/crates/nu-command/src/database/commands/group_by.rs new file mode 100644 index 0000000000..ceb2ed3856 --- /dev/null +++ b/crates/nu-command/src/database/commands/group_by.rs @@ -0,0 +1,102 @@ +use crate::database::values::dsl::ExprDb; + +use super::super::SQLiteDatabase; +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, SyntaxShape, Value, +}; +use sqlparser::ast::{SetExpr, Statement}; + +#[derive(Clone)] +pub struct GroupByDb; + +impl Command for GroupByDb { + fn name(&self) -> &str { + "db group-by" + } + + fn usage(&self) -> &str { + "Group by query" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .rest( + "select", + SyntaxShape::Any, + "Select expression(s) on the table", + ) + .category(Category::Custom("database".into())) + } + + fn search_terms(&self) -> Vec<&str> { + vec!["database", "select"] + } + + fn examples(&self) -> Vec { + vec![Example { + description: "orders query by a column", + example: r#"db open db.mysql + | db from table_a + | db select a + | db group-by a + | db describe"#, + result: None, + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + let vals: Vec = call.rest(engine_state, stack, 0)?; + let value = Value::List { + vals, + span: call.head, + }; + let expressions = ExprDb::extract_exprs(value)?; + + let mut db = SQLiteDatabase::try_from_pipeline(input, call.head)?; + match db.statement.as_mut() { + Some(statement) => match statement { + Statement::Query(ref mut query) => match &mut query.body { + SetExpr::Select(ref mut select) => select.group_by = expressions, + s => { + return Err(ShellError::GenericError( + "Connection doesnt define a select".into(), + format!("Expected a connection with select query. Got {}", s), + Some(call.head), + None, + Vec::new(), + )) + } + }, + s => { + return Err(ShellError::GenericError( + "Connection doesnt define a query".into(), + format!("Expected a connection with query. Got {}", s), + Some(call.head), + None, + Vec::new(), + )) + } + }, + None => { + return Err(ShellError::GenericError( + "Connection without statement".into(), + "The connection needs a statement defined".into(), + Some(call.head), + None, + Vec::new(), + )) + } + }; + + Ok(db.into_value(call.head).into_pipeline_data()) + } +} diff --git a/crates/nu-command/src/database/commands/limit.rs b/crates/nu-command/src/database/commands/limit.rs index d2b7f14558..d46b7c48c0 100644 --- a/crates/nu-command/src/database/commands/limit.rs +++ b/crates/nu-command/src/database/commands/limit.rs @@ -6,6 +6,7 @@ use nu_protocol::{ engine::{Command, EngineState, Stack}, Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, SyntaxShape, Value, }; +use sqlparser::ast::Statement; #[derive(Clone)] pub struct LimitDb; @@ -56,11 +57,19 @@ impl Command for LimitDb { let expr = ExprDb::try_from_value(&limit)?.into_native(); let mut db = SQLiteDatabase::try_from_pipeline(input, call.head)?; - db.query = match db.query { - Some(mut query) => { - query.limit = Some(expr); - Some(query) - } + match db.statement { + Some(ref mut statement) => match statement { + Statement::Query(query) => query.as_mut().limit = Some(expr), + s => { + return Err(ShellError::GenericError( + "Connection doesnt define a statement".into(), + format!("Expected a connection with query. Got {}", s), + Some(call.head), + None, + Vec::new(), + )) + } + }, None => { return Err(ShellError::GenericError( "Connection without query".into(), diff --git a/crates/nu-command/src/database/commands/mod.rs b/crates/nu-command/src/database/commands/mod.rs index 7906918319..12a2b7aea5 100644 --- a/crates/nu-command/src/database/commands/mod.rs +++ b/crates/nu-command/src/database/commands/mod.rs @@ -1,12 +1,17 @@ +mod alias; mod and; +mod col; mod collect; mod command; mod describe; mod from; +mod function; +mod group_by; mod limit; mod open; mod or; mod order_by; +mod over; mod query; mod schema; mod select; @@ -18,21 +23,26 @@ use testing::TestingDb; use nu_protocol::engine::StateWorkingSet; +use alias::AliasExpr; use and::AndDb; +use col::ColExpr; use collect::CollectDb; use command::Database; use describe::DescribeDb; use from::FromDb; +use function::FunctionExpr; +use group_by::GroupByDb; use limit::LimitDb; use open::OpenDb; use or::OrDb; use order_by::OrderByDb; +use over::OverExpr; use query::QueryDb; use schema::SchemaDb; use select::ProjectionDb; use where_::WhereDb; -pub fn add_commands_decls(working_set: &mut StateWorkingSet) { +pub fn add_database_decls(working_set: &mut StateWorkingSet) { macro_rules! bind_command { ( $command:expr ) => { working_set.add_decl(Box::new($command)); @@ -44,17 +54,22 @@ pub fn add_commands_decls(working_set: &mut StateWorkingSet) { // Series commands bind_command!( + AliasExpr, AndDb, CollectDb, + ColExpr, Database, DescribeDb, FromDb, + FunctionExpr, + GroupByDb, QueryDb, LimitDb, ProjectionDb, OpenDb, OrderByDb, OrDb, + OverExpr, SchemaDb, TestingDb, WhereDb diff --git a/crates/nu-command/src/database/commands/or.rs b/crates/nu-command/src/database/commands/or.rs index 62ee26bbe4..4593421939 100644 --- a/crates/nu-command/src/database/commands/or.rs +++ b/crates/nu-command/src/database/commands/or.rs @@ -8,7 +8,7 @@ use nu_protocol::{ Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, Span, SyntaxShape, Value, }; -use sqlparser::ast::{BinaryOperator, Expr, Query, Select, SetExpr}; +use sqlparser::ast::{BinaryOperator, Expr, Query, Select, SetExpr, Statement}; #[derive(Clone)] pub struct OrDb; @@ -78,12 +78,23 @@ impl Command for OrDb { Ok(expression.into_value(call.head).into_pipeline_data()) } else if let Ok(mut db) = SQLiteDatabase::try_from_value(value.clone()) { - db.query = match db.query { - Some(query) => Some(modify_query(query, expr, call.head)?), + match db.statement { + Some(ref mut statement) => match statement { + Statement::Query(query) => modify_query(query, expr, call.head)?, + s => { + return Err(ShellError::GenericError( + "Connection doesnt define a query".into(), + format!("Expected a connection with query. Got {}", s), + Some(call.head), + None, + Vec::new(), + )) + } + }, None => { return Err(ShellError::GenericError( - "Connection without query".into(), - "Missing query in the connection".into(), + "Connection without statement".into(), + "The connection needs a statement defined".into(), Some(call.head), None, Vec::new(), @@ -103,26 +114,24 @@ impl Command for OrDb { } } -fn modify_query(mut query: Query, expression: Expr, span: Span) -> Result { - query.body = match query.body { - SetExpr::Select(select) => Ok(SetExpr::Select(modify_select(select, expression, span)?)), - _ => Err(ShellError::GenericError( - "Query without a select".into(), - "Missing a WHERE clause before an OR clause".into(), - Some(span), - None, - Vec::new(), - )), - }?; +fn modify_query(query: &mut Box, expression: Expr, span: Span) -> Result<(), ShellError> { + match query.body { + SetExpr::Select(ref mut select) => modify_select(select, expression, span)?, + _ => { + return Err(ShellError::GenericError( + "Query without a select".into(), + "Missing a WHERE clause before an OR clause".into(), + Some(span), + None, + Vec::new(), + )) + } + }; - Ok(query) + Ok(()) } -fn modify_select( - mut select: Box, expression: Expr, span: Span) -> Result<(), ShellError> { let new_expression = match &select.selection { Some(expr) => Ok(Expr::BinaryOp { left: Box::new(expr.clone()), @@ -139,5 +148,5 @@ fn modify_select( }?; select.as_mut().selection = Some(new_expression); - Ok(select) + Ok(()) } diff --git a/crates/nu-command/src/database/commands/order_by.rs b/crates/nu-command/src/database/commands/order_by.rs index 7ce65ceae7..6d8a726ede 100644 --- a/crates/nu-command/src/database/commands/order_by.rs +++ b/crates/nu-command/src/database/commands/order_by.rs @@ -7,7 +7,7 @@ use nu_protocol::{ engine::{Command, EngineState, Stack}, Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, SyntaxShape, Value, }; -use sqlparser::ast::OrderByExpr; +use sqlparser::ast::{Expr, OrderByExpr, Statement}; #[derive(Clone)] pub struct OrderByDb; @@ -58,40 +58,100 @@ impl Command for OrderByDb { ) -> Result { let asc = call.has_flag("ascending"); let nulls_first = call.has_flag("nulls_first"); - - let vals: Vec = call.rest(engine_state, stack, 0)?; - let value = Value::List { - vals, + let expressions: Vec = call.rest(engine_state, stack, 0)?; + let expressions = Value::List { + vals: expressions, span: call.head, }; - let expressions = ExprDb::extract_exprs(value)?; + let expressions = ExprDb::extract_exprs(expressions)?; + let expressions: Vec = expressions + .into_iter() + .map(|expr| OrderByExpr { + expr, + asc: if asc { Some(asc) } else { None }, + nulls_first: if nulls_first { Some(nulls_first) } else { None }, + }) + .collect(); - let mut db = SQLiteDatabase::try_from_pipeline(input, call.head)?; - db.query = match db.query { - Some(mut query) => { - let mut order_expr: Vec = expressions - .into_iter() - .map(|expr| OrderByExpr { - expr, - asc: if asc { Some(asc) } else { None }, - nulls_first: if nulls_first { Some(nulls_first) } else { None }, - }) - .collect(); + let value = input.into_value(call.head); - query.order_by.append(&mut order_expr); - Some(query) - } + if let Ok(expr) = ExprDb::try_from_value(&value) { + update_expressions(expr, expressions, call) + } else if let Ok(db) = SQLiteDatabase::try_from_value(value.clone()) { + update_connection(db, expressions, call) + } else { + Err(ShellError::CantConvert( + "expression or query".into(), + value.get_type().to_string(), + value.span()?, + None, + )) + } + } +} + +fn update_expressions( + mut expr: ExprDb, + mut expressions: Vec, + call: &Call, +) -> Result { + match expr.as_mut() { + Expr::Function(function) => match &mut function.over { + Some(over) => over.order_by.append(&mut expressions), None => { return Err(ShellError::GenericError( - "Connection without query".into(), - "The connection needs a query defined".into(), + "Expression doesnt define a partition to order".into(), + "Expected an expression with partition".into(), Some(call.head), None, Vec::new(), )) } - }; + }, + s => { + return Err(ShellError::GenericError( + "Expression doesnt define a function".into(), + format!("Expected an expression with a function. Got {}", s), + Some(call.head), + None, + Vec::new(), + )) + } + }; - Ok(db.into_value(call.head).into_pipeline_data()) - } + Ok(expr.into_value(call.head).into_pipeline_data()) +} + +fn update_connection( + mut db: SQLiteDatabase, + mut expressions: Vec, + call: &Call, +) -> Result { + match db.statement.as_mut() { + Some(statement) => match statement { + Statement::Query(query) => { + query.order_by.append(&mut expressions); + } + s => { + return Err(ShellError::GenericError( + "Connection doesnt define a query".into(), + format!("Expected a connection with query. Got {}", s), + Some(call.head), + None, + Vec::new(), + )) + } + }, + None => { + return Err(ShellError::GenericError( + "Connection without statement".into(), + "The connection needs a statement defined".into(), + Some(call.head), + None, + Vec::new(), + )) + } + }; + + Ok(db.into_value(call.head).into_pipeline_data()) } diff --git a/crates/nu-command/src/database/commands/over.rs b/crates/nu-command/src/database/commands/over.rs new file mode 100644 index 0000000000..b47f2b5db1 --- /dev/null +++ b/crates/nu-command/src/database/commands/over.rs @@ -0,0 +1,80 @@ +use crate::database::values::dsl::ExprDb; +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, SyntaxShape, Value, +}; +use sqlparser::ast::{Expr, WindowSpec}; + +#[derive(Clone)] +pub struct OverExpr; + +impl Command for OverExpr { + fn name(&self) -> &str { + "db over" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .rest( + "partition-by", + SyntaxShape::Any, + "columns to partition the window function", + ) + .category(Category::Custom("database".into())) + } + + fn usage(&self) -> &str { + "Adds a partition to an expression function" + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Adds a partition to a function expresssion", + example: "db function avg col_a | db over col_b", + result: None, + }] + } + + fn search_terms(&self) -> Vec<&str> { + vec!["database", "column", "expression"] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + let vals: Vec = call.rest(engine_state, stack, 0)?; + let value = Value::List { + vals, + span: call.head, + }; + let partitions = ExprDb::extract_exprs(value)?; + + let mut expression = ExprDb::try_from_pipeline(input, call.head)?; + match expression.as_mut() { + Expr::Function(function) => { + function.over = Some(WindowSpec { + partition_by: partitions, + order_by: Vec::new(), + window_frame: None, + }); + } + s => { + return Err(ShellError::GenericError( + "Expression doesnt define a function".into(), + format!("Expected an expression with a function. Got {}", s), + Some(call.head), + None, + Vec::new(), + )) + } + }; + + Ok(expression.into_value(call.head).into_pipeline_data()) + } +} diff --git a/crates/nu-command/src/database/commands/select.rs b/crates/nu-command/src/database/commands/select.rs index 78a366a8b1..e550b510b9 100644 --- a/crates/nu-command/src/database/commands/select.rs +++ b/crates/nu-command/src/database/commands/select.rs @@ -3,9 +3,10 @@ use nu_engine::CallExt; use nu_protocol::{ ast::Call, engine::{Command, EngineState, Stack}, - Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, SyntaxShape, Value, + Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, Span, SyntaxShape, + Value, }; -use sqlparser::ast::{Query, Select, SelectItem, SetExpr}; +use sqlparser::ast::{Query, Select, SelectItem, SetExpr, Statement}; #[derive(Clone)] pub struct ProjectionDb; @@ -63,17 +64,17 @@ impl Command for ProjectionDb { let projection = SelectDb::extract_selects(value)?; let mut db = SQLiteDatabase::try_from_pipeline(input, call.head)?; - db.query = match db.query { - None => Some(create_query(projection)), - Some(query) => Some(modify_query(query, projection)), + db.statement = match db.statement { + None => Some(create_statement(projection)), + Some(statement) => Some(modify_statement(statement, projection, call.head)?), }; Ok(db.into_value(call.head).into_pipeline_data()) } } -fn create_query(expressions: Vec) -> Query { - Query { +fn create_statement(expressions: Vec) -> Statement { + let query = Query { with: None, body: SetExpr::Select(Box::new(create_select(expressions))), order_by: Vec::new(), @@ -81,21 +82,35 @@ fn create_query(expressions: Vec) -> Query { offset: None, fetch: None, lock: None, - } -} - -fn modify_query(mut query: Query, expressions: Vec) -> Query { - query.body = match query.body { - SetExpr::Select(select) => SetExpr::Select(modify_select(select, expressions)), - _ => SetExpr::Select(Box::new(create_select(expressions))), }; - query + Statement::Query(Box::new(query)) } -fn modify_select(mut select: Box { - select.as_mut().projection = projection; - select +fn modify_statement( + mut statement: Statement, + expressions: Vec, + span: Span, +) -> Result { + match statement { + Statement::Query(ref mut query) => { + match query.body { + SetExpr::Select(ref mut select) => select.as_mut().projection = expressions, + _ => { + query.as_mut().body = SetExpr::Select(Box::new(create_select(expressions))); + } + }; + + Ok(statement) + } + s => Err(ShellError::GenericError( + "Connection doesnt define a statement".into(), + format!("Expected a connection with query. Got {}", s), + Some(span), + None, + Vec::new(), + )), + } } fn create_select(projection: Vec) -> Select { diff --git a/crates/nu-command/src/database/commands/where_.rs b/crates/nu-command/src/database/commands/where_.rs index 47e6f9ed1e..f14a22afa0 100644 --- a/crates/nu-command/src/database/commands/where_.rs +++ b/crates/nu-command/src/database/commands/where_.rs @@ -7,7 +7,7 @@ use nu_protocol::{ engine::{Command, EngineState, Stack}, Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, SyntaxShape, Value, }; -use sqlparser::ast::{Expr, Query, Select, SetExpr}; +use sqlparser::ast::{Expr, Query, Select, SetExpr, Statement}; #[derive(Clone)] pub struct WhereDb; @@ -54,12 +54,23 @@ impl Command for WhereDb { let expr = ExprDb::try_from_value(&value)?.into_native(); let mut db = SQLiteDatabase::try_from_pipeline(input, call.head)?; - db.query = match db.query { - Some(query) => Some(modify_query(query, expr)), + match db.statement.as_mut() { + Some(statement) => match statement { + Statement::Query(query) => modify_query(query, expr), + s => { + return Err(ShellError::GenericError( + "Connection doesnt define a query".into(), + format!("Expected a connection with query. Got {}", s), + Some(call.head), + None, + Vec::new(), + )) + } + }, None => { return Err(ShellError::GenericError( - "Connection without query".into(), - "The connection needs a query defined".into(), + "Connection without statement".into(), + "The connection needs a statement defined".into(), Some(call.head), None, Vec::new(), @@ -71,18 +82,17 @@ impl Command for WhereDb { } } -fn modify_query(mut query: Query, expression: Expr) -> Query { - query.body = match query.body { - SetExpr::Select(select) => SetExpr::Select(modify_select(select, expression)), - _ => SetExpr::Select(Box::new(create_select(expression))), +fn modify_query(query: &mut Box, expression: Expr) { + match query.body { + SetExpr::Select(ref mut select) => modify_select(select, expression), + _ => { + query.as_mut().body = SetExpr::Select(Box::new(create_select(expression))); + } }; - - query } -fn modify_select(mut select: Box { +fn modify_select(select: &mut Box