diff --git a/integration/rust/tests/integration/subqueries/mod.rs b/integration/rust/tests/integration/subqueries/mod.rs index 392304113..ef93e7260 100644 --- a/integration/rust/tests/integration/subqueries/mod.rs +++ b/integration/rust/tests/integration/subqueries/mod.rs @@ -180,6 +180,7 @@ async fn test_subqueries_suite() { let (mut regular_cases, mut data_modifying_cases) = (Vec::new(), Vec::new()); // Run a single case by name, e.g. SUBQUERY_TEST= let test_target: Option = std::env::var("SUBQUERY_TEST").ok(); + println!("Using test target: {test_target:?}"); // Load all cases from the 'tests' folder let list_of_files: fs::ReadDir = fs::read_dir("tests/integration/subqueries/tests").unwrap(); diff --git a/integration/rust/tests/integration/subqueries/test_case.rs b/integration/rust/tests/integration/subqueries/test_case.rs index 2f9d9a0f7..eb23489ae 100644 --- a/integration/rust/tests/integration/subqueries/test_case.rs +++ b/integration/rust/tests/integration/subqueries/test_case.rs @@ -43,19 +43,16 @@ pub async fn run_tests_for_case( break; }; - if execute_test_case( + let result = execute_test_case( &job, worker_name.as_str(), &mut direct_conn, &mut sharded_conn, ) - .await - .is_ok() - { - success += 1; - } else { - println!("FAILED HERE: {} {}", worker_name.as_str(), job.sql); - println!("FAILED HERE!"); + .await; + + if let Err(err) = result { + println!("FAILED HERE: {} {}\n{err:?}", worker_name.as_str(), job.sql); // TRUNCATE in case leftover state from error setup_subqueries_suite(worker_name.as_str(), false).await; @@ -64,6 +61,8 @@ pub async fn run_tests_for_case( sharded_conn = wire::connect(6432, "pgdog_sharded_subqueries").await; fail.push(job); + } else { + success += 1; } } @@ -225,16 +224,18 @@ pub enum StatementType { impl StatementType { pub fn parse(query: &str) -> Self { let query = query.to_lowercase(); - if query.contains("select") { - StatementType::Select - } else if query.contains("insert") { + if query.starts_with("insert") { StatementType::Insert - } else if query.contains("update") { + } else if query.starts_with("update") { StatementType::Update - } else if query.contains("delete") { + } else if query.starts_with("delete") { StatementType::Delete + } else if query.contains("select") { + // This deliberately uses `contains` instead of `starts_with` to match on CTEs. + // Also deliberately matched last to match INSERT / UPDATE / DELETE beforehand. + StatementType::Select } else { - panic!("Unknown query type!"); + panic!("Unknown query type! {query:?}"); } } } diff --git a/integration/rust/tests/integration/subqueries/tests/delete.sql b/integration/rust/tests/integration/subqueries/tests/delete.sql index 40bc1ade1..94d647287 100644 --- a/integration/rust/tests/integration/subqueries/tests/delete.sql +++ b/integration/rust/tests/integration/subqueries/tests/delete.sql @@ -107,7 +107,7 @@ DELETE FROM subusers WHERE id = 3 AND status IN (SELECT status FROM subusers WHE -- case -- what: .. WHERE = AND IN (), key in both, same shard --- route select: 1 +-- route delete: 1 -- verify: SELECT * FROM subusers ORDER BY id DELETE FROM subusers WHERE id = 3 AND status IN (SELECT status FROM subusers WHERE id = 3) diff --git a/integration/rust/tests/integration/subqueries/tests/insert.sql b/integration/rust/tests/integration/subqueries/tests/insert.sql index b9d76d8e0..becdf856b 100644 --- a/integration/rust/tests/integration/subqueries/tests/insert.sql +++ b/integration/rust/tests/integration/subqueries/tests/insert.sql @@ -3,9 +3,7 @@ -- route select: 0,1 -- route insert: 0 -- verify: SELECT * FROM subpackages ORDER BY id --- name: test INSERT INTO subpackages (id, user_id, order_id) SELECT id, user_id, id FROM suborders - -- case -- what: INSERT INTO SELECT .. FROM , source rows are on other shards than their destination -- route select: 0,1 @@ -18,7 +16,6 @@ INSERT INTO subpackages (id, user_id, order_id) SELECT id, user_id + 2, NULL FRO -- route select: 0,1 -- route insert: 0,1 -- verify: SELECT * FROM subpackages ORDER BY id --- name: test2 INSERT INTO subpackages (id, user_id, order_id) SELECT id, id + 1, NULL FROM suborders RETURNING * -- case @@ -103,7 +100,7 @@ INSERT INTO subpackages (id, user_id) SELECT max(id), max(user_id) FROM suborder -- route select: 0,1 -- route insert: 0 -- verify: SELECT * FROM subpackages ORDER BY id -INSERT INTO subpackages (id, user_id) SELECT avg(id), avg(id) FROM suborders +INSERT INTO subpackages (id, user_id) SELECT avg(id), max(user_id) FROM suborders -- case -- what: INSERT INTO SELECT .. FROM generate_series(..) @@ -115,7 +112,6 @@ INSERT INTO subusers(id) SELECT id FROM generate_series(10, 1000) as id -- case -- what: INSERT INTO SELECT .. FROM unnest($1, ..) -- $1 = {10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36} --- route select: any -- route insert: 0,1 -- verify: SELECT * FROM subusers ORDER BY id INSERT INTO subusers(id) SELECT id FROM UNNEST($1::int[]) as id diff --git a/integration/rust/tests/integration/subqueries/tests/update.sql b/integration/rust/tests/integration/subqueries/tests/update.sql index 199fa0a72..da78eaa2c 100644 --- a/integration/rust/tests/integration/subqueries/tests/update.sql +++ b/integration/rust/tests/integration/subqueries/tests/update.sql @@ -79,7 +79,7 @@ UPDATE subusers SET status = 'changed' WHERE id = 3 AND status IN (SELECT status -- case -- what: .. WHERE = AND IN (), key in both, same shard --- route select: 1 +-- route update: 1 -- verify: SELECT * FROM subusers ORDER BY id UPDATE subusers SET status = 'changed' WHERE id = 3 AND status IN (SELECT status FROM subusers WHERE id = 3) diff --git a/pgdog-postgres-types/src/datum.rs b/pgdog-postgres-types/src/datum.rs index b04096383..d392f9c85 100644 --- a/pgdog-postgres-types/src/datum.rs +++ b/pgdog-postgres-types/src/datum.rs @@ -376,6 +376,29 @@ impl DataType { } } + // Returns None for arrays, Vector and Other, since we can't send those types as one-dimensional arrays (unnest would flatten arrays of arrays) + pub fn into_array_type(&self) -> Option { + match self { + DataType::Bigint => Some(DataType::Array(1016)), + DataType::Integer => Some(DataType::Array(1007)), + DataType::Text => Some(DataType::Array(1009)), + DataType::Interval => Some(DataType::Array(1187)), + DataType::Timestamp => Some(DataType::Array(1115)), + DataType::TimestampTz => Some(DataType::Array(1185)), + DataType::Real => Some(DataType::Array(1021)), + DataType::DoublePrecision => Some(DataType::Array(1022)), + DataType::Bool => Some(DataType::Array(1000)), + DataType::SmallInt => Some(DataType::Array(1005)), + DataType::Numeric => Some(DataType::Array(1231)), + DataType::Uuid => Some(DataType::Array(2951)), + DataType::Oid => Some(DataType::Array(1028)), + DataType::Array(_) | DataType::Vector | DataType::Other(_) => None, + DataType::Cidr => todo!(), + DataType::Money => todo!(), + DataType::Inet => todo!(), + } + } + /// Format the `DataType` into a type that can be /// expressed as pg_catalog. /// diff --git a/pgdog/src/frontend/client/query_engine/context.rs b/pgdog/src/frontend/client/query_engine/context.rs index 2f783e947..847b6b7bc 100644 --- a/pgdog/src/frontend/client/query_engine/context.rs +++ b/pgdog/src/frontend/client/query_engine/context.rs @@ -47,6 +47,11 @@ pub(crate) struct QueryEngineContext<'a> { pub(super) query_size_limit: Option, /// When we received the first message of the request. pub(super) statement_start: DateTime, + + /// Is there a subquery in the pipeline that was + /// originally from the simple protocol? + /// TODO(jkaczman): (There may be a better way to do this!) + pub(super) drop_extended_parse_and_bind: bool, } impl<'a> QueryEngineContext<'a> { @@ -70,6 +75,7 @@ impl<'a> QueryEngineContext<'a> { query_log_stdout: client.query_log_stdout, query_size_limit: client.query_size_limit, statement_start: client.statement_start, + drop_extended_parse_and_bind: false, }, &mut client.client_request, ) @@ -100,6 +106,7 @@ impl<'a> QueryEngineContext<'a> { query_log_stdout: false, query_size_limit: None, statement_start: Utc::now(), + drop_extended_parse_and_bind: false, } } diff --git a/pgdog/src/frontend/client/query_engine/multi_step/error.rs b/pgdog/src/frontend/client/query_engine/multi_step/error.rs index 8f82c2db4..b7d723a6e 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/error.rs +++ b/pgdog/src/frontend/client/query_engine/multi_step/error.rs @@ -7,6 +7,9 @@ pub(crate) enum Error { #[error("{0}")] Update(#[from] UpdateError), + #[error("{0}")] + Subquery(#[from] SubqueryError), + #[error("parser: {0}")] Parser(Box), @@ -32,6 +35,12 @@ pub(crate) enum Error { SqlParser(#[from] pg_raw_parse::error::Error), } +#[derive(Debug, Error)] +pub(crate) enum SubqueryError { + #[error("subquery is returning a column with an unsupported data type")] + SubqueryReturningUnsupportedColumn, +} + #[derive(Debug, Error)] pub(crate) enum UpdateError { #[error("sharding key updates are forbidden")] diff --git a/pgdog/src/frontend/client/query_engine/multi_step/insert.rs b/pgdog/src/frontend/client/query_engine/multi_step/insert.rs index 1dd3ac0e8..08b8a0232 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/insert.rs +++ b/pgdog/src/frontend/client/query_engine/multi_step/insert.rs @@ -81,7 +81,7 @@ impl<'a> InsertMulti<'a> { // All tuples map to the same shard: send the original multi-row INSERT // as a single statement, skipping the multi-step path entirely. - if self.requests.same_shard().is_some() { + if self.requests.len() >= 2 && self.requests.same_shard().is_some() { self.engine .backend .handle_client_request( diff --git a/pgdog/src/frontend/client/query_engine/multi_step/mod.rs b/pgdog/src/frontend/client/query_engine/multi_step/mod.rs index 049406c23..aff20c018 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/mod.rs +++ b/pgdog/src/frontend/client/query_engine/multi_step/mod.rs @@ -6,7 +6,7 @@ pub(crate) mod state; pub(crate) mod subquery; pub(crate) mod update; -pub(crate) use error::{Error, UpdateError}; +pub(crate) use error::{Error, SubqueryError, UpdateError}; pub(crate) use fetch::fetch_rows; pub(crate) use forward_check::*; pub(crate) use insert::InsertMulti; diff --git a/pgdog/src/frontend/client/query_engine/multi_step/subquery/from_clause.rs b/pgdog/src/frontend/client/query_engine/multi_step/subquery/from_clause.rs index d82cdbaf7..bd6506118 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/subquery/from_clause.rs +++ b/pgdog/src/frontend/client/query_engine/multi_step/subquery/from_clause.rs @@ -3,14 +3,15 @@ use std::ops::Deref; use pg_raw_parse::{ Node, NodeMut, make::{MemoryToken, Unique}, - nodes::{self, Alias, SelectStmt}, + nodes::{self, Alias, RangeFunction, SelectStmt}, transform, }; -use pgdog_postgres_types::{DataType, Datum}; +use pgdog_postgres_types::{Array, Datum}; use crate::{ frontend::client::query_engine::multi_step::SubqueryMulti, - net::{DataRow, Decoder, RowDescription}, + frontend::router::parser::rewrite::statement::insert::is_constant, + net::{Bind, DataRow, Decoder, RowDescription, bind::Parameter}, }; impl<'a> SubqueryMulti<'a> { @@ -20,6 +21,7 @@ impl<'a> SubqueryMulti<'a> { mut statement: Unique<'mem, Node<'mem>>, row_description: Option<&RowDescription>, data_rows: &[DataRow], + new_bind: &mut Bind, ) -> Unique<'mem, Node<'mem>> { let Some(row_description) = row_description else { return statement; @@ -37,15 +39,17 @@ impl<'a> SubqueryMulti<'a> { let mut range_subselect = mem.make_unique(range_subselect.deref()); // Replace the subquery in the statement with the DataRows (via VALUES) { - let explicit_values_list = Self::new_select_with_explicit_values( + let new_subselect = Self::new_select_with_unnest( + new_bind, mem, - row_description.clone(), + range_subselect.subquery(), data_rows, + row_description, ); range_subselect .as_mut() - .set_subquery(explicit_values_list.uncast()); + .set_subquery(new_subselect.uncast()); } { @@ -99,110 +103,152 @@ impl<'a> SubqueryMulti<'a> { statement } - /// Construct a new `SelectStmt` to be set as the `subselect` of our new `RangeSubselect` node, - /// whose inner-statement now is an *explicit* VALUES list of the returning results - /// of the executed inner-query (as in, the returned `DataRow` list articulated explicitly - /// within the outer statement) - /// - /// TODO(jkaczman): we need to account for too much data being present, and split - /// apart the query into multiple when that occurs. This TODO also exists elsewhere for subqueries. - pub(crate) fn new_select_with_explicit_values<'mem>( + pub(crate) fn new_select_with_unnest<'mem>( + old_bind: &mut Bind, mem: MemoryToken<'mem>, - row_description: RowDescription, + original_ast: Node<'_>, data_rows: &[DataRow], + row_description: &RowDescription, ) -> Unique<'mem, &'mem SelectStmt> { - let mut values_select = mem.make_node::(); - let mut decoder = Decoder::default(); - decoder.set_row_description(row_description); - - let mut elems = Vec::new(); + decoder.set_row_description(row_description.clone()); + + let base_param_len = old_bind.params_raw().len() as i32; + + let original_targets = match original_ast { + Node::SelectStmt(original_select) + if original_select.target_list().len() == row_description.fields.len() => + { + original_select + .target_list() + .into_iter() + .map(|target| Some(target.val())) + .collect() + } + _ => vec![None; row_description.fields.len()], + }; - if data_rows.is_empty() { - // We need to create a fake, all-NULL row, representing all columns - // within our subquery, so that we can tell Postgres what the - // RowDescription types to be to the client - // - // TODO(jkaczman): See the matching TODO for this in the area - // below where we do the explicit type casts for the first row - // when we don't have `DataRows::is_empty()` - let mut singular_empty_top_level_list = Vec::new(); - for index in 0..decoder.row_description().fields.len() { - if let Some(field) = decoder.row_description().field(index) { - let data_type: DataType = field.data_type(); - let pg_catalog_type: &str = data_type.as_pg_catalog_type(); - - let mut node = mem.make_node::(); - node.as_mut().set_isnull(true); - - let node = mem - .make_type_cast( - node.uncast(), - mem.make_list(&[ - mem.make_string(Some("pg_catalog")), - mem.make_string(Some(pg_catalog_type)), - ]), - ) - .uncast(); - - singular_empty_top_level_list.push(node); - } + // Form the bind params + let mut param_refs = Vec::new(); + let mut select_targets = Vec::new(); + let mut unnest_column_names = Vec::new(); + for (idx, field) in row_description.fields.iter().enumerate() { + if let Some(constant) = original_targets[idx].filter(|target| is_constant(*target)) { + select_targets.push(mem.make_res_target( + None, + mem.empty(), + mem.make_unique(constant), + )); + continue; } - // Now form a singular VALUES list containing NULL entries for all applicable - // columns - let new_list = mem.make_list(singular_empty_top_level_list.as_slice()); - elems.push(new_list); - } else { - for (row_idx, row) in data_rows.iter().enumerate() { - let mut second = Vec::new(); - - for col in 0..row.len() { - let datum: Datum = row - .get_column(col, &decoder) - .expect("get_column never returns an Error") - .unwrap() - .value; - - let mut node = Self::convert_datum_to_a_const(datum.clone(), mem); - - // TODO(jkaczman): Right now, this does explicit casts for each type on the - // first row as a workaround to properly set the RowDescription - // type OIDs in the response to the client. This is more query - // verbosity, which saves us backend logic to intercept the - // RowDescription response and modify it, however, at the expense - // of (well, making the query more verbose). Aka, this works fine, - // but could be represented better, perhaps by in the future intercepting - // when this is more mature. - if row_idx == 0 { - let data_type_str: &str = datum.data_type().as_pg_catalog_type(); - - // TODO(jkaczman): this is a frequently used construct; - // I think that we can create a method to share this throughout - // all the kinds... they all use pg_catalog casts - node = mem - .make_type_cast( - node.uncast(), - mem.make_list(&[ - mem.make_string(Some("pg_catalog")), - mem.make_string(Some(data_type_str)), - ]), - ) - .uncast(); - } + let unnest_column_name = format!("c{}", param_refs.len() + 1); + select_targets.push( + mem.make_res_target( + None, + mem.empty(), + mem.make_column_ref(mem.make_list(&[ + mem.make_string(Some("u")).uncast(), + mem.make_string(Some(unnest_column_name.as_str())).uncast(), + ])) + .uncast(), + ), + ); - second.push(node.uncast()); - } + unnest_column_names.push(mem.make_string(Some(unnest_column_name.as_str())).uncast()); + + let pg_catalog_type: &str = field + .data_type() + .into_array_type() + .expect("rewrite_outer_statement_with_subquery_response checked the column types") + .as_pg_catalog_type(); - let new_list = mem.make_list(second.as_slice()); - elems.push(new_list); + let node = { + let mut node = mem.make_node::(); + + node.as_mut() + .set_number(base_param_len + param_refs.len() as i32 + 1); + + node + }; + + let mut list_of_rows_for_col: Vec = Vec::new(); + for row in data_rows { + // TODO(jkaczman): replace all unwraps() + let col = row.get_column(idx, &decoder).unwrap().unwrap(); + list_of_rows_for_col.push(col.value); } + + let array = Array::new(list_of_rows_for_col, field.type_oid); + + // TODO(jkaczman): wonder what the performance tradeoffs are btwn. Binary and Text + let default_fmt = old_bind.default_param_format(); + old_bind.push_param( + Parameter::new(array.encode(default_fmt).unwrap().as_ref()), + default_fmt, + ); + + let node = mem + .make_type_cast( + node.uncast(), + mem.make_list(&[ + mem.make_string(Some("pg_catalog")), + mem.make_string(Some(pg_catalog_type)), + ]), + ) + .uncast(); + + param_refs.push(node); } - let new_list = mem.make_list(elems.as_slice()); - values_select.as_mut().set_values_lists(new_list); + let list_of_param_refs = mem.make_list(param_refs.as_slice()); + let unnest_func_call = mem.make_func_call( + mem.make_list(&[mem.make_string(Some("unnest")).uncast()]), + list_of_param_refs, + 0, + ); + + let unnest_alias = { + let mut unnest_alias = mem.make_node::(); + + unnest_alias + .as_mut() + .set_aliasname(Some(mem.copy_string("u"))); + + unnest_alias + .as_mut() + .set_colnames(mem.make_list(unnest_column_names.as_slice())); + + unnest_alias + }; + + let unnest_func_wrapper = { + let mut unnest_func_wrapper = mem.make_node::(); - values_select + unnest_func_wrapper + .as_mut() + .set_functions(mem.make_list(&[mem.make_list(&[unnest_func_call.uncast()])])); + + unnest_func_wrapper + .as_mut() + .set_alias(unnest_alias.as_option()); + + unnest_func_wrapper + }; + + { + let mut new_select = mem.make_node::(); + + new_select + .as_mut() + .set_from_clause(mem.make_list(&[unnest_func_wrapper.uncast()])); + + new_select + .as_mut() + .set_target_list(mem.make_list(select_targets.as_slice())); + + new_select + } } /// Assemble a new alias to be used within the `RangeSubselect` node we're re-writing, diff --git a/pgdog/src/frontend/client/query_engine/multi_step/subquery/insert_select.rs b/pgdog/src/frontend/client/query_engine/multi_step/subquery/insert_select.rs index c27a8649e..1e428beb2 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/subquery/insert_select.rs +++ b/pgdog/src/frontend/client/query_engine/multi_step/subquery/insert_select.rs @@ -1,12 +1,12 @@ use pg_raw_parse::{ Node, NodeMut, make::{MemoryToken, Unique}, - nodes::{SelectStmt, WithClause}, + nodes::WithClause, }; use crate::{ frontend::client::query_engine::multi_step::SubqueryMulti, - net::{DataRow, Decoder, RowDescription}, + net::{Bind, DataRow, RowDescription}, }; impl<'a> SubqueryMulti<'a> { @@ -18,34 +18,23 @@ impl<'a> SubqueryMulti<'a> { mut stmt: Unique<'mem, Node<'mem>>, row_description: Option<&RowDescription>, data_rows: &[DataRow], + new_bind: &mut Bind, ) -> Unique<'mem, Node<'mem>> { let NodeMut::InsertStmt(mut insert) = stmt.as_mut() else { return stmt; }; - // Replace with SelectStmt with new VALUES - let mut new_select = mem.make_node::(); - let mut new_select_mut = new_select.as_mut(); - - let mut decoder = Decoder::default(); - decoder.set_row_description(row_description.unwrap().clone()); - - // Build the VALUES list from the DataRows - let mut lists = Vec::new(); - for row in data_rows { - let mut elems = Vec::new(); - - for col in 0..row.len() { - let datum = row.get_column(col, &decoder).unwrap().unwrap().value; - elems.push(Self::convert_datum_to_a_const(datum, mem).uncast()); - } - - let tuple = mem.make_list(elems.as_ref()).uncast(); - lists.push(tuple); - } + let Some(row_description) = row_description else { + return stmt; + }; - let values_lists = mem.make_list(lists.as_ref()); - new_select_mut.set_values_lists(values_lists); + let new_select = Self::new_select_with_unnest( + new_bind, + mem, + insert.select_stmt(), + data_rows, + row_description, + ); insert.set_select_stmt(new_select.uncast()); insert.set_with_clause(mem.make_unique(None::<&WithClause>)); diff --git a/pgdog/src/frontend/client/query_engine/multi_step/subquery/mod.rs b/pgdog/src/frontend/client/query_engine/multi_step/subquery/mod.rs index 488b2451a..d3057a3e5 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/subquery/mod.rs +++ b/pgdog/src/frontend/client/query_engine/multi_step/subquery/mod.rs @@ -22,7 +22,7 @@ use super::Error; use crate::{ frontend::{ ClientRequest, - client::query_engine::{QueryEngine, QueryEngineContext}, + client::query_engine::{QueryEngine, QueryEngineContext, multi_step::error::SubqueryError}, router::parser::rewrite::statement::{ insert::Subquery, subquery::{ @@ -45,6 +45,11 @@ pub(crate) struct SubqueryMulti<'ast> { // within the client's Bind message. pub(super) top_level_bind: Option<&'ast Bind>, + // Start off with the very outer statement's Bind message. + // Add our own parameters as we need to inject DataRows into the statement as + // unnest(..) + pub(super) current_level_bind: Bind, + // Store the results here, so that can reference previously executed subqueries (e.g., if we need to // use earlier resolved CTEs to resolve a CTE occurring after those) pub(super) subquery_results: Vec<(Vec, Option)>, @@ -95,9 +100,12 @@ impl<'ast> SubqueryMulti<'ast> { None }; + let current_level_bind = top_level_bind.cloned().unwrap_or(Bind::default()); + Self { subqueries, top_level_bind, + current_level_bind, subquery_results: Vec::new(), is_extended_protocol, top_level_statement, @@ -112,7 +120,11 @@ impl<'ast> SubqueryMulti<'ast> { Self { subqueries, top_level_bind: parent_engine.top_level_bind, - subquery_results: parent_engine.subquery_results.clone(), + current_level_bind: parent_engine + .top_level_bind + .cloned() + .unwrap_or(Bind::default()), + subquery_results: Vec::new(), is_extended_protocol: parent_engine.is_extended_protocol, top_level_statement: new_parent_ast, } @@ -128,13 +140,13 @@ impl<'ast> SubqueryMulti<'ast> { context: &mut QueryEngineContext<'_>, ) -> Result<(), Error> { // Send the inner-subquery and re-write the outer statement with its response. - let rewritten_outer_statement = self + let (new_ast, new_bind): (Owned, Bind) = self .evaluate_and_rewrite_outer_statement(engine, context) .await?; - let rewritten_outer_statement = deparse(rewritten_outer_statement.stmt())? - .as_str() - .to_string(); + self.current_level_bind = new_bind; + + let rewritten_outer_statement = deparse(new_ast.stmt())?.as_str().to_string(); // Send the re-written statement self.send_query_string( @@ -157,7 +169,7 @@ impl<'ast> SubqueryMulti<'ast> { &mut self, engine: &mut QueryEngine, context: &mut QueryEngineContext<'_>, - ) -> Result, Error> { + ) -> Result<(Owned, Bind), Error> { // TODO(jkaczman): we could execute some of these in parallel (in groups) e.g., if CTEs are not referencing earlier ones let current_statement: Node<'ast> = self.top_level_statement; @@ -193,9 +205,10 @@ impl<'ast> SubqueryMulti<'ast> { } // Re-write outer statement with the response - // TODO(jkaczman): have this return the AST as well ig - let new_ast: Owned = self.rewrite_outer_statement_with_subquery_response()?; - Ok(new_ast) + let (new_ast, new_bind): (Owned, Bind) = + self.rewrite_outer_statement_with_subquery_response()?; + + Ok((new_ast, new_bind)) } fn rewrite_subquery_ast_with_earlier_subquery_responses( @@ -337,9 +350,36 @@ impl<'ast> SubqueryMulti<'ast> { }.expect("We already ran these methods in the parser originally. No reason they shouldn't pass here.") } - fn rewrite_outer_statement_with_subquery_response(&self) -> Result, Error> { + fn rewrite_outer_statement_with_subquery_response( + &self, + ) -> Result<(Owned, Bind), Error> { // TODO(jkaczman): if we're always passing these three, might as well make a new type or // hold these in the existing one + let mut new_bind = self.current_level_bind.clone(); + + // We send the rows of subqueries as unnest(..) arrays, so we check the column types before rewriting to ensure that all columns have array types. + for (subquery, (_, row_description)) in + self.subqueries.iter().zip(self.subquery_results.iter()) + { + if matches!(subquery, Subquery::WhereClause) { + continue; + } + + let Some(row_description) = row_description else { + continue; + }; + + if row_description + .fields + .iter() + .find(|field| field.data_type().into_array_type().is_none()) + .is_some() + { + return Err(Error::Subquery( + SubqueryError::SubqueryReturningUnsupportedColumn, + )); + } + } let new_statement: Owned = owned(|mem| { let mut statement: Unique<'_, Node<'_>> = mem.make_unique(self.top_level_statement); @@ -357,13 +397,16 @@ impl<'ast> SubqueryMulti<'ast> { subquery, row_description.as_ref(), data_rows.as_slice(), + &mut new_bind, ); } + // We also need to capture Bind parameters from this + mem.make_raw_stmt(statement) }); - Ok(new_statement) + Ok((new_statement, new_bind)) } fn rewrite_with_subquery_response<'mem>( @@ -373,20 +416,23 @@ impl<'ast> SubqueryMulti<'ast> { subquery: &Subquery, row_description: Option<&RowDescription>, data_rows: &[DataRow], + new_bind: &mut Bind, ) -> Unique<'mem, Node<'mem>> { match subquery { - Subquery::FromClause => self.handle_from_clause(mem, stmt, row_description, data_rows), + Subquery::FromClause => { + self.handle_from_clause(mem, stmt, row_description, data_rows, new_bind) + } Subquery::WhereClause => { - self.handle_where_clause(mem, stmt, row_description, data_rows) + self.handle_where_clause(mem, stmt, row_description, data_rows, new_bind) } Subquery::InsertSelect => { - self.handle_insert_select(mem, stmt, row_description, data_rows) + self.handle_insert_select(mem, stmt, row_description, data_rows, new_bind) } Subquery::WhereClauseCte(alias) => { - self.handle_where_clause_cte(alias, mem, stmt, row_description, data_rows) + self.handle_where_clause_cte(alias, mem, stmt, row_description, data_rows, new_bind) } Subquery::Cte(alias) => { - self.handle_with_clause(alias, mem, stmt, row_description, data_rows) + self.handle_with_clause(alias, mem, stmt, row_description, data_rows, new_bind) } } } diff --git a/pgdog/src/frontend/client/query_engine/multi_step/subquery/request.rs b/pgdog/src/frontend/client/query_engine/multi_step/subquery/request.rs index 4a9606053..fc5c55555 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/subquery/request.rs +++ b/pgdog/src/frontend/client/query_engine/multi_step/subquery/request.rs @@ -46,6 +46,7 @@ impl<'a> SubqueryMulti<'a> { &cluster.schema(), rewrite_result.offset_plan(), )?; + rewrite_result.apply_after_route(&mut client_request)?; Ok(Some(multi_step::fetch_rows(engine, &client_request).await?)) @@ -58,7 +59,7 @@ impl<'a> SubqueryMulti<'a> { // // Then, we send the re-written subquery here, and return the relevant DataRows / RD - let rewritten_subquery_ast = { + let (new_ast, new_bind) = { let mut subquery_multi = SubqueryMulti::new_sub_engine(self, query_ast, subqueries); @@ -68,12 +69,13 @@ impl<'a> SubqueryMulti<'a> { .await? }; - let rewritten_subquery_sql = - deparse(rewritten_subquery_ast.stmt())?.as_str().to_string(); + self.current_level_bind = new_bind; + + let rewritten_subquery_sql = deparse(new_ast.stmt())?.as_str().to_string(); Box::pin(self.send_query_string( engine, - rewritten_subquery_ast.stmt(), + new_ast.stmt(), rewritten_subquery_sql, context, true, @@ -103,6 +105,14 @@ impl<'a> SubqueryMulti<'a> { todo!() } + // We send outer queries that we want evaluated in the extended protocol, + // even if the original query was sent in the simple protocol, as we need + // to send a Bind message with the subquery contents in a byte array vs. + // hand-rolling it into the SQL. + // + // So, this just drops all Parse/Binds that we come across, as we know we're handling subqueries. + context.drop_extended_parse_and_bind = !self.is_extended_protocol; + // TODO(jkaczman): We rely on INSERT split for insert statements. // If a user has that rewrite disabled, but subqueries // enabled, then INSERTs will still fail with @@ -122,35 +132,35 @@ impl<'a> SubqueryMulti<'a> { mut query_string: String, context: &mut QueryEngineContext<'_>, ) -> Result { - let mut client_request = if self.is_extended_protocol { + let mut client_request = { let mut bind = Bind::default(); - if let Some(top_level_bind) = self.top_level_bind { - let mut actual_param_index = 1; - let mut actual_param_index_str = format!("${actual_param_index}"); - - for index in 1..=top_level_bind.params_raw().len() { - // If we have a Bind parameter within the subquery, then we should - // include it in our Bind message. This also means that we have to rename - // it to a local scheme to the subquery itself (e.g., if the first - // one we encounter in the subquery is $2, needs to be renamed to $1) - // - // TODO(jkaczman): we need to also update the original (top-lvl) ClientRequest's Bind - // to get rid of the evaluated params - let param_str = format!("${}", index); - if query_string.contains(param_str.as_str()) { - query_string = query_string - .replace(param_str.as_str(), actual_param_index_str.as_str()); - - actual_param_index += 1; - actual_param_index_str = format!("${}", actual_param_index); - - let old_bind_param = top_level_bind.parameter(index - 1).unwrap().unwrap(); - - bind.push_param( - old_bind_param.parameter().clone(), - old_bind_param.format(), - ); - } + let mut actual_param_index = 1; + let mut actual_param_index_str = format!("${actual_param_index}"); + + for index in 1..=self.current_level_bind.params_raw().len() { + // If we have a Bind parameter within the subquery, then we should + // include it in our Bind message. This also means that we have to rename + // it to a local scheme to the subquery itself (e.g., if the first + // one we encounter in the subquery is $2, needs to be renamed to $1) + // + // TODO(jkaczman): we need to also update the original (top-lvl) ClientRequest's Bind + // to get rid of the evaluated params that are extra (self.top_level_bind) + // TODO(jkaczman): This also won't work if we have 10 Bind params ($1 matching on $10) + let param_str = format!("${}", index); + if query_string.contains(param_str.as_str()) { + query_string = + query_string.replace(param_str.as_str(), actual_param_index_str.as_str()); + + actual_param_index += 1; + actual_param_index_str = format!("${}", actual_param_index); + + let old_bind_param = self + .current_level_bind + .parameter(index - 1) + .unwrap() + .unwrap(); + + bind.push_param(old_bind_param.parameter().clone(), old_bind_param.format()); } } @@ -161,8 +171,6 @@ impl<'a> SubqueryMulti<'a> { ProtocolMessage::Execute(Default::default()), ProtocolMessage::Sync(Default::default()), ]) - } else { - ClientRequest::from(vec![ProtocolMessage::Query(Query::new(&query_string))]) }; let cluster = engine.backend.cluster()?; diff --git a/pgdog/src/frontend/client/query_engine/multi_step/subquery/where_clause.rs b/pgdog/src/frontend/client/query_engine/multi_step/subquery/where_clause.rs index 7b9b2e7e3..af718b163 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/subquery/where_clause.rs +++ b/pgdog/src/frontend/client/query_engine/multi_step/subquery/where_clause.rs @@ -12,7 +12,7 @@ use tracing::debug; use crate::{ frontend::{client::query_engine::multi_step::SubqueryMulti, router::parser::rewrite}, - net::{DataRow, Decoder, RowDescription}, + net::{Bind, DataRow, Decoder, RowDescription}, }; impl<'a> SubqueryMulti<'a> { @@ -22,6 +22,7 @@ impl<'a> SubqueryMulti<'a> { mut statement: Unique<'mem, Node<'mem>>, row_description: Option<&RowDescription>, data_rows: &[DataRow], + _new_bind: &Bind, ) -> Unique<'mem, Node<'mem>> { let Some((sublink, lexpr)) = rewrite::statement::subquery::where_clause::has_sub_link_in_where(statement.as_ref()) @@ -68,6 +69,7 @@ impl<'a> SubqueryMulti<'a> { statement } + // TODO: can use unnest() here too fn rewrite_sublink_in_where_clause<'mem>( &self, sublink: &SubLink, diff --git a/pgdog/src/frontend/client/query_engine/multi_step/subquery/with_clause.rs b/pgdog/src/frontend/client/query_engine/multi_step/subquery/with_clause.rs index ba329ff90..4cc7dbd86 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/subquery/with_clause.rs +++ b/pgdog/src/frontend/client/query_engine/multi_step/subquery/with_clause.rs @@ -9,7 +9,7 @@ use crate::{ client::query_engine::multi_step::SubqueryMulti, router::parser::rewrite::statement::subquery::where_clause::has_sub_link_in_where, }, - net::{DataRow, RowDescription}, + net::{Bind, DataRow, RowDescription}, }; impl<'a> SubqueryMulti<'a> { @@ -20,6 +20,7 @@ impl<'a> SubqueryMulti<'a> { mut statement: Unique<'mem, Node<'mem>>, row_description: Option<&'a RowDescription>, data_rows: &'a [DataRow], + new_bind: &mut Bind, ) -> Unique<'mem, Node<'mem>> { let Some((sublink, _)) = has_sub_link_in_where(statement.as_ref()) else { return statement; @@ -31,6 +32,7 @@ impl<'a> SubqueryMulti<'a> { mem.make_unique(sublink.subselect()), row_description, data_rows, + new_bind, ); let mut rewritten_sublink = mem.make_unique(sublink); @@ -76,6 +78,7 @@ impl<'a> SubqueryMulti<'a> { mut statement: Unique<'mem, Node<'mem>>, row_description: Option<&'a RowDescription>, data_rows: &'a [DataRow], + new_bind: &mut Bind, ) -> Unique<'mem, Node<'mem>> { let Some(row_description) = row_description else { return statement; @@ -115,8 +118,13 @@ impl<'a> SubqueryMulti<'a> { // // Any memory implications from storing it? Is it good // as it's reusable? What about the length of the strings? - let new_cte_query = - Self::new_select_with_explicit_values(mem, row_description.clone(), data_rows); + let new_cte_query = Self::new_select_with_unnest( + new_bind, + mem, + cte.ctequery(), + data_rows, + row_description, + ); let mut alias_col_names = mem.make_unique(cte.aliascolnames()); if alias_col_names.is_empty() { diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs index 0e3abc8ce..b9d5ba851 100644 --- a/pgdog/src/frontend/client/query_engine/query.rs +++ b/pgdog/src/frontend/client/query_engine/query.rs @@ -140,9 +140,24 @@ impl QueryEngine { // Thus, we should now fetch the Bind, if present, and pass it in (at the top-level), // so that we can supply the very top level's Bind to all sub-levels. - multi_step::SubqueryMulti::new(subquery, client_request) + match multi_step::SubqueryMulti::new(subquery, client_request) .execute(self, context) - .await?; + .await + { + Ok(()) => {} + Err(err @ multi_step::Error::Subquery(_)) => { + if let Some(transaction) = context.transaction.as_mut() { + transaction.set_error(); + } + self.error_response( + context, + client_request, + ErrorResponse::unsupported(&err), + ) + .await?; + } + Err(err) => return Err(err.into()), + } } Some(RewriteResult::InPlace { .. }) | None @@ -350,11 +365,21 @@ impl QueryEngine { // Pipelined requests only return // one ReadyForQuery message. - let drop_message = message.code() == 'Z' + let mut drop_message = message.code() == 'Z' && !context.pipeline.is_done() && context.pipeline.is_simple() && !context.in_error(); // On error, pipeline is done executing. + // TODO(jkaczman): these constants can go in a new crate (waiting to rebase off main for the whole stack first) + const BIND_COMPLETE_CODE: char = '2'; + const PARSE_COMPLETE_CODE: char = '1'; + const NO_DATA_CODE: char = 'n'; + let not_allowed = &[BIND_COMPLETE_CODE, PARSE_COMPLETE_CODE, NO_DATA_CODE]; + + drop_message = drop_message + || (context.drop_extended_parse_and_bind + && (not_allowed.iter().find(|e| message.code() == **e).is_some())); + if !drop_message { trace!("{:#?} >>> {:?}", message, context.stream.peer_addr()); diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index 1ac2c5ef9..632aa1467 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -668,6 +668,7 @@ impl QueryParser { user: context.router_context.cluster.user(), search_path: context.router_context.parameter_hints.search_path, }; + let insert_rewritten = context.router_context.ast.is_some_and(|ast| { !ast.rewrite_plan.subqueries.is_empty() || matches!( @@ -681,6 +682,7 @@ impl QueryParser { .with_schema_lookup(schema_lookup) .with_explain(self.recorder_mut().is_some()) .with_insert_rewritten(insert_rewritten); + parser.set_resolved_lookups(&context.router_context.resolved_lookups); let is_sharded = parser.is_sharded( diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs index e11913ad9..374f370fb 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs @@ -331,6 +331,7 @@ pub(crate) fn group_split_requests_by_shard( statement_name, } .build_request(request)?; + new_request.route = requests[tuples[0]].route.clone(); grouped.push(new_request); } @@ -557,7 +558,7 @@ impl<'a> UnnestOutputNames<'a> { } // Returns true for constants with or without a type cast, e.g., NULL, 'VIP' and 'VIP'::text. -fn is_constant(node: Node<'_>) -> bool { +pub(crate) fn is_constant(node: Node<'_>) -> bool { match node { Node::A_Const(_) => true, Node::TypeCast(type_cast) => matches!(type_cast.arg(), Node::A_Const(_)), diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/subquery/from_clause.rs b/pgdog/src/frontend/router/parser/rewrite/statement/subquery/from_clause.rs index 1584d0970..40dd14f18 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/subquery/from_clause.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/subquery/from_clause.rs @@ -26,8 +26,25 @@ pub fn parse_from_clause<'a>(outer_stmt: Node<'a>) -> Option> { let subquery = range_subselect.subquery(); if range_subselect.alias().is_some() && let Node::SelectStmt(inner_subselect) = subquery - && inner_subselect.values_lists().is_empty() { + // TODO(jkaczman): Redundant with INSERT SELECT + + // We don't want to recurse back through for a re-written inner query. + if let Some(Node::RangeFunction(range_func)) = inner_subselect.from_clause().first() + && let Some(Node::NodeList(node_list)) = range_func.functions().first() + && let Some(Node::FuncCall(func_call)) = node_list.first() + && let Some(func_name) = func_call.funcname().first() + && let Some(func_name) = func_name.as_str() + && func_name.eq("unnest") + { + return None; + } + + // If VALUES is specified, no need for us to re-write. + if !inner_subselect.values_lists().is_empty() { + return None; + } + // There's no `is_sharded` check here, because we could have // a case where we need to execute a subquery that's NOT sharded, // such as generate_series, which would resolve into a diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/subquery/insert_select.rs b/pgdog/src/frontend/router/parser/rewrite/statement/subquery/insert_select.rs index ed2194ead..6b896199c 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/subquery/insert_select.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/subquery/insert_select.rs @@ -8,9 +8,7 @@ use crate::frontend::router::parser::{ impl<'a, 'mem> SubqueryRewrite<'a, 'mem> { pub(crate) fn handle_subquery_in_insert_select(&self) -> Option { if let Node::InsertStmt(insert_stmt) = self.outer_stmt - && let Node::SelectStmt(select) = insert_stmt.select_stmt() - // This won't be a SubLink node as the select statement is naturally apart of the insert's AST. - && !select.from_clause().is_empty() + && parse_insert_select(self.outer_stmt).is_some() { // Is it inserting into a sharded table? let mut parser = StatementParser::new( @@ -38,10 +36,26 @@ impl<'a, 'mem> SubqueryRewrite<'a, 'mem> { pub(crate) fn parse_insert_select<'a>(entire_ast: Node<'a>) -> Option> { if let Node::InsertStmt(insert_stmt) = entire_ast - && let Node::SelectStmt(select) = insert_stmt.select_stmt() - // This won't be a SubLink node as the select statement is naturally apart of the insert's AST. - && !select.from_clause().is_empty() + && let Node::SelectStmt(select) = insert_stmt.select_stmt() + // This won't be a SubLink node as the select statement is naturally apart of the insert's AST. { + // We don't want to recurse back through for a re-written query. + // TODO(jkaczman): this is a redundant check that we use in multiple places; could + // put this in a separate helper method in a util mod + if let Some(Node::RangeFunction(range_func)) = select.from_clause().first() + && let Some(Node::NodeList(node_list)) = range_func.functions().first() + && let Some(Node::FuncCall(func_call)) = node_list.first() + && let Some(funcname) = func_call.funcname().first() + && let Some(funcname) = funcname.as_str() + && funcname.eq("unnest") + { + return None; + } + + if !select.values_lists().is_empty() { + return None; + } + return Some(insert_stmt.select_stmt()); }