Add more object types

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

@@ -21,7 +21,7 @@ pub async fn get_schema(
 ) -> Result<models::SchemaResponse, connector::SchemaError> {
     let configuration::RuntimeConfiguration { metadata, .. } = config;
 
-    let scalar_types: BTreeMap<String, models::ScalarType> =
+    let mut scalar_types: BTreeMap<String, models::ScalarType> =
         configuration::occurring_scalar_types(&metadata.tables, &metadata.native_queries)
             .iter()
             .map(|scalar_type| {
@@ -260,7 +260,9 @@ pub async fn get_schema(
             config.mutations_version,
         )
         .iter()
-        .map(|(name, mutation)| mutation_to_procedure(name, mutation, &mut more_object_types))
+        .map(|(name, mutation)| {
+            mutation_to_procedure(name, mutation, &mut more_object_types, &mut scalar_types)
+        })
         .collect();
 
     procedures.extend(generated_procedures);
@@ -301,17 +303,25 @@ fn mutation_to_procedure(
     name: &String,
     mutation: &generate::Mutation,
     object_types: &mut BTreeMap<String, models::ObjectType>,
+    scalar_types: &mut BTreeMap<String, models::ScalarType>,
 ) -> models::ProcedureInfo {
     match mutation {
-        generate::Mutation::DeleteMutation(delete) => delete_to_procedure(name, delete),
+        generate::Mutation::DeleteMutation(delete) => {
+            delete_to_procedure(name, delete, object_types, scalar_types)
+        }
         generate::Mutation::InsertMutation(insert) => {
-            insert_to_procedure(name, insert, object_types)
+            insert_to_procedure(name, insert, object_types, scalar_types)
         }
     }
 }
 
 /// given a `DeleteMutation`, turn it into a `ProcedureInfo` to be output in the schema
-fn delete_to_procedure(name: &String, delete: &delete::DeleteMutation) -> models::ProcedureInfo {
+fn delete_to_procedure(
+    name: &String,
+    delete: &delete::DeleteMutation,
+    object_types: &mut BTreeMap<String, models::ObjectType>,
+    scalar_types: &mut BTreeMap<String, models::ScalarType>,
+) -> models::ProcedureInfo {
     match delete {
         delete::DeleteMutation::DeleteByKey {
             by_column,
@@ -329,14 +339,16 @@ fn delete_to_procedure(name: &String, delete: &delete::DeleteMutation) -> models
                 },
             );
 
-            models::ProcedureInfo {
-                name: name.to_string(),
-                description: Some(description.to_string()),
+            make_procedure_type(
+                name.to_string(),
+                Some(description.to_string()),
                 arguments,
-                result_type: models::Type::Named {
+                models::Type::Named {
                     name: collection_name.to_string(),
                 },
-            }
+                object_types,
+                scalar_types,
+            )
         }
     }
 }
@@ -379,6 +391,7 @@ fn insert_to_procedure(
     name: &String,
     insert: &insert::InsertMutation,
     object_types: &mut BTreeMap<String, models::ObjectType>,
+    scalar_types: &mut BTreeMap<String, models::ScalarType>,
 ) -> models::ProcedureInfo {
     let mut arguments = BTreeMap::new();
     let object_type = make_object_type(&insert.columns);
@@ -393,12 +406,78 @@ fn insert_to_procedure(
         },
     );
 
+    make_procedure_type(
+        name.to_string(),
+        Some(insert.description.to_string()),
+        arguments,
+        models::Type::Named {
+            name: insert.collection_name.to_string(),
+        },
+        object_types,
+        scalar_types,
+    )
+}
+
+/// Build a `ProcedureInfo` type from the given parameters.
+///
+/// Because procedures return an `affected_rows` count alongside the result type that it's
+/// `returning`, we have to generate a separate object type for its result. As part of that, we may
+/// also have to include the `int4` scalar type (if it isn't included for another reason elsewhere
+/// in the schema). So, this function creates that object type, optionally adds that scalar type,
+/// and then returns a `ProcedureInfo` that points to the correct object type.
+fn make_procedure_type(
+    name: String,
+    description: Option<String>,
+    arguments: BTreeMap<String, models::ArgumentInfo>,
+    result_type: models::Type,
+
+    object_types: &mut BTreeMap<String, models::ObjectType>,
+    scalar_types: &mut BTreeMap<String, models::ScalarType>,
+) -> models::ProcedureInfo {
+    let mut fields = BTreeMap::new();
+    let object_type_name = format!("{name}_response");
+
+    scalar_types
+        .entry("int4".to_string())
+        .or_insert(models::ScalarType {
+            aggregate_functions: BTreeMap::new(),
+            comparison_operators: BTreeMap::new(),
+        });
+
+    fields.insert(
+        "affected_rows".to_string(),
+        models::ObjectField {
+            description: Some("The number of rows affected by the mutation".to_string()),
+            r#type: models::Type::Named {
+                name: "int4".to_string(),
+            },
+        },
+    );
+
+    fields.insert(
+        "returning".to_string(),
+        models::ObjectField {
+            description: Some("Data from rows affected by the mutation".to_string()),
+            r#type: models::Type::Array {
+                element_type: Box::from(result_type),
+            },
+        },
+    );
+
+    object_types.insert(
+        object_type_name.clone(),
+        models::ObjectType {
+            description: Some(format!("Responses from the '{name}' procedure")),
+            fields,
+        },
+    );
+
     models::ProcedureInfo {
-        name: name.to_string(),
-        description: Some(insert.description.to_string()),
+        name,
+        description,
         arguments,
         result_type: models::Type::Named {
-            name: insert.collection_name.to_string(),
+            name: object_type_name,
         },
     }
 }

crates/query-engine/translation/src/translation/mutation/translate.rs

@@ -300,7 +300,7 @@ pub fn parse_procedure_fields(
 ) -> Result<
     (
         Option<IndexMap<String, models::Aggregate>>, // Contains "affected_rows"
-        Option<IndexMap<String, models::Field>>, // Contains "returning"
+        Option<IndexMap<String, models::Field>>,     // Contains "returning"
     ),
     Error,
 > {