More refactoring

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),

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.