From 33658c4e25369719445ea91f7b87fafedd661a0c Mon Sep 17 00:00:00 2001
From: Tom Harding <i.am.tom.harding@gmail.com>
Date: Wed, 14 Feb 2024 13:28:41 +0100
Subject: [PATCH] It lives!

---
 crates/query-engine/sql/src/sql/helpers.rs    |  61 ++++---
 .../translation/src/translation/error.rs      |   7 +
 .../src/translation/mutation/translate.rs     | 150 ++++++++++++++----
 .../goldenfiles/mutations/simple/request.json |  23 ++-
 .../mutations/v1_insert/request.json          |  31 +++-
 .../snapshots/tests__mutations__simple.snap   |  14 +-
 .../tests__mutations__v1_insert.snap          |  11 +-
 ...ests__mutation__delete_playlist_track.snap |   9 +-
 ...ion_tests__basic__delete_invoice_line.snap |  21 ++-
 ...n_tests__basic__delete_playlist_track.snap |  25 ++-
 ...on_tests__basic__v1_insert_custom_dog.snap |  59 +++----
 .../mutations/delete_invoice_line.json        |  31 +++-
 .../mutations/delete_playlist_track.json      |  23 ++-
 .../mutations/insert_artist_album.json        |  68 +++++---
 .../mutations/insert_artist_album_bad.json    |  86 +++++++---
 .../mutations/v1_insert_custom_dog.json       | 126 ++++++++++-----
 .../v1_insert_custom_dog_missing_column.json  |  63 +++++---
 17 files changed, 557 insertions(+), 251 deletions(-)

diff --git a/crates/query-engine/sql/src/sql/helpers.rs b/crates/query-engine/sql/src/sql/helpers.rs
index 3d1941a3..c55230db 100644
--- a/crates/query-engine/sql/src/sql/helpers.rs
+++ b/crates/query-engine/sql/src/sql/helpers.rs
@@ -401,15 +401,22 @@ pub fn select_mutation_rowset(
     (output_table_alias, output_column_alias): (TableAlias, ColumnAlias),
     (row_table_alias, row_column_alias): (TableAlias, ColumnAlias),
     aggregate_table_alias: TableAlias,
-    row_select: Select,
-    aggregate_select: Option<Select>,
+    select: SelectSet,
 ) -> Select {
     let row = vec![(
         output_column_alias,
         Expression::JsonBuildObject(BTreeMap::from([
-            ("type".to_string(), Box::new(Expression::Value(Value::String("procedure".to_string())))),
-            ("result".to_string(), Box::new(Expression::RowToJson(TableReference::AliasedTable(output_table_alias.clone()))))
-        ]))
+            (
+                "type".to_string(),
+                Box::new(Expression::Value(Value::String("procedure".to_string()))),
+            ),
+            (
+                "result".to_string(),
+                Box::new(Expression::RowToJson(TableReference::AliasedTable(
+                    output_table_alias.clone(),
+                ))),
+            ),
+        ])),
     )];
 
     let mut final_select = simple_select(row);
@@ -418,25 +425,43 @@ pub fn select_mutation_rowset(
         select_rows_as_json_for_mutation(row_sel, row_column_alias, row_table_alias.clone())
     };
 
-    let mut select_star = star_select(From::Select {
-        alias: row_table_alias.clone(),
-        select: Box::new(wrap_row(row_select)),
-    });
+    match select {
+        SelectSet::Rows(row_select) => {
+            let select_star = star_select(From::Select {
+                alias: row_table_alias.clone(),
+                select: Box::new(wrap_row(row_select)),
+            });
+
+            final_select.from = Some(From::Select {
+                alias: output_table_alias,
+                select: Box::new(select_star),
+            });
+        }
+
+        SelectSet::Aggregates(aggregate_select) => {
+            final_select.from = Some(From::Select {
+                alias: output_table_alias,
+                select: Box::new(aggregate_select),
+            });
+        }
+
+        SelectSet::RowsAndAggregates(row_select, aggregate_select) => {
+            let mut select_star = star_select(From::Select {
+                alias: row_table_alias.clone(),
+                select: Box::new(wrap_row(row_select)),
+            });
 
-    match aggregate_select {
-        None => {}
-        Some(aggregate_select) => {
             select_star.joins = vec![Join::CrossJoin(CrossJoin {
                 select: Box::new(aggregate_select),
                 alias: aggregate_table_alias.clone(),
             })];
-        }
-    }
 
-    final_select.from = Some(From::Select {
-        alias: output_table_alias,
-        select: Box::new(select_star),
-    });
+            final_select.from = Some(From::Select {
+                alias: output_table_alias,
+                select: Box::new(select_star),
+            });
+        }
+    };
 
     final_select
 }
diff --git a/crates/query-engine/translation/src/translation/error.rs b/crates/query-engine/translation/src/translation/error.rs
index 62cd313e..6791c90d 100644
--- a/crates/query-engine/translation/src/translation/error.rs
+++ b/crates/query-engine/translation/src/translation/error.rs
@@ -28,6 +28,8 @@ pub enum Error {
     ColumnIsIdentityAlways(String),
     MissingColumnInInsert(String, String),
     NotImplementedYet(String),
+    NoProcedureResultFieldsRequested,
+    UnexpectedStructure(String),
     InternalError(String),
 }
 
@@ -119,6 +121,11 @@ impl std::fmt::Display for Error {
             Error::NotImplementedYet(thing) => {
                 write!(f, "Queries containing {} are not supported.", thing)
             }
+            Error::NoProcedureResultFieldsRequested => write!(
+                f,
+                "Procedure requests must ask for 'affected_rows' or use the 'returning' clause."
+            ),
+            Error::UnexpectedStructure(structure) => write!(f, "Unexpected {}.", structure),
             Error::InternalError(thing) => {
                 write!(f, "Internal error: {}.", thing)
             }
diff --git a/crates/query-engine/translation/src/translation/mutation/translate.rs b/crates/query-engine/translation/src/translation/mutation/translate.rs
index 0da6644e..5b3fc837 100644
--- a/crates/query-engine/translation/src/translation/mutation/translate.rs
+++ b/crates/query-engine/translation/src/translation/mutation/translate.rs
@@ -64,16 +64,49 @@ fn translate_mutation(
         alias: table_alias.clone(),
     };
 
-    let fields = match fields {
-        Some(models::NestedField::Object(models::NestedObject{ fields })) => Ok(Some(fields)),
-        Some(models::NestedField::Array(_)) => Err(Error::NotImplementedYet("nested array fields".to_string())),
-        None => Ok(None)
+    let (aggregates, fields) = match fields {
+        Some(models::NestedField::Object(models::NestedObject { fields })) => {
+            let affected_rows = fields
+                .get("affected_rows")
+                .map(|_| indexmap!("affected_rows".to_string() => models::Aggregate::StarCount{}));
+
+            let returning = match fields.get("returning") {
+                Some(field) => match field {
+                    models::Field::Column { fields, .. } => match fields {
+                        Some(nested_fields) => match nested_fields {
+                            models::NestedField::Object(models::NestedObject { fields }) => {
+                                Some(fields.clone())
+                            }
+                            models::NestedField::Array(_) => {
+                                Err(Error::NotImplementedYet("nested array fields".to_string()))?
+                            }
+                        },
+                        None => None,
+                    },
+                    models::Field::Relationship { .. } => Err(Error::UnexpectedStructure(
+                        "relationship variant in procedure fields".to_string(),
+                    ))?,
+                },
+                None => None,
+            };
+
+            if affected_rows.is_none() && returning.is_none() {
+                Err(Error::NoProcedureResultFieldsRequested)?
+            }
+
+            Ok((affected_rows, returning))
+        }
+
+        Some(models::NestedField::Array(_)) => {
+            Err(Error::NotImplementedYet("nested array fields".to_string()))
+        }
+        None => Err(Error::NoProcedureResultFieldsRequested)?,
     }?;
 
     // define the query selecting from the native query,
     // selecting the affected_rows as aggregate and the fields.
     let query = models::Query {
-        aggregates: Some(indexmap!("affected_rows".to_string() => models::Aggregate::StarCount{})),
+        aggregates,
         fields,
         limit: None,
         offset: None,
@@ -113,13 +146,21 @@ fn translate_mutation(
     };
 
     // fields
-    let (_, returning_select) = crate::translation::query::root::translate_rows_query(
-        env,
-        &mut state,
-        &current_table,
-        &from_clause,
-        &query,
-    )?;
+    let returning_select = match &query.fields {
+        Some(_) => {
+            let (_, returning_select) = crate::translation::query::root::translate_rows_query(
+                env,
+                &mut state,
+                &current_table,
+                &from_clause,
+                &query,
+            )?;
+
+            Some(returning_select)
+        }
+
+        None => None,
+    };
 
     // affected rows
     let aggregate_select = crate::translation::query::root::translate_aggregate_query(
@@ -141,8 +182,14 @@ fn translate_mutation(
             sql::helpers::make_column_alias("returning".to_string()),
         ),
         state.make_table_alias("aggregates".to_string()),
-        returning_select,
-        aggregate_select,
+        match (returning_select, aggregate_select) {
+            (Some(returning_select), None) => sql::helpers::SelectSet::Rows(returning_select),
+            (None, Some(aggregate_select)) => sql::helpers::SelectSet::Aggregates(aggregate_select),
+            (Some(returning_select), Some(aggregate_select)) => {
+                sql::helpers::SelectSet::RowsAndAggregates(returning_select, aggregate_select)
+            }
+            (None, None) => Err(Error::NoProcedureResultFieldsRequested)?,
+        },
     );
 
     let common_table_expression = sql::ast::CommonTableExpression {
@@ -200,16 +247,49 @@ fn translate_native_query(
         alias: table_alias.clone(),
     };
 
-    let fields = match fields {
-        Some(models::NestedField::Object(models::NestedObject{ fields })) => Ok(Some(fields)),
-        Some(models::NestedField::Array(_)) => Err(Error::NotImplementedYet("nested array fields".to_string())),
-        None => Ok(None)
+    let (aggregates, fields) = match fields {
+        Some(models::NestedField::Object(models::NestedObject { fields })) => {
+            let affected_rows = fields
+                .get("affected_rows")
+                .map(|_| indexmap!("affected_rows".to_string() => models::Aggregate::StarCount{}));
+
+            let returning = match fields.get("returning") {
+                Some(field) => match field {
+                    models::Field::Column { fields, .. } => match fields {
+                        Some(nested_fields) => match nested_fields {
+                            models::NestedField::Object(models::NestedObject { fields }) => {
+                                Some(fields.clone())
+                            }
+                            models::NestedField::Array(_) => {
+                                Err(Error::NotImplementedYet("nested array fields".to_string()))?
+                            }
+                        },
+                        None => None,
+                    },
+                    models::Field::Relationship { .. } => Err(Error::UnexpectedStructure(
+                        "relationship variant in procedure fields".to_string(),
+                    ))?,
+                },
+                None => None,
+            };
+
+            if affected_rows.is_none() && returning.is_none() {
+                Err(Error::NoProcedureResultFieldsRequested)?
+            }
+
+            Ok((affected_rows, returning))
+        }
+
+        Some(models::NestedField::Array(_)) => {
+            Err(Error::NotImplementedYet("nested array fields".to_string()))
+        }
+        None => Err(Error::NoProcedureResultFieldsRequested)?,
     }?;
 
     // define the query selecting from the native query,
     // selecting the affected_rows as aggregate and the fields.
     let query = models::Query {
-        aggregates: Some(indexmap!("affected_rows".to_string() => models::Aggregate::StarCount{})),
+        aggregates,
         fields,
         limit: None,
         offset: None,
@@ -223,13 +303,21 @@ fn translate_native_query(
     };
 
     // fields
-    let (_, returning_select) = crate::translation::query::root::translate_rows_query(
-        env,
-        &mut state,
-        &current_table,
-        &from_clause,
-        &query,
-    )?;
+    let returning_select = match &query.fields {
+        Some(_) => {
+            let (_, returning_select) = crate::translation::query::root::translate_rows_query(
+                env,
+                &mut state,
+                &current_table,
+                &from_clause,
+                &query,
+            )?;
+
+            Some(returning_select)
+        }
+
+        None => None,
+    };
 
     // affected rows
     let aggregate_select = crate::translation::query::root::translate_aggregate_query(
@@ -251,8 +339,14 @@ fn translate_native_query(
             sql::helpers::make_column_alias("returning".to_string()),
         ),
         state.make_table_alias("aggregates".to_string()),
-        returning_select,
-        aggregate_select,
+        match (returning_select, aggregate_select) {
+            (Some(returning_select), None) => sql::helpers::SelectSet::Rows(returning_select),
+            (None, Some(aggregate_select)) => sql::helpers::SelectSet::Aggregates(aggregate_select),
+            (Some(returning_select), Some(aggregate_select)) => {
+                sql::helpers::SelectSet::RowsAndAggregates(returning_select, aggregate_select)
+            }
+            (None, None) => Err(Error::NoProcedureResultFieldsRequested)?,
+        },
     );
 
     // add the procedure native query definition is a with clause.
diff --git a/crates/query-engine/translation/tests/goldenfiles/mutations/simple/request.json b/crates/query-engine/translation/tests/goldenfiles/mutations/simple/request.json
index 9a2aa5a1..da92a2a1 100644
--- a/crates/query-engine/translation/tests/goldenfiles/mutations/simple/request.json
+++ b/crates/query-engine/translation/tests/goldenfiles/mutations/simple/request.json
@@ -7,9 +7,26 @@
         "track_id": 90
       },
       "fields": {
-        "playlist_id": {
-          "type": "column",
-          "column": "PlaylistId"
+        "type": "object",
+        "fields": {
+          "affected_rows": {
+            "column": "affected_rows",
+            "type": "column"
+          },
+
+          "returning": {
+            "type": "column",
+            "column": "returning",
+            "fields": {
+              "type": "object",
+              "fields": {
+                "playlist_id": {
+                  "type": "column",
+                  "column": "PlaylistId"
+                }
+              }
+            }
+          }
         }
       }
     }
diff --git a/crates/query-engine/translation/tests/goldenfiles/mutations/v1_insert/request.json b/crates/query-engine/translation/tests/goldenfiles/mutations/v1_insert/request.json
index b644e5f0..6ca84989 100644
--- a/crates/query-engine/translation/tests/goldenfiles/mutations/v1_insert/request.json
+++ b/crates/query-engine/translation/tests/goldenfiles/mutations/v1_insert/request.json
@@ -10,13 +10,30 @@
         }
       },
       "fields": {
-        "artist_id": {
-          "type": "column",
-          "column": "id"
-        },
-        "name": {
-          "type": "column",
-          "column": "name"
+        "type": "object",
+        "fields": {
+          "affected_rows": {
+            "column": "affected_rows",
+            "type": "column"
+          },
+
+          "returning": {
+            "type": "column",
+            "column": "returning",
+            "fields": {
+              "type": "object",
+              "fields": {
+                "artist_id": {
+                  "type": "column",
+                  "column": "id"
+                },
+                "name": {
+                  "type": "column",
+                  "column": "name"
+                }
+              }
+            }
+          }
         }
       }
     }
diff --git a/crates/query-engine/translation/tests/snapshots/tests__mutations__simple.snap b/crates/query-engine/translation/tests/snapshots/tests__mutations__simple.snap
index 87a07fad..89cce2c7 100644
--- a/crates/query-engine/translation/tests/snapshots/tests__mutations__simple.snap
+++ b/crates/query-engine/translation/tests/snapshots/tests__mutations__simple.snap
@@ -12,10 +12,7 @@ WITH "%0_NATIVE_QUERY_delete_playlist_track" AS (
     "TrackId" = 90 RETURNING *
 )
 SELECT
-  json_build_object(
-    "result", row_to_json("%2_universe"),
-    "type", "Procedure"
-  ) as "universe"
+  json_build_object('result', row_to_json("%2_universe"), 'type', $1) AS "universe"
 FROM
   (
     SELECT
@@ -23,12 +20,7 @@ FROM
     FROM
       (
         SELECT
-          json_build_array(
-            json_build_object(
-              '__value',
-              coalesce(json_agg(row_to_json("%3_returning")), '[]')
-            )
-          ) AS "returning"
+          coalesce(json_agg(row_to_json("%3_returning")), '[]') AS "returning"
         FROM
           (
             SELECT
@@ -47,4 +39,4 @@ FROM
 
 COMMIT;
 
-[[]]
+[[(1, String("procedure"))]]
diff --git a/crates/query-engine/translation/tests/snapshots/tests__mutations__v1_insert.snap b/crates/query-engine/translation/tests/snapshots/tests__mutations__v1_insert.snap
index 196e671f..36455e23 100644
--- a/crates/query-engine/translation/tests/snapshots/tests__mutations__v1_insert.snap
+++ b/crates/query-engine/translation/tests/snapshots/tests__mutations__v1_insert.snap
@@ -12,7 +12,7 @@ WITH "%0_generated_mutation" AS (
     (276, cast($1 as varchar)) RETURNING *
 )
 SELECT
-  row_to_json("%2_universe") AS "universe"
+  json_build_object('result', row_to_json("%2_universe"), 'type', $2) AS "universe"
 FROM
   (
     SELECT
@@ -20,12 +20,7 @@ FROM
     FROM
       (
         SELECT
-          json_build_array(
-            json_build_object(
-              '__value',
-              coalesce(json_agg(row_to_json("%3_returning")), '[]')
-            )
-          ) AS "returning"
+          coalesce(json_agg(row_to_json("%3_returning")), '[]') AS "returning"
         FROM
           (
             SELECT
@@ -45,4 +40,4 @@ FROM
 
 COMMIT;
 
-[[(1, String("Olympians"))]]
+[[(1, String("Olympians")), (2, String("procedure"))]]
diff --git a/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__explain_tests__mutation__delete_playlist_track.snap b/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__explain_tests__mutation__delete_playlist_track.snap
index 073749c4..d5ec8a9b 100644
--- a/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__explain_tests__mutation__delete_playlist_track.snap
+++ b/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__explain_tests__mutation__delete_playlist_track.snap
@@ -9,7 +9,7 @@ EXPLAIN WITH "%0_NATIVE_QUERY_delete_playlist_track" AS (
     "TrackId" = 90 RETURNING *
 )
 SELECT
-  row_to_json("%2_universe") AS "universe"
+  json_build_object('result', row_to_json("%2_universe"), 'type', $1) AS "universe"
 FROM
   (
     SELECT
@@ -17,12 +17,7 @@ FROM
     FROM
       (
         SELECT
-          json_build_array(
-            json_build_object(
-              '__value',
-              coalesce(json_agg(row_to_json("%3_returning")), '[]')
-            )
-          ) AS "returning"
+          coalesce(json_agg(row_to_json("%3_returning")), '[]') AS "returning"
         FROM
           (
             SELECT
diff --git a/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__mutation_tests__basic__delete_invoice_line.snap b/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__mutation_tests__basic__delete_invoice_line.snap
index a45af1ac..a9067c9c 100644
--- a/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__mutation_tests__basic__delete_invoice_line.snap
+++ b/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__mutation_tests__basic__delete_invoice_line.snap
@@ -5,17 +5,16 @@ expression: result
 {
   "operation_results": [
     {
-      "affected_rows": 1,
-      "returning": [
-        {
-          "__value": [
-            {
-              "invoice_line_id": 90,
-              "quantity": 1
-            }
-          ]
-        }
-      ]
+      "type": "procedure",
+      "result": {
+        "returning": [
+          {
+            "invoice_line_id": 90,
+            "quantity": 1
+          }
+        ],
+        "affected_rows": 1
+      }
     }
   ]
 }
diff --git a/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__mutation_tests__basic__delete_playlist_track.snap b/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__mutation_tests__basic__delete_playlist_track.snap
index f0196ce5..d246539e 100644
--- a/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__mutation_tests__basic__delete_playlist_track.snap
+++ b/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__mutation_tests__basic__delete_playlist_track.snap
@@ -5,19 +5,18 @@ expression: result
 {
   "operation_results": [
     {
-      "affected_rows": 2,
-      "returning": [
-        {
-          "__value": [
-            {
-              "playlist_id": 1
-            },
-            {
-              "playlist_id": 8
-            }
-          ]
-        }
-      ]
+      "type": "procedure",
+      "result": {
+        "returning": [
+          {
+            "playlist_id": 1
+          },
+          {
+            "playlist_id": 8
+          }
+        ],
+        "affected_rows": 2
+      }
     }
   ]
 }
diff --git a/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__mutation_tests__basic__v1_insert_custom_dog.snap b/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__mutation_tests__basic__v1_insert_custom_dog.snap
index bcf8b78b..b4debd69 100644
--- a/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__mutation_tests__basic__v1_insert_custom_dog.snap
+++ b/crates/tests/databases-tests/src/postgres/snapshots/databases_tests__postgres__mutation_tests__basic__v1_insert_custom_dog.snap
@@ -5,36 +5,39 @@ expression: result
 {
   "operation_results": [
     {
-      "affected_rows": 1,
-      "returning": [
-        {
-          "__value": [
-            {
-              "id": 1,
-              "name": "Cremebo",
-              "birthday": "2024-01-17",
-              "height_cm": 160,
-              "height_inch": 62.99212598425197,
-              "adopter_name": null
-            }
-          ]
-        }
-      ]
+      "type": "procedure",
+      "result": {
+        "returning": [
+          {
+            "id": 1,
+            "name": "Cremebo",
+            "birthday": "2024-01-17",
+            "height_cm": 160,
+            "height_inch": 62.99212598425197,
+            "adopter_name": null
+          }
+        ],
+        "affected_rows": 1
+      }
     },
     {
-      "affected_rows": 1,
-      "returning": [
-        {
-          "__value": [
-            {
-              "name": "Manbo",
-              "birthday": "2024-01-01",
-              "height_inch": 135.8267716535433,
-              "adopter_name": "Joseph"
-            }
-          ]
-        }
-      ]
+      "type": "procedure",
+      "result": {
+        "returning": [
+          {
+            "name": "Manbo",
+            "birthday": "2024-01-01",
+            "height_inch": 135.8267716535433,
+            "adopter_name": "Joseph"
+          }
+        ]
+      }
+    },
+    {
+      "type": "procedure",
+      "result": {
+        "affected_rows": 1
+      }
     }
   ]
 }
diff --git a/crates/tests/tests-common/goldenfiles/mutations/delete_invoice_line.json b/crates/tests/tests-common/goldenfiles/mutations/delete_invoice_line.json
index 8530cbf4..9e7ec8ab 100644
--- a/crates/tests/tests-common/goldenfiles/mutations/delete_invoice_line.json
+++ b/crates/tests/tests-common/goldenfiles/mutations/delete_invoice_line.json
@@ -7,13 +7,30 @@
         "InvoiceLineId": 90
       },
       "fields": {
-        "invoice_line_id": {
-          "type": "column",
-          "column": "InvoiceLineId"
-        },
-        "quantity": {
-          "type": "column",
-          "column": "Quantity"
+        "type": "object",
+        "fields": {
+          "affected_rows": {
+            "column": "affected_rows",
+            "type": "column"
+          },
+
+          "returning": {
+            "type": "column",
+            "column": "returning",
+            "fields": {
+              "type": "object",
+              "fields": {
+                "invoice_line_id": {
+                  "type": "column",
+                  "column": "InvoiceLineId"
+                },
+                "quantity": {
+                  "type": "column",
+                  "column": "Quantity"
+                }
+              }
+            }
+          }
         }
       }
     }
diff --git a/crates/tests/tests-common/goldenfiles/mutations/delete_playlist_track.json b/crates/tests/tests-common/goldenfiles/mutations/delete_playlist_track.json
index 9a2aa5a1..da92a2a1 100644
--- a/crates/tests/tests-common/goldenfiles/mutations/delete_playlist_track.json
+++ b/crates/tests/tests-common/goldenfiles/mutations/delete_playlist_track.json
@@ -7,9 +7,26 @@
         "track_id": 90
       },
       "fields": {
-        "playlist_id": {
-          "type": "column",
-          "column": "PlaylistId"
+        "type": "object",
+        "fields": {
+          "affected_rows": {
+            "column": "affected_rows",
+            "type": "column"
+          },
+
+          "returning": {
+            "type": "column",
+            "column": "returning",
+            "fields": {
+              "type": "object",
+              "fields": {
+                "playlist_id": {
+                  "type": "column",
+                  "column": "PlaylistId"
+                }
+              }
+            }
+          }
         }
       }
     }
diff --git a/crates/tests/tests-common/goldenfiles/mutations/insert_artist_album.json b/crates/tests/tests-common/goldenfiles/mutations/insert_artist_album.json
index 8aca8819..9fdfd73e 100644
--- a/crates/tests/tests-common/goldenfiles/mutations/insert_artist_album.json
+++ b/crates/tests/tests-common/goldenfiles/mutations/insert_artist_album.json
@@ -10,13 +10,27 @@
       "fields": {
         "type": "object",
         "fields": {
-          "artist_id": {
-            "type": "column",
-            "column": "ArtistId"
+          "affected_rows": {
+            "column": "affected_rows",
+            "type": "column"
           },
-          "name": {
+
+          "returning": {
             "type": "column",
-            "column": "Name"
+            "column": "returning",
+            "fields": {
+              "type": "object",
+              "fields": {
+                "artist_id": {
+                  "type": "column",
+                  "column": "ArtistId"
+                },
+                "name": {
+                  "type": "column",
+                  "column": "Name"
+                }
+              }
+            }
           }
         }
       }
@@ -32,28 +46,42 @@
       "fields": {
         "type": "object",
         "fields": {
-          "album_id": {
-            "type": "column",
-            "column": "AlbumId"
+          "affected_rows": {
+            "column": "affected_rows",
+            "type": "column"
           },
-          "title": {
+
+          "returning": {
             "type": "column",
-            "column": "Title"
-          },
-          "artist": {
-            "type": "relationship",
-            "column": "Title",
-            "relationship": "AlbumToArtist",
-            "query": {
+            "column": "returning",
+            "fields": {
+              "type": "object",
               "fields": {
-                "name": {
+                "album_id": {
+                  "type": "column",
+                  "column": "AlbumId"
+                },
+                "title": {
                   "type": "column",
-                  "column": "Name",
+                  "column": "Title"
+                },
+                "artist": {
+                  "type": "relationship",
+                  "column": "Title",
+                  "relationship": "AlbumToArtist",
+                  "query": {
+                    "fields": {
+                      "name": {
+                        "type": "column",
+                        "column": "Name",
+                        "arguments": {}
+                      }
+                    }
+                  },
                   "arguments": {}
                 }
               }
-            },
-            "arguments": {}
+            }
           }
         }
       }
diff --git a/crates/tests/tests-common/goldenfiles/mutations/insert_artist_album_bad.json b/crates/tests/tests-common/goldenfiles/mutations/insert_artist_album_bad.json
index 1b143090..a997b1ad 100644
--- a/crates/tests/tests-common/goldenfiles/mutations/insert_artist_album_bad.json
+++ b/crates/tests/tests-common/goldenfiles/mutations/insert_artist_album_bad.json
@@ -8,13 +8,30 @@
         "name": "Olympians"
       },
       "fields": {
-        "artist_id": {
-          "type": "column",
-          "column": "ArtistId"
-        },
-        "name": {
-          "type": "column",
-          "column": "Name"
+        "type": "object",
+        "fields": {
+          "affected_rows": {
+            "column": "affected_rows",
+            "type": "column"
+          },
+
+          "returning": {
+            "type": "column",
+            "column": "returning",
+            "fields": {
+              "type": "object",
+              "fields": {
+                "artist_id": {
+                  "type": "column",
+                  "column": "ArtistId"
+                },
+                "name": {
+                  "type": "column",
+                  "column": "Name"
+                }
+              }
+            }
+          }
         }
       }
     },
@@ -27,28 +44,45 @@
         "artist_id": 276
       },
       "fields": {
-        "album_id": {
-          "type": "column",
-          "column": "AlbumId"
-        },
-        "title": {
-          "type": "column",
-          "column": "Title"
-        },
-        "artist": {
-          "type": "relationship",
-          "column": "Title",
-          "relationship": "AlbumToArtist",
-          "query": {
+        "type": "object",
+        "fields": {
+          "affected_rows": {
+            "column": "affected_rows",
+            "type": "column"
+          },
+
+          "returning": {
+            "type": "column",
+            "column": "returning",
             "fields": {
-              "name": {
-                "type": "column",
-                "column": "Name",
-                "arguments": {}
+              "type": "object",
+              "fields": {
+                "album_id": {
+                  "type": "column",
+                  "column": "AlbumId"
+                },
+                "title": {
+                  "type": "column",
+                  "column": "Title"
+                },
+                "artist": {
+                  "type": "relationship",
+                  "column": "Title",
+                  "relationship": "AlbumToArtist",
+                  "query": {
+                    "fields": {
+                      "name": {
+                        "type": "column",
+                        "column": "Name",
+                        "arguments": {}
+                      }
+                    }
+                  },
+                  "arguments": {}
+                }
               }
             }
-          },
-          "arguments": {}
+          }
         }
       }
     }
diff --git a/crates/tests/tests-common/goldenfiles/mutations/v1_insert_custom_dog.json b/crates/tests/tests-common/goldenfiles/mutations/v1_insert_custom_dog.json
index 99d1c6fd..021d0b43 100644
--- a/crates/tests/tests-common/goldenfiles/mutations/v1_insert_custom_dog.json
+++ b/crates/tests/tests-common/goldenfiles/mutations/v1_insert_custom_dog.json
@@ -10,29 +10,46 @@
         }
       },
       "fields": {
-        "id": {
-          "type": "column",
-          "column": "id"
-        },
-        "name": {
-          "type": "column",
-          "column": "name"
-        },
-        "birthday": {
-          "type": "column",
-          "column": "birthday"
-        },
-        "height_cm": {
-          "type": "column",
-          "column": "height_cm"
-        },
-        "height_inch": {
-          "type": "column",
-          "column": "height_in"
-        },
-        "adopter_name": {
-          "type": "column",
-          "column": "adopter_name"
+        "type": "object",
+        "fields": {
+          "affected_rows": {
+            "column": "affected_rows",
+            "type": "column"
+          },
+
+          "returning": {
+            "type": "column",
+            "column": "returning",
+            "fields": {
+              "type": "object",
+              "fields": {
+                "id": {
+                  "type": "column",
+                  "column": "id"
+                },
+                "name": {
+                  "type": "column",
+                  "column": "name"
+                },
+                "birthday": {
+                  "type": "column",
+                  "column": "birthday"
+                },
+                "height_cm": {
+                  "type": "column",
+                  "column": "height_cm"
+                },
+                "height_inch": {
+                  "type": "column",
+                  "column": "height_in"
+                },
+                "adopter_name": {
+                  "type": "column",
+                  "column": "adopter_name"
+                }
+              }
+            }
+          }
         }
       }
     },
@@ -48,21 +65,54 @@
         }
       },
       "fields": {
-        "name": {
-          "type": "column",
-          "column": "name"
-        },
-        "birthday": {
-          "type": "column",
-          "column": "birthday"
-        },
-        "height_inch": {
-          "type": "column",
-          "column": "height_in"
-        },
-        "adopter_name": {
-          "type": "column",
-          "column": "adopter_name"
+        "type": "object",
+        "fields": {
+          "returning": {
+            "type": "column",
+            "column": "returning",
+            "fields": {
+              "type": "object",
+              "fields": {
+                "name": {
+                  "type": "column",
+                  "column": "name"
+                },
+                "birthday": {
+                  "type": "column",
+                  "column": "birthday"
+                },
+                "height_inch": {
+                  "type": "column",
+                  "column": "height_in"
+                },
+                "adopter_name": {
+                  "type": "column",
+                  "column": "adopter_name"
+                }
+              }
+            }
+          }
+        }
+      }
+    },
+    {
+      "type": "procedure",
+      "name": "v1_insert_custom_dog",
+      "arguments": {
+        "_object": {
+          "name": "Kremboli",
+          "height_cm": 345,
+          "adopter_name": "Yaakov",
+          "birthday": "2024-02-14"
+        }
+      },
+      "fields": {
+        "type": "object",
+        "fields": {
+          "affected_rows": {
+            "column": "affected_rows",
+            "type": "column"
+          }
         }
       }
     }
diff --git a/crates/tests/tests-common/goldenfiles/mutations/v1_insert_custom_dog_missing_column.json b/crates/tests/tests-common/goldenfiles/mutations/v1_insert_custom_dog_missing_column.json
index 7be1a59c..da529003 100644
--- a/crates/tests/tests-common/goldenfiles/mutations/v1_insert_custom_dog_missing_column.json
+++ b/crates/tests/tests-common/goldenfiles/mutations/v1_insert_custom_dog_missing_column.json
@@ -9,29 +9,46 @@
         }
       },
       "fields": {
-        "id": {
-          "type": "column",
-          "column": "id"
-        },
-        "name": {
-          "type": "column",
-          "column": "name"
-        },
-        "birthday": {
-          "type": "column",
-          "column": "birthday"
-        },
-        "height_cm": {
-          "type": "column",
-          "column": "height_cm"
-        },
-        "height_inch": {
-          "type": "column",
-          "column": "height_in"
-        },
-        "adopter_name": {
-          "type": "column",
-          "column": "adopter_name"
+        "type": "object",
+        "fields": {
+          "affected_rows": {
+            "column": "affected_rows",
+            "type": "column"
+          },
+
+          "returning": {
+            "type": "column",
+            "column": "returning",
+            "fields": {
+              "type": "object",
+              "fields": {
+                "id": {
+                  "type": "column",
+                  "column": "id"
+                },
+                "name": {
+                  "type": "column",
+                  "column": "name"
+                },
+                "birthday": {
+                  "type": "column",
+                  "column": "birthday"
+                },
+                "height_cm": {
+                  "type": "column",
+                  "column": "height_cm"
+                },
+                "height_inch": {
+                  "type": "column",
+                  "column": "height_in"
+                },
+                "adopter_name": {
+                  "type": "column",
+                  "column": "adopter_name"
+                }
+              }
+            }
+          }
         }
       }
     }