feat: support createManyAndReturn (#4842)

query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/top_level_mutations/mod.rs

@@ -1,6 +1,7 @@
 mod create;
 mod create_list;
 mod create_many;
+mod create_many_and_return;
 mod default_value;
 mod delete;
 mod delete_many;

query-engine/connector-test-kit-rs/query-tests-setup/src/runner/mod.rs

@@ -22,6 +22,7 @@ use request_handlers::{
 use serde_json::json;
 use std::{
     env,
+    fmt::Display,
     sync::{atomic::AtomicUsize, Arc},
 };
 
@@ -349,11 +350,8 @@ impl Runner {
         Ok(result)
     }
 
-    pub async fn query_json<T>(&self, query: T) -> TestResult<QueryResult>
-    where
-        T: Into<String>,
-    {
-        let query = query.into();
+    pub async fn query_json(&self, query: impl Display) -> TestResult<QueryResult> {
+        let query = query.to_string();
 
         tracing::debug!("Querying: {}", query.clone().green());
 
@@ -364,6 +362,7 @@ impl Runner {
             RunnerExecutor::Builtin(e) => e,
             RunnerExecutor::External(external) => {
                 let response = external.query(query, self.current_tx_id.as_ref()).await?;
+
                 return Ok(response);
             }
         };

query-engine/connectors/sql-query-connector/src/database/operations/write.rs

@@ -270,11 +270,14 @@ pub(crate) async fn create_records_returning(
     let meta = column_metadata::create(&field_names, &idents);
     let mut records = ManyRecords::new(field_names.clone());
     let inserts = generate_insert_statements(model, args, skip_duplicates, Some(&selected_fields.into()), ctx);
+
     for insert in inserts {
         let result_set = conn.query(insert.into()).await?;
+
         for result_row in result_set {
             let sql_row = result_row.to_sql_row(&meta)?;
             let record = Record::from(sql_row);
+
             records.push(record);
         }
     }

query-engine/core/src/interpreter/query_interpreters/read.rs

@@ -137,6 +137,7 @@ fn read_many_by_queries(
             record_not_found()
         } else {
             let nested: Vec<QueryResult> = process_nested(tx, query.nested, Some(&records)).await?;
+
             Ok(RecordSelection {
                 name: query.name,
                 fields: query.selection_order,
@@ -268,7 +269,7 @@ async fn aggregate(
     }))
 }
 
-fn process_nested<'conn>(
+pub(crate) fn process_nested<'conn>(
     tx: &'conn mut dyn ConnectionLike,
     nested: Vec<ReadQuery>,
     parent_result: Option<&'conn ManyRecords>,

query-engine/core/src/interpreter/query_interpreters/write.rs

@@ -72,18 +72,21 @@ async fn create_many(
             .create_records_returning(&q.model, q.args, q.skip_duplicates, selected_fields.fields, trace_id)
             .await?;
 
+        let nested: Vec<QueryResult> = super::read::process_nested(tx, selected_fields.nested, Some(&records)).await?;
+
         let selection = RecordSelection {
             name: q.name,
             fields: selected_fields.order,
             records,
-            nested: vec![],
+            nested,
             model: q.model,
             virtual_fields: vec![],
         };
 
         Ok(QueryResult::RecordSelection(Some(Box::new(selection))))
     } else {
         let affected_records = tx.create_records(&q.model, q.args, q.skip_duplicates, trace_id).await?;
+
         Ok(QueryResult::Count(affected_records))
     }
 }
@@ -108,6 +111,7 @@ async fn create_many_split_by_shape(
 
     if let Some(selected_fields) = q.selected_fields {
         let mut result: Option<ManyRecords> = None;
+
         for args in args_by_shape.into_values() {
             let current_batch = tx
                 .create_records_returning(
@@ -137,24 +141,29 @@ async fn create_many_split_by_shape(
                 .await?
         };
 
+        let nested: Vec<QueryResult> =
+            super::read::process_nested(tx, selected_fields.nested.clone(), Some(&records)).await?;
+
         let selection = RecordSelection {
             name: q.name,
             fields: selected_fields.order,
             records,
-            nested: vec![],
+            nested,
             model: q.model,
             virtual_fields: vec![],
         };
 
         Ok(QueryResult::RecordSelection(Some(Box::new(selection))))
     } else {
         let mut result = 0;
+
         for args in args_by_shape.into_values() {
             let affected_records = tx
                 .create_records(&q.model, args, q.skip_duplicates, trace_id.clone())
                 .await?;
             result += affected_records;
         }
+
         Ok(QueryResult::Count(result))
     }
 }

query-engine/core/src/query_ast/write.rs

@@ -1,6 +1,6 @@
 //! Write query AST
 use super::{FilteredNestedMutation, FilteredQuery};
-use crate::{RecordQuery, ToGraphviz};
+use crate::{ReadQuery, RecordQuery, ToGraphviz};
 use connector::{DatasourceFieldName, NativeUpsert, RecordFilter, WriteArgs};
 use query_structure::{prelude::*, Filter};
 use std::collections::HashMap;
@@ -286,6 +286,7 @@ pub struct CreateManyRecords {
 pub struct CreateManyRecordsFields {
     pub fields: FieldSelection,
     pub order: Vec<String>,
+    pub nested: Vec<ReadQuery>,
 }
 
 #[derive(Debug, Clone)]

query-engine/core/src/query_graph_builder/builder.rs

@@ -113,7 +113,8 @@ impl<'a> QueryGraphBuilder<'a> {
             (QueryTag::Aggregate, Some(m)) => read::aggregate(parsed_field, m).map(Into::into),
             (QueryTag::GroupBy, Some(m)) => read::group_by(parsed_field, m).map(Into::into),
             (QueryTag::CreateOne, Some(m)) => QueryGraph::root(|g| write::create_record(g, query_schema, m, parsed_field)),
-            (QueryTag::CreateMany, Some(m)) => QueryGraph::root(|g| write::create_many_records(g, query_schema,m, parsed_field)),
+            (QueryTag::CreateMany, Some(m)) => QueryGraph::root(|g| write::create_many_records(g, query_schema, m, false, parsed_field)),
+            (QueryTag::CreateManyAndReturn, Some(m)) => QueryGraph::root(|g| write::create_many_records(g, query_schema, m, true, parsed_field)),
             (QueryTag::UpdateOne, Some(m)) => QueryGraph::root(|g| write::update_record(g, query_schema, m, parsed_field)),
             (QueryTag::UpdateMany, Some(m)) => QueryGraph::root(|g| write::update_many_records(g, query_schema, m, parsed_field)),
             (QueryTag::UpsertOne, Some(m)) => QueryGraph::root(|g| write::upsert_record(g, query_schema, m, parsed_field)),

query-engine/core/src/query_graph_builder/mod.rs

@@ -3,8 +3,8 @@
 mod builder;
 mod error;
 mod extractors;
-mod read;
 
+pub(crate) mod read;
 pub(crate) mod write;
 pub(crate) use extractors::*;
 

query-engine/core/src/query_graph_builder/read/one.rs

@@ -44,11 +44,8 @@ fn find_unique_with_options(
 
     let name = field.name;
     let alias = field.alias;
-    let nested_fields = field.nested_fields.unwrap().fields;
-    let selection_order = utils::collect_selection_order(&nested_fields);
-    let selected_fields = utils::collect_selected_fields(&nested_fields, None, &model, query_schema)?;
-    let nested = utils::collect_nested_queries(nested_fields, &model, query_schema)?;
-    let selected_fields = utils::merge_relation_selections(selected_fields, None, &nested);
+    let (selected_fields, selection_order, nested) =
+        utils::extract_selected_fields(field.nested_fields.unwrap().fields, &model, query_schema)?;
 
     let relation_load_strategy = get_relation_load_strategy(requested_rel_load_strategy, None, &nested, query_schema)?;
 

query-engine/core/src/query_graph_builder/read/utils.rs

@@ -306,3 +306,16 @@ fn query_can_be_resolved_with_joins(cursor: Option<&SelectionResult>, nested_que
             _ => false,
         })
 }
+
+pub(crate) fn extract_selected_fields(
+    nested_fields: Vec<FieldPair<'_>>,
+    model: &Model,
+    query_schema: &QuerySchema,
+) -> crate::QueryGraphBuilderResult<(FieldSelection, Vec<String>, Vec<ReadQuery>)> {
+    let selection_order = utils::collect_selection_order(&nested_fields);
+    let selected_fields = utils::collect_selected_fields(&nested_fields, None, model, query_schema)?;
+    let nested = utils::collect_nested_queries(nested_fields, model, query_schema)?;
+    let selected_fields = utils::merge_relation_selections(selected_fields, None, &nested);
+
+    Ok((selected_fields, selection_order, nested))
+}

query-engine/core/src/query_graph_builder/write/create.rs

@@ -68,6 +68,7 @@ pub(crate) fn create_many_records(
     graph: &mut QueryGraph,
     query_schema: &QuerySchema,
     model: Model,
+    with_field_selection: bool,
     mut field: ParsedField<'_>,
 ) -> QueryGraphBuilderResult<()> {
     graph.flag_transactional();
@@ -93,16 +94,30 @@ pub(crate) fn create_many_records(
         })
         .collect::<QueryGraphBuilderResult<Vec<_>>>()?;
 
+    let selected_fields = if with_field_selection {
+        let (selected_fields, selection_order, nested_read) =
+            super::read::utils::extract_selected_fields(field.nested_fields.unwrap().fields, &model, query_schema)?;
+
+        Some(CreateManyRecordsFields {
+            fields: selected_fields,
+            order: selection_order,
+            nested: nested_read,
+        })
+    } else {
+        None
+    };
+
     let query = CreateManyRecords {
         name: field.name,
         model,
         args,
         skip_duplicates,
-        selected_fields: None,
+        selected_fields,
         split_by_shape: !query_schema.has_capability(ConnectorCapability::SupportsDefaultInInsert),
     };
 
     graph.create_node(Query::Write(WriteQuery::CreateManyRecords(query)));
+
     Ok(())
 }
 

query-engine/core/src/query_graph_builder/write/nested/create_nested.rs

@@ -65,6 +65,7 @@ pub fn nested_create(
             Some(CreateManyRecordsFields {
                 fields: selected_fields,
                 order: selection_order,
+                nested: Vec::new(),
             })
         } else {
             None

query-engine/schema/src/build/mutations/create_many.rs

@@ -2,7 +2,7 @@ use super::*;
 use crate::{Identifier, IdentifierType, InputField, InputType, OutputField, OutputType, QueryInfo, QueryTag};
 use constants::*;
 use input_types::{fields::data_input_mapper::*, list_union_type};
-use output_types::objects;
+use output_types::{field, objects};
 use psl::datamodel_connector::ConnectorCapability;
 use query_structure::{Model, RelationFieldRef};
 
@@ -22,6 +22,50 @@ pub(crate) fn create_many(ctx: &'_ QuerySchema, model: Model) -> OutputField<'_>
     )
 }
 
+/// Builds a create many mutation field (e.g. createManyUsers) for given model.
+pub(crate) fn create_many_and_return(ctx: &'_ QuerySchema, model: Model) -> OutputField<'_> {
+    let field_name = format!("createMany{}AndReturn", model.name());
+    let model_id = model.id;
+    let object_type = create_many_and_return_output_type(ctx, model.clone());
+
+    field(
+        field_name,
+        move || create_many_arguments(ctx, model),
+        OutputType::list(InnerOutputType::Object(object_type)),
+        Some(QueryInfo {
+            model: Some(model_id),
+            tag: QueryTag::CreateManyAndReturn,
+        }),
+    )
+}
+
+pub(crate) fn create_many_and_return_output_type(ctx: &'_ QuerySchema, model: Model) -> ObjectType<'_> {
+    let model_id = model.id;
+    let mut obj = ObjectType::new(
+        Identifier::new_model(IdentifierType::CreateManyAndReturnOutput(model.clone())),
+        move || {
+            let mut fields: Vec<_> = model
+                .fields()
+                .scalar()
+                .map(|sf| field::map_output_field(ctx, sf.into()))
+                .collect();
+
+            // If the relation is inlined in the enclosing model, that means the foreign keys can be set at creation
+            // and thus it makes sense to enable querying this relation.
+            for rf in model.fields().relation() {
+                if rf.is_inlined_on_enclosing_model() {
+                    fields.push(field::map_output_field(ctx, rf.into()));
+                }
+            }
+
+            fields
+        },
+    );
+
+    obj.model = Some(model_id);
+    obj
+}
+
 /// Builds "skip_duplicates" and "data" arguments intended for the create many field.
 pub(crate) fn create_many_arguments(ctx: &'_ QuerySchema, model: Model) -> Vec<InputField<'_>> {
     let create_many_type = InputType::object(create_many_object_type(ctx, model, None));

query-engine/schema/src/build/mutations/mod.rs

@@ -1,7 +1,7 @@
 pub(crate) mod create_many;
 pub(crate) mod create_one;
 
-pub(crate) use create_many::create_many;
+pub(crate) use create_many::{create_many, create_many_and_return};
 pub(crate) use create_one::create_one;
 
 use super::*;

query-engine/schema/src/build/output_types/mutation_type.rs

@@ -1,6 +1,6 @@
 use super::*;
 use input_types::fields::arguments;
-use mutations::{create_many, create_one};
+use mutations::{create_many, create_many_and_return, create_one};
 use psl::datamodel_connector::ConnectorCapability;
 use query_structure::{DefaultKind, PrismaValue};
 
@@ -20,8 +20,13 @@ pub(crate) fn mutation_fields(ctx: &QuerySchema) -> Vec<FieldFn> {
             field!(create_one, model);
 
             field!(upsert_item_field, model);
+
             if ctx.has_capability(ConnectorCapability::CreateMany) {
                 field!(create_many, model);
+
+                if ctx.has_capability(ConnectorCapability::InsertReturning) {
+                    field!(create_many_and_return, model);
+                }
             }
         }
 

query-engine/schema/src/build/output_types/query_type.rs

@@ -15,7 +15,7 @@ pub(crate) fn query_fields(ctx: &QuerySchema) -> Vec<FieldFn> {
     for model in ctx.internal_data_model.models() {
         field!(find_first_field, model);
         field!(find_first_or_throw_field, model);
-        field!(all_items_field, model);
+        field!(find_many_field, model);
         field!(plain_aggregation_field, model);
         field!(group_by_aggregation_field, model);
         field!(find_unique_field, model);
@@ -113,7 +113,7 @@ fn find_first_or_throw_field(ctx: &QuerySchema, model: Model) -> OutputField<'_>
 }
 
 /// Builds a "multiple" query arity items field (e.g. "users", "posts", ...) for given model.
-fn all_items_field(ctx: &QuerySchema, model: Model) -> OutputField<'_> {
+fn find_many_field(ctx: &QuerySchema, model: Model) -> OutputField<'_> {
     let field_name = format!("findMany{}", model.name());
     let object_type = objects::model::model_object_type(ctx, model.clone());
     let model_id = model.id;

query-engine/schema/src/identifier_type.rs

@@ -22,6 +22,7 @@ pub enum IdentifierType {
     CompositeUpdateManyInput(CompositeType),
     CompositeUpsertObjectInput(CompositeType),
     CreateManyInput(Model, Option<RelationField>),
+    CreateManyAndReturnOutput(Model),
     CreateOneScalarList(ScalarField),
     Enum(InternalEnum),
     FieldUpdateOperationsInput(bool, String),
@@ -296,6 +297,9 @@ impl std::fmt::Display for IdentifierType {
                 Some(ref rf) => write!(f, "{}CreateMany{}Input", model.name(), capitalize(rf.name())),
                 _ => write!(f, "{}CreateManyInput", model.name()),
             },
+            IdentifierType::CreateManyAndReturnOutput(model) => {
+                write!(f, "CreateMany{}AndReturnOutputType", model.name())
+            }
             IdentifierType::UncheckedUpdateManyInput(model, related_field) => match related_field {
                 Some(rf) => write!(
                     f,

query-engine/schema/src/query_schema.rs

@@ -232,6 +232,7 @@ pub enum QueryTag {
     FindMany,
     CreateOne,
     CreateMany,
+    CreateManyAndReturn,
     UpdateOne,
     UpdateMany,
     DeleteOne,
@@ -257,6 +258,7 @@ impl fmt::Display for QueryTag {
             Self::FindMany => "findMany",
             Self::CreateOne => "createOne",
             Self::CreateMany => "createMany",
+            Self::CreateManyAndReturn => "createManyAndReturn",
             Self::UpdateOne => "updateOne",
             Self::UpdateMany => "updateMany",
             Self::DeleteOne => "deleteOne",
@@ -285,6 +287,7 @@ impl From<&str> for QueryTag {
             "findMany" => Self::FindMany,
             "createOne" => Self::CreateOne,
             "createMany" => Self::CreateMany,
+            "createManyAndReturn" => Self::CreateManyAndReturn,
             "updateOne" => Self::UpdateOne,
             "updateMany" => Self::UpdateMany,
             "deleteOne" => Self::DeleteOne,