Invert the update columns syntax for update procedures (#494)

### What

Slack thread for more context:
https://hasurahq.slack.com/archives/C04NS5JCD8A/p1717760957475579

We want to introduce a different syntax for column update. Instead of
bucketing by operations, we bucket by column names.

#### Old API

```json
{ 
  "_set": { "name": "Al", "address": "street"},
  "_inc": { "age": 1}, 
  "_concat": { "..." }
}
```

#### New API

```json
{ 
  "update_columns": { 
    "name": { "_set": "Al" },
    "address": { "_set": "street" },
    "age": { "_inc": 1 }
  }
}
```

This syntax is more consistent with other operations in other places,
and makes it easy to augment the existing update procedure with new
operations without adding arguments.

### How

Warning: this code was written at night, extra care is advised.

1. We change the name of the `_set` argument in the schema to
`update_columns`
2. We introduce a new object type for each column update with the
structure `{ _set: value }`
3. The type for `update_columns` is all the columns in the table mapped
to their object type.
4. We parse the `update_columns` object by parsing each column and its
operation, and generate a `MutationValueExpression`.

## Versioning and changelog

No need to worry about versioning. This is all in the
`veryExperimentalWIP` mutationVersion.

crates/connectors/ndc-postgres/src/schema.rs

@@ -577,32 +577,25 @@ fn experimental_insert_to_procedure(
 
 /// Given an experimental `UpdateMutation`, turn it into a `ProcedureInfo` to be output in the schema.
 fn experimental_update_to_procedure(
-    name: &String,
+    procedure_name: &str,
     update: &mutation::experimental::update::UpdateMutation,
     object_types: &mut BTreeMap<String, models::ObjectType>,
     scalar_types: &mut BTreeMap<String, models::ScalarType>,
 ) -> models::ProcedureInfo {
     let mutation::experimental::update::UpdateMutation::UpdateByKey(update_by_key) = update;
 
     let mut arguments = BTreeMap::new();
-    let object_type = make_object_type(&update_by_key.columns);
-    let object_name = format!("{name}_object");
-    object_types.insert(object_name.clone(), object_type);
 
+    // by column argument.
     arguments.insert(
         update_by_key.by_column.name.clone(),
         models::ArgumentInfo {
             argument_type: column_to_type(&update_by_key.by_column),
             description: update_by_key.by_column.description.clone(),
         },
     );
-    arguments.insert(
-        update_by_key.set_argument_name.clone(),
-        models::ArgumentInfo {
-            argument_type: models::Type::Named { name: object_name },
-            description: None,
-        },
-    );
+
+    // pre check argument.
     arguments.insert(
         update_by_key.pre_check.argument_name.clone(),
         models::ArgumentInfo {
@@ -612,6 +605,8 @@ fn experimental_update_to_procedure(
             description: Some(update_by_key.pre_check.description.clone()),
         },
     );
+
+    // post check argument.
     arguments.insert(
         update_by_key.post_check.argument_name.clone(),
         models::ArgumentInfo {
@@ -622,8 +617,70 @@ fn experimental_update_to_procedure(
         },
     );
 
+    // update columns argument.
+    // Is of the form update_columns: { <column_name>: { <operation>: <value> }, ... }.
+    //
+    // 1. We need to create an object type for each `{ <operation>: <value> }` object
+    // 2. We need to create an object type for the mapping of column to column update object
+    // 3. We need to add the argument to match the name of the object in (2)
+
+    // Keep track of the fields.
+    let mut fields = BTreeMap::new();
+
+    // Make an object type for each column's update object.
+    for (column_name, column_info) in &update_by_key.table_columns {
+        let (object_name, object_type) = make_update_column_type(
+            &update_by_key.collection_name,
+            column_name,
+            column_to_type(column_info),
+        );
+        // add to object types
+        object_types.insert(object_name.clone(), object_type.clone());
+        // Remember for the update_columns type
+        fields.insert(
+            column_name.clone(),
+            models::ObjectField {
+                description: Some(format!(
+                    "Update the '{column_name}' column in the '{}' collection.",
+                    update_by_key.collection_name
+                )),
+                r#type: models::Type::Named { name: object_name },
+                arguments: BTreeMap::new(),
+            },
+        );
+    }
+
+    // Create the update columns object type.
+    let update_columns_object_type_name = format!(
+        "{procedure_name}_{}",
+        update_by_key.update_columns_argument_name
+    );
+
+    object_types.insert(
+        update_columns_object_type_name.clone(),
+        models::ObjectType {
+            description: Some(format!(
+                "Update the columns of the '{}' collection",
+                update_by_key.collection_name
+            )),
+            fields,
+        },
+    );
+
+    // Insert the update columns argument.
+    arguments.insert(
+        update_by_key.update_columns_argument_name.clone(),
+        models::ArgumentInfo {
+            argument_type: models::Type::Named {
+                name: update_columns_object_type_name,
+            },
+            description: None,
+        },
+    );
+
+    // Make a type for the procedure.
     make_procedure_type(
-        name.to_string(),
+        procedure_name.to_string(),
         Some(update_by_key.description.to_string()),
         arguments,
         models::Type::Named {
@@ -705,6 +762,36 @@ fn make_procedure_type(
     }
 }
 
+/// Build an `ObjectType` for an update column.
+fn make_update_column_type(
+    collection_name: &str,
+    column_name: &str,
+    column_type: models::Type,
+) -> (String, models::ObjectType) {
+    let mut fields = BTreeMap::new();
+    let object_type_name = format!("update_column_{collection_name}_{column_name}");
+
+    // Right now we only support set
+    fields.insert(
+        "_set".to_string(),
+        models::ObjectField {
+            description: Some("Set the column to this value".to_string()),
+            r#type: column_type,
+            arguments: BTreeMap::new(),
+        },
+    );
+
+    (
+        object_type_name.clone(),
+        models::ObjectType {
+            description: Some(format!(
+                "Update the '{column_name}' column in the '{collection_name}' collection"
+            )),
+            fields,
+        },
+    )
+}
+
 /// Map our local type representation to ndc-spec type representation.
 #[allow(clippy::match_same_arms)] // merging arms would require changing the order, making this harder to understand
 fn map_type_representation(

crates/query-engine/translation/src/translation/error.rs

@@ -35,6 +35,11 @@ pub enum Error {
     NotImplementedYet(String),
     NoProcedureResultFieldsRequested,
     UnexpectedStructure(String),
+    UnexpectedOperation {
+        column_name: String,
+        operation: String,
+        available_operations: Vec<String>,
+    },
     InternalError(String),
     NestedArrayTypesNotSupported,
     NestedArraysNotSupported {
@@ -146,6 +151,22 @@ impl std::fmt::Display for Error {
                 f,
                 "Procedure requests must ask for 'affected_rows' or use the 'returning' clause."
             ),
+            Error::UnexpectedOperation {
+                column_name,
+                operation,
+                available_operations,
+            } => {
+                let mut string = format!(
+                    "Unexpected operation '{operation}' on column {column_name}. Expected one of: "
+                );
+                for (index, op) in available_operations.iter().enumerate() {
+                    string.push_str(op);
+                    if index < available_operations.len() - 1 {
+                        string.push_str(", ");
+                    }
+                }
+                write!(f, "{string}")
+            }
             Error::UnexpectedStructure(structure) => write!(f, "Unexpected {structure}."),
             Error::InternalError(thing) => {
                 write!(f, "Internal error: {thing}.")

crates/query-engine/translation/src/translation/mutation/experimental/update.rs

@@ -28,10 +28,10 @@ pub struct UpdateByKey {
     pub schema_name: sql::ast::SchemaName,
     pub table_name: sql::ast::TableName,
     pub by_column: metadata::database::ColumnInfo,
-    pub set_argument_name: String,
+    pub update_columns_argument_name: String,
     pub pre_check: Constraint,
     pub post_check: Constraint,
-    pub columns: BTreeMap<String, metadata::database::ColumnInfo>,
+    pub table_columns: BTreeMap<String, metadata::database::ColumnInfo>,
 }
 
 /// The name and description of the constraint input argument.
@@ -65,7 +65,7 @@ pub fn generate_update_by_unique(
                 table_name: sql::ast::TableName(table_info.table_name.clone()),
                 collection_name: collection_name.clone(),
                 by_column: unique_column.clone(),
-                set_argument_name: "_set".to_string(),
+                update_columns_argument_name: "update_columns".to_string(),
                 pre_check: Constraint {
                     argument_name: "pre_check".to_string(),
                     description: format!(
@@ -78,7 +78,7 @@ pub fn generate_update_by_unique(
                 "Update permission post-condition predicate over the '{collection_name}' collection"
             ),
                 },
-                columns: table_info.columns.clone(),
+                table_columns: table_info.columns.clone(),
 
                 description,
             });
@@ -99,10 +99,12 @@ pub fn translate(
     match mutation {
         UpdateMutation::UpdateByKey(mutation) => {
             let object = arguments
-                .get(&mutation.set_argument_name)
-                .ok_or(Error::ArgumentNotFound("_set".to_string()))?;
+                .get(&mutation.update_columns_argument_name)
+                .ok_or(Error::ArgumentNotFound(
+                    mutation.update_columns_argument_name.clone(),
+                ))?;
 
-            let set = parse_set(env, state, mutation, object)?;
+            let set = parse_update_columns(env, state, mutation, object)?;
 
             let table_name_and_reference = TableNameAndReference {
                 name: mutation.collection_name.clone(),
@@ -204,7 +206,7 @@ pub fn translate(
 }
 
 /// Translate a single update object into a mapping from column names to values.
-fn parse_set(
+fn parse_update_columns(
     env: &crate::translation::helpers::Env,
     state: &mut crate::translation::helpers::State,
     mutation: &UpdateByKey,
@@ -214,11 +216,12 @@ fn parse_set(
 
     match object {
         serde_json::Value::Object(object) => {
-            // For each field, look up the column name in the table and update it and the value into the map.
+            // For each field, look up the column name in the table
+            // and update it and the value into the map.
             for (name, value) in object {
                 let column_info =
                     mutation
-                        .columns
+                        .table_columns
                         .get(name)
                         .ok_or(Error::ColumnNotFoundInCollection(
                             name.clone(),
@@ -227,30 +230,79 @@ fn parse_set(
 
                 columns_to_values.insert(
                     sql::ast::ColumnName(column_info.name.clone()),
-                    sql::ast::MutationValueExpression::Expression(translate_json_value(
-                        env,
-                        state,
-                        value,
-                        &column_info.r#type,
-                    )?),
+                    parse_update_column(env, state, name, column_info, value)?,
                 );
             }
             Ok(())
         }
-        serde_json::Value::Array(_) => Err(Error::UnexpectedStructure(
-            "array structure in update _set argument. Expecting an object.".to_string(),
-        )),
-        _ => Err(Error::UnexpectedStructure(
-            "value structure in update _set argument. Expecting an object.".to_string(),
-        )),
+        serde_json::Value::Array(_) => Err(Error::UnexpectedStructure(format!(
+            "array structure in update '{}' argument. Expecting an object.",
+            mutation.update_columns_argument_name
+        ))),
+        _ => Err(Error::UnexpectedStructure(format!(
+            "value structure in update '{}' argument. Expecting an object.",
+            mutation.update_columns_argument_name
+        ))),
     }?;
 
     check_columns::check_columns(
-        &mutation.columns,
+        &mutation.table_columns,
         &columns_to_values,
         &mutation.collection_name,
         &check_columns::CheckMissingColumns::No,
     )?;
 
     Ok(columns_to_values)
 }
+
+/// Translate the operation object of a column to a mutation value expression.
+fn parse_update_column(
+    env: &crate::translation::helpers::Env,
+    state: &mut crate::translation::helpers::State,
+    column_name: &str,
+    column_info: &metadata::database::ColumnInfo,
+    object: &serde_json::Value,
+) -> Result<sql::ast::MutationValueExpression, Error> {
+    match object {
+        serde_json::Value::Object(object) => {
+            let vec = object.into_iter().collect::<Vec<_>>();
+
+            // We expect exactly one operation.
+            match vec.first() {
+                None => Err(unexpected_operation_error(column_name, vec.len())),
+                Some((operation, value)) => {
+                    if vec.len() != 1 {
+                        Err(unexpected_operation_error(column_name, vec.len()))?;
+                    }
+                    // _set operation.
+                    if *operation == "_set" {
+                        Ok(sql::ast::MutationValueExpression::Expression(
+                            translate_json_value(env, state, value, &column_info.r#type)?,
+                        ))
+                    }
+                    // Operation is not supported.
+                    else {
+                        Err(Error::UnexpectedOperation {
+                            column_name: column_name.to_string(),
+                            operation: (*operation).clone(),
+                            available_operations: vec!["_set".to_string()],
+                        })
+                    }
+                }
+            }
+        }
+        // Unexpected structures.
+        serde_json::Value::Array(_) => Err(Error::UnexpectedStructure(format!(
+            "array structure in update column '{column_name}' argument. Expecting an object.",
+        ))),
+        _ => Err(Error::UnexpectedStructure(format!(
+            "value structure in update '{column_name}' argument. Expecting an object.",
+        ))),
+    }
+}
+
+fn unexpected_operation_error(column_name: &str, len: usize) -> Error {
+    Error::UnexpectedStructure(
+        format!("Column mapping in update for column '{column_name}' should contain exactly 1 operation, but got {len}.")
+    )
+}