perf: migrate to unnest (from VALUES) for CTEs, INSERT .. SELECT FROM, SELECT/UPDATE FROM

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=<case_name>
     let test_target: Option<String> = 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();

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:?}");
         }
     }
 }

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 <key> = <k1> AND <col> IN (<sq>), 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)
 

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 <st1> SELECT .. FROM <st2>, 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 <st1> 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 <st1> 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

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 <key> = <k1> AND <col> IN (<sq>), 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)
 

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<DataType> {
+        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.<TYPE>
     ///

pgdog/src/frontend/client/query_engine/context.rs

@@ -47,6 +47,11 @@ pub(crate) struct QueryEngineContext<'a> {
     pub(super) query_size_limit: Option<usize>,
     /// When we received the first message of the request.
     pub(super) statement_start: DateTime<Utc>,
+
+    /// 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,
         }
     }
 

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<crate::frontend::router::parser::Error>),
 
@@ -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")]

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(

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;