feat: split INSERT .. SELECT <columns> FROM unnest(..) into sharded tables (not only SELECT *)

pgdog/src/frontend/client/query_engine/test/rewrite_insert_split.rs

@@ -188,13 +188,16 @@ fn unnest_insert_messages(
 }
 
 // Returns the (id, value) rows that the requests insert (one list per request)
-fn inserted_rows_by_request(requests: &[ClientRequest]) -> Vec<Vec<(i64, Option<String>)>> {
+fn inserted_rows_by_request(
+    requests: &[ClientRequest],
+    unnest_insert: &str,
+) -> Vec<Vec<(i64, Option<String>)>> {
     requests
         .iter()
         .map(|request| {
             std::assert_matches!(
                 &request.messages[0],
-                ProtocolMessage::Parse(parse) if parse.query() == UNNEST_INSERT,
+                ProtocolMessage::Parse(parse) if parse.query() == unnest_insert,
                 "expected the original unnest INSERT"
             );
 
@@ -245,7 +248,7 @@ async fn test_insert_split_unnest() {
     assert_eq!(requests.len(), 2, "expected one INSERT per shard");
     assert_ne!(requests[0].route().shard(), requests[1].route().shard());
 
-    let mut rows: Vec<_> = inserted_rows_by_request(&requests)
+    let mut rows: Vec<_> = inserted_rows_by_request(&requests, UNNEST_INSERT)
         .into_iter()
         .flatten()
         .collect();
@@ -305,7 +308,7 @@ async fn test_insert_split_unnest_binary_one_shard() {
     assert_eq!(requests.len(), 2, "expected one INSERT per shard");
 
     // Two ids that the router sent to the same shard.
-    let same_shard_ids: Vec<i64> = inserted_rows_by_request(&requests)[0]
+    let same_shard_ids: Vec<i64> = inserted_rows_by_request(&requests, UNNEST_INSERT)[0]
         .iter()
         .take(2)
         .map(|(id, _)| *id)
@@ -321,10 +324,61 @@ async fn test_insert_split_unnest_binary_one_shard() {
 
     assert_eq!(requests.len(), 1, "expected one INSERT for one shard");
     assert_eq!(
-        inserted_rows_by_request(&requests)[0],
+        inserted_rows_by_request(&requests, UNNEST_INSERT)[0],
         vec![
             (same_shard_ids[0], Some("a".to_string())),
             (same_shard_ids[1], Some("b".to_string())),
         ]
     );
 }
+
+// We send all 20 rows into the table sharded (sharded by id) two times.
+// Both below use unnest($1::pg_catalog._int8, $2::pg_catalog._text); $1 = array of ids, $2 = array of values
+// - UNNEST_INSERT runs `INSERT INTO sharded (id, value) SELECT * FROM unnest(..)`.
+// - UNNEST_INSERT_WITH_TARGETS runs `INSERT INTO sharded (value, status, id) SELECT u.value, NULL, u.id FROM unnest(..) AS u(id, value)` (the SELECT lists the unnest(..) output columns in a different order than the unnest(..) arguments and adds a NULL for status).
+#[tokio::test]
+async fn test_insert_split_unnest_select_targets() {
+    const UNNEST_INSERT_WITH_TARGETS: &str = "INSERT INTO sharded (value, status, id) SELECT u.value, NULL, u.id FROM unnest($1::pg_catalog._int8, $2::pg_catalog._text) AS u(id, value)";
+
+    let ids: Vec<i64> = (1..=20).collect();
+    let values: Vec<String> = ids.iter().map(|id| format!("v{id}")).collect();
+    let values: Vec<Option<&str>> = values.iter().map(|value| Some(value.as_str())).collect();
+
+    let rows_by_shard = |requests: &[ClientRequest], unnest_insert: &str| {
+        let mut rows_by_shard: Vec<_> = requests
+            .iter()
+            .map(|request| request.route().shard().clone())
+            .zip(inserted_rows_by_request(requests, unnest_insert))
+            .collect();
+
+        rows_by_shard.sort_by_key(|(_, rows)| rows[0].0);
+        rows_by_shard
+    };
+
+    let select_star_requests = run_test(unnest_insert_messages(
+        UNNEST_INSERT,
+        &ids,
+        &values,
+        Format::Text,
+    ))
+    .await;
+
+    let select_targets_requests = run_test(unnest_insert_messages(
+        UNNEST_INSERT_WITH_TARGETS,
+        &ids,
+        &values,
+        Format::Text,
+    ))
+    .await;
+
+    assert_eq!(
+        select_targets_requests.len(),
+        2,
+        "expected one INSERT per shard"
+    );
+
+    assert_eq!(
+        rows_by_shard(&select_targets_requests, UNNEST_INSERT_WITH_TARGETS),
+        rows_by_shard(&select_star_requests, UNNEST_INSERT)
+    );
+}