From 7cf1695b8011cc83f60edc7cb01ceea67ecedc63 Mon Sep 17 00:00:00 2001
From: Tom Harding <i.am.tom.harding@gmail.com>
Date: Thu, 15 Feb 2024 15:18:04 +0100
Subject: [PATCH] More refactoring

---
 crates/query-engine/sql/src/sql/helpers.rs    |  6 +++-
 .../src/translation/mutation/translate.rs     | 36 ++++++++++---------
 2 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/crates/query-engine/sql/src/sql/helpers.rs b/crates/query-engine/sql/src/sql/helpers.rs
index c55230db..3f2d99a8 100644
--- a/crates/query-engine/sql/src/sql/helpers.rs
+++ b/crates/query-engine/sql/src/sql/helpers.rs
@@ -382,7 +382,8 @@ pub fn select_rowset_with_variables(
     final_select
 }
 
-/// Given a set of rows and aggregate queries, combine them into one Select.
+/// Build a `Select` query using a `SelectSet` of row fields and aggregates according to the
+/// following SQL template:
 ///
 /// ```sql
 /// SELECT row_to_json(<output_table_alias>) AS <output_column_alias>
@@ -397,6 +398,9 @@ pub fn select_rowset_with_variables(
 ///     ) AS <aggregate_table_alias>
 /// ) AS <output_table_alias>
 /// ```
+///
+/// The `SelectSet` determines whether we select from both the rows and the aggregates, or just the
+/// rows, or just the aggregates.
 pub fn select_mutation_rowset(
     (output_table_alias, output_column_alias): (TableAlias, ColumnAlias),
     (row_table_alias, row_column_alias): (TableAlias, ColumnAlias),
diff --git a/crates/query-engine/translation/src/translation/mutation/translate.rs b/crates/query-engine/translation/src/translation/mutation/translate.rs
index 1827586c..22af7cb6 100644
--- a/crates/query-engine/translation/src/translation/mutation/translate.rs
+++ b/crates/query-engine/translation/src/translation/mutation/translate.rs
@@ -145,14 +145,7 @@ fn translate_mutation(
             sql::helpers::make_column_alias("returning".to_string()),
         ),
         state.make_table_alias("aggregates".to_string()),
-        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)?,
-        },
+        rows_and_aggregates_to_select_set(returning_select, aggregate_select)?
     );
 
     let common_table_expression = sql::ast::CommonTableExpression {
@@ -175,6 +168,24 @@ fn translate_mutation(
     })
 }
 
+/// Procedures can return a number of affected rows and/or some fields from the rows that are
+/// affected, but it must return at least one. A `SelectSet` describes this as a type, so we can
+/// convert an optional returning `Select` and an optional aggregate `Select` to a `SelectSet`,
+/// failing if neither exists.
+fn rows_and_aggregates_to_select_set(
+    returning_select: Option<sql::ast::Select>,
+    aggregate_select: Option<sql::ast::Select>,
+) -> Result<sql::helpers::SelectSet, Error> {
+    match (returning_select, aggregate_select) {
+        (Some(returning_select), None) => Ok(sql::helpers::SelectSet::Rows(returning_select)),
+        (None, Some(aggregate_select)) => Ok(sql::helpers::SelectSet::Aggregates(aggregate_select)),
+        (Some(returning_select), Some(aggregate_select)) => Ok(
+            sql::helpers::SelectSet::RowsAndAggregates(returning_select, aggregate_select),
+        ),
+        (None, None) => Err(Error::NoProcedureResultFieldsRequested),
+    }
+}
+
 /// Translate a Native Query mutation into an ExecutionPlan (SQL) to be run against the database.
 fn translate_native_query(
     env: &Env,
@@ -265,14 +276,7 @@ fn translate_native_query(
             sql::helpers::make_column_alias("returning".to_string()),
         ),
         state.make_table_alias("aggregates".to_string()),
-        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)?,
-        },
+        rows_and_aggregates_to_select_set(returning_select, aggregate_select)?
     );
 
     // add the procedure native query definition is a with clause.