diff --git a/prisma-fmt/src/get_dmmf.rs b/prisma-fmt/src/get_dmmf.rs index 26ec932242fc..a798a6d9661c 100644 --- a/prisma-fmt/src/get_dmmf.rs +++ b/prisma-fmt/src/get_dmmf.rs @@ -5533,6 +5533,50 @@ mod tests { "isList": false } }, + { + "name": "updateManyAAndReturn", + "args": [ + { + "name": "data", + "isRequired": true, + "isNullable": false, + "inputTypes": [ + { + "type": "AUpdateManyMutationInput", + "namespace": "prisma", + "location": "inputObjectTypes", + "isList": false + }, + { + "type": "AUncheckedUpdateManyInput", + "namespace": "prisma", + "location": "inputObjectTypes", + "isList": false + } + ] + }, + { + "name": "where", + "isRequired": false, + "isNullable": false, + "inputTypes": [ + { + "type": "AWhereInput", + "namespace": "prisma", + "location": "inputObjectTypes", + "isList": false + } + ] + } + ], + "isNullable": false, + "outputType": { + "type": "UpdateManyAAndReturnOutputType", + "namespace": "model", + "location": "outputObjectTypes", + "isList": true + } + }, { "name": "deleteManyA", "args": [ @@ -5851,6 +5895,50 @@ mod tests { "isList": false } }, + { + "name": "updateManyBAndReturn", + "args": [ + { + "name": "data", + "isRequired": true, + "isNullable": false, + "inputTypes": [ + { + "type": "BUpdateManyMutationInput", + "namespace": "prisma", + "location": "inputObjectTypes", + "isList": false + }, + { + "type": "BUncheckedUpdateManyInput", + "namespace": "prisma", + "location": "inputObjectTypes", + "isList": false + } + ] + }, + { + "name": "where", + "isRequired": false, + "isNullable": false, + "inputTypes": [ + { + "type": "BWhereInput", + "namespace": "prisma", + "location": "inputObjectTypes", + "isList": false + } + ] + } + ], + "isNullable": false, + "outputType": { + "type": "UpdateManyBAndReturnOutputType", + "namespace": "model", + "location": "outputObjectTypes", + "isList": true + } + }, { "name": "deleteManyB", "args": [ @@ -6399,6 +6487,42 @@ mod tests { } ] }, + { + "name": "UpdateManyAAndReturnOutputType", + "fields": [ + { + "name": "id", + "args": [], + "isNullable": false, + "outputType": { + "type": "String", + "location": "scalar", + "isList": false + } + }, + { + "name": "b_id", + "args": [], + "isNullable": false, + "outputType": { + "type": "String", + "location": "scalar", + "isList": false + } + }, + { + "name": "b", + "args": [], + "isNullable": false, + "outputType": { + "type": "B", + "namespace": "model", + "location": "outputObjectTypes", + "isList": false + } + } + ] + }, { "name": "CreateManyBAndReturnOutputType", "fields": [ @@ -6413,6 +6537,21 @@ mod tests { } } ] + }, + { + "name": "UpdateManyBAndReturnOutputType", + "fields": [ + { + "name": "id", + "args": [], + "isNullable": false, + "outputType": { + "type": "String", + "location": "scalar", + "isList": false + } + } + ] } ] }, @@ -6622,6 +6761,7 @@ mod tests { "findUniqueOrThrow": "findUniqueAOrThrow", "groupBy": "groupByA", "updateMany": "updateManyA", + "updateManyAndReturn": "updateManyAAndReturn", "updateOne": "updateOneA", "upsertOne": "upsertOneA" }, @@ -6640,6 +6780,7 @@ mod tests { "findUniqueOrThrow": "findUniqueBOrThrow", "groupBy": "groupByB", "updateMany": "updateManyB", + "updateManyAndReturn": "updateManyBAndReturn", "updateOne": "updateOneB", "upsertOne": "upsertOneB" } diff --git a/query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/top_level_mutations/mod.rs b/query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/top_level_mutations/mod.rs index a34936908cf1..e1342917d890 100644 --- a/query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/top_level_mutations/mod.rs +++ b/query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/top_level_mutations/mod.rs @@ -11,4 +11,5 @@ mod insert_null_in_required_field; mod non_embedded_upsert; mod update; mod update_many; +mod update_many_and_return; mod upsert; diff --git a/query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/top_level_mutations/update_many_and_return.rs b/query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/top_level_mutations/update_many_and_return.rs new file mode 100644 index 000000000000..caced8231b3f --- /dev/null +++ b/query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/top_level_mutations/update_many_and_return.rs @@ -0,0 +1,536 @@ +use query_engine_tests::*; + +#[test_suite(schema(schema), capabilities(UpdateReturning))] +mod update_many_and_return { + use indoc::indoc; + use query_engine_tests::{is_one_of, run_query, run_query_json}; + + fn schema() -> String { + let schema = indoc! { + r#"model TestModel { + #id(id, Int, @id) + optStr String? + optInt Int? + optFloat Float? + }"# + }; + + schema.to_owned() + } + + // "An updateManyAndReturn mutation" should "update the records matching the where clause" + #[connector_test] + async fn update_recs_matching_where(runner: Runner) -> TestResult<()> { + create_row(&runner, r#"{ id: 1, optStr: "str1" }"#).await?; + create_row(&runner, r#"{ id: 2, optStr: "str2" }"#).await?; + + insta::assert_snapshot!( + run_query!(&runner, r#"mutation { + updateManyTestModelAndReturn( + where: { optStr: { equals: "str1" } } + data: { optStr: { set: "str1new" }, optInt: { set: 1 }, optFloat: { multiply: 2 } } + ) { + optStr optInt optFloat + } + }"#), + @r###"{"data":{"updateManyTestModelAndReturn":[{"optStr":"str1new","optInt":1,"optFloat":null}]}}"### + ); + + Ok(()) + } + + // "An updateMany mutation" should "update the records matching the where clause using shorthands" + #[connector_test] + async fn update_recs_matching_where_shorthands(runner: Runner) -> TestResult<()> { + create_row(&runner, r#"{ id: 1, optStr: "str1" }"#).await?; + create_row(&runner, r#"{ id: 2, optStr: "str2" }"#).await?; + + insta::assert_snapshot!( + run_query!(&runner, r#"mutation { + updateManyTestModelAndReturn( + where: { optStr: "str1" } + data: { optStr: "str1new", optInt: null, optFloat: { multiply: 2 } } + ) { + optStr optInt optFloat + } + }"#), + @r###"{"data":{"updateManyTestModelAndReturn":[{"optStr":"str1new","optInt":null,"optFloat":null}]}}"### + ); + + Ok(()) + } + + // "An updateManyAndReturn mutation" should "update all items if the where clause is empty" + #[connector_test] + async fn update_all_items_if_where_empty(runner: Runner) -> TestResult<()> { + create_row(&runner, r#"{ id: 2, optStr: "str2", optInt: 2 }"#).await?; + create_row(&runner, r#"{ id: 3, optStr: "str3", optInt: 3, optFloat: 3.1 }"#).await?; + + insta::assert_snapshot!( + run_query!(&runner, r#"mutation { + updateManyTestModelAndReturn( + where: { } + data: { optStr: { set: "updated" }, optFloat: { divide: 2 }, optInt: { decrement: 1 } } + ){ + optStr optInt optFloat + } + }"#), + @r###"{"data":{"updateManyTestModelAndReturn":[{"optStr":"updated","optInt":1,"optFloat":null},{"optStr":"updated","optInt":2,"optFloat":1.55}]}}"### + ); + + Ok(()) + } + + // "An updateManyAndReturn mutation" should "correctly apply all number operations for Int" + #[connector_test(exclude(CockroachDb))] + async fn apply_number_ops_for_int(runner: Runner) -> TestResult<()> { + create_row(&runner, r#"{ id: 1, optStr: "str1" }"#).await?; + create_row(&runner, r#"{ id: 2, optStr: "str2", optInt: 2 }"#).await?; + create_row(&runner, r#"{ id: 3, optStr: "str3", optInt: 3, optFloat: 3.1 }"#).await?; + + is_one_of!( + query_number_operation(&runner, "optInt", "increment", "10").await?, + [ + r#"[{"optInt":null},{"optInt":12},{"optInt":13}]"#, + r#"[{"optInt":10},{"optInt":12},{"optInt":13}]"# + ] + ); + + // optInts before this op are now: null/10, 12, 13 + is_one_of!( + query_number_operation(&runner, "optInt", "decrement", "10").await?, + [ + r#"[{"optInt":null},{"optInt":2},{"optInt":3}]"#, + r#"[{"optInt":0},{"optInt":2},{"optInt":3}]"# + ] + ); + + // optInts before this op are now: null/0, 2, 3 + is_one_of!( + query_number_operation(&runner, "optInt", "multiply", "2").await?, + [ + r#"[{"optInt":null},{"optInt":4},{"optInt":6}]"#, + r#"[{"optInt":0},{"optInt":4},{"optInt":6}]"# + ] + ); + + // Todo: Mongo divisions are broken + if !matches!(runner.connector_version(), ConnectorVersion::MongoDb(_)) { + // optInts before this op are now: null/0, 4, 6 + is_one_of!( + query_number_operation(&runner, "optInt", "divide", "3").await?, + [ + r#"[{"optInt":null},{"optInt":1},{"optInt":2}]"#, + r#"[{"optInt":0},{"optInt":1},{"optInt":2}]"# + ] + ); + } + + is_one_of!( + query_number_operation(&runner, "optInt", "set", "5").await?, + [ + r#"[{"optInt":5},{"optInt":5},{"optInt":5}]"#, + r#"[{"optInt":5},{"optInt":5},{"optInt":5}]"# + ] + ); + + is_one_of!( + query_number_operation(&runner, "optInt", "set", "null").await?, + [ + r#"[{"optInt":null},{"optInt":null},{"optInt":null}]"#, + r#"[{"optInt":null},{"optInt":null},{"optInt":null}]"# + ] + ); + + Ok(()) + } + + // CockroachDB does not support the "divide" operator as is. + // See https://github.com/cockroachdb/cockroach/issues/41448. + #[connector_test(only(CockroachDb))] + async fn apply_number_ops_for_int_cockroach(runner: Runner) -> TestResult<()> { + create_row(&runner, r#"{ id: 1, optStr: "str1" }"#).await?; + create_row(&runner, r#"{ id: 2, optStr: "str2", optInt: 2 }"#).await?; + create_row(&runner, r#"{ id: 3, optStr: "str3", optInt: 3, optFloat: 3.1 }"#).await?; + + is_one_of!( + query_number_operation(&runner, "optInt", "increment", "10").await?, + [ + r#"[{"optInt":null},{"optInt":12},{"optInt":13}]"#, + r#"[{"optInt":10},{"optInt":12},{"optInt":13}]"# + ] + ); + + // optInts before this op are now: null/10, 12, 13 + is_one_of!( + query_number_operation(&runner, "optInt", "decrement", "10").await?, + [ + r#"[{"optInt":null},{"optInt":2},{"optInt":3}]"#, + r#"[{"optInt":0},{"optInt":2},{"optInt":3}]"# + ] + ); + + // optInts before this op are now: null/0, 2, 3 + is_one_of!( + query_number_operation(&runner, "optInt", "multiply", "2").await?, + [ + r#"[{"optInt":null},{"optInt":4},{"optInt":6}]"#, + r#"[{"optInt":0},{"optInt":4},{"optInt":6}]"# + ] + ); + + is_one_of!( + query_number_operation(&runner, "optInt", "set", "5").await?, + [ + r#"[{"optInt":5},{"optInt":5},{"optInt":5}]"#, + r#"[{"optInt":5},{"optInt":5},{"optInt":5}]"# + ] + ); + + is_one_of!( + query_number_operation(&runner, "optInt", "set", "null").await?, + [ + r#"[{"optInt":null},{"optInt":null},{"optInt":null}]"#, + r#"[{"optInt":null},{"optInt":null},{"optInt":null}]"# + ] + ); + + Ok(()) + } + + // "An updateManyAndReturn mutation" should "correctly apply all number operations for Float" + #[connector_test] + async fn apply_number_ops_for_float(runner: Runner) -> TestResult<()> { + create_row(&runner, r#"{ id: 1, optStr: "str1" }"#).await?; + create_row(&runner, r#"{ id: 2, optStr: "str2", optFloat: 2 }"#).await?; + create_row(&runner, r#"{ id: 3, optStr: "str3", optFloat: 3.1 }"#).await?; + + insta::assert_snapshot!( + query_number_operation(&runner, "optFloat", "increment", "1.1").await?, + @r###"[{"optFloat":null},{"optFloat":3.1},{"optFloat":4.2}]"### + ); + + insta::assert_snapshot!( + query_number_operation(&runner, "optFloat", "decrement", "1.1").await?, + @r###"[{"optFloat":null},{"optFloat":2.0},{"optFloat":3.1}]"### + ); + + insta::assert_snapshot!( + query_number_operation(&runner, "optFloat", "multiply", "5.5").await?, + @r###"[{"optFloat":null},{"optFloat":11.0},{"optFloat":17.05}]"### + ); + + insta::assert_snapshot!( + query_number_operation(&runner, "optFloat", "divide", "2").await?, + @r###"[{"optFloat":null},{"optFloat":5.5},{"optFloat":8.525}]"### + ); + + insta::assert_snapshot!( + query_number_operation(&runner, "optFloat", "set", "5").await?, + @r###"[{"optFloat":5.0},{"optFloat":5.0},{"optFloat":5.0}]"### + ); + + insta::assert_snapshot!( + query_number_operation(&runner, "optFloat", "set", "null").await?, + @r###"[{"optFloat":null},{"optFloat":null},{"optFloat":null}]"### + ); + + Ok(()) + } + + fn schema_11_child() -> String { + let schema = indoc! { + r#"model Test { + #id(id, Int, @id) + + child Child? + } + + model Child { + #id(id, Int, @id) + + testId Int? @unique + test Test? @relation(fields: [testId], references: [id]) + + }"# + }; + + schema.to_owned() + } + + #[connector_test(schema(schema_1m_child))] + async fn update_many_11_inline_rel_read_works(runner: Runner) -> TestResult<()> { + insta::assert_snapshot!( + run_query!(&runner,r#"mutation { createManyTest(data: [{ id: 1 }, { id: 2 }]) { count } }"#), + @r###"{"data":{"createManyTest":{"count":2}}}"### + ); + insta::assert_snapshot!( + run_query!(&runner, r#"mutation { createManyChild(data: [{ id: 1, testId: 1 }, { id: 2, testId: 2 }]) { count } }"#), + @r###"{"data":{"createManyChild":{"count":2}}}"### + ); + + insta::assert_snapshot!( + run_query!(&runner, r#"mutation { + updateManyChildAndReturn( + where: {} + data: { testId: 1 } + ) { id test { id } } + }"#), + @r###"{"data":{"updateManyChildAndReturn":[{"id":1,"test":{"id":1}},{"id":2,"test":{"id":1}}]}}"### + ); + + Ok(()) + } + + #[connector_test(schema(schema_11_child))] + async fn update_many_11_non_inline_rel_read_fails(runner: Runner) -> TestResult<()> { + runner + .query_json(serde_json::json!({ + "modelName": "Test", + "action": "updateManyAndReturn", + "query": { + "arguments": { "data": { "id": 1 } }, + "selection": { + "id": true, + "child": true + } + } + })) + .await? + .assert_failure(2009, Some("Field 'child' not found in enclosing type".to_string())); + + Ok(()) + } + + fn schema_1m_child() -> String { + let schema = indoc! { + r#"model Test { + #id(id, Int, @id) + str1 String? + str2 String? + str3 String? @default("SOME_DEFAULT") + + children Child[] + } + + model Child { + #id(id, Int, @id) + str1 String? + str2 String? + str3 String? @default("SOME_DEFAULT") + + testId Int? + test Test? @relation(fields: [testId], references: [id]) + + }"# + }; + + schema.to_owned() + } + + #[connector_test(schema(schema_1m_child))] + async fn update_many_1m_inline_rel_read_works(runner: Runner) -> TestResult<()> { + insta::assert_snapshot!( + run_query!(&runner, r#"mutation { createManyTest(data: [{ id: 1, str1: "1" }, { id: 2, str1: "2" }]) { count } }"#), + @r###"{"data":{"createManyTest":{"count":2}}}"### + ); + insta::assert_snapshot!( + run_query!(&runner, r#"mutation { createManyChild(data: [{ id: 1, str1: "1", testId: 1 }, { id: 2, str1: "2", testId: 2 }]) { count } }"#), + @r###"{"data":{"createManyChild":{"count":2}}}"### + ); + + insta::assert_snapshot!( + run_query!(&runner, r#"mutation { + updateManyChildAndReturn( + where: {} + data: { str1: "updated" } + ) { id str1 } + }"#), + @r###"{"data":{"updateManyChildAndReturn":[{"id":1,"str1":"updated"},{"id":2,"str1":"updated"}]}}"### + ); + + Ok(()) + } + + #[connector_test(schema(schema_1m_child))] + async fn update_many_1m_non_inline_rel_read_fails(runner: Runner) -> TestResult<()> { + runner + .query_json(serde_json::json!({ + "modelName": "Test", + "action": "updateManyAndReturn", + "query": { + "arguments": { "data": { "id": 1 } }, + "selection": { + "id": true, + "children": true + } + } + })) + .await? + .assert_failure(2009, Some("Field 'children' not found in enclosing type".to_string())); + + Ok(()) + } + + fn schema_m2m_child() -> String { + let schema = indoc! { + r#"model Test { + #id(id, Int, @id) + str1 String? + + #m2m(children, Child[], id, Int) + } + + model Child { + #id(id, Int, @id) + + #m2m(tests, Test[], id, Int) + + }"# + }; + + schema.to_owned() + } + + #[connector_test(schema(schema_m2m_child))] + async fn update_many_m2m_rel_read_fails(runner: Runner) -> TestResult<()> { + runner + .query_json(serde_json::json!({ + "modelName": "Test", + "action": "updateManyAndReturn", + "query": { + "arguments": { "data": { "id": 1 } }, + "selection": { + "id": true, + "children": true + } + } + })) + .await? + .assert_failure(2009, Some("Field 'children' not found in enclosing type".to_string())); + + runner + .query_json(serde_json::json!({ + "modelName": "Child", + "action": "updateManyAndReturn", + "query": { + "arguments": { "data": { "id": 1 } }, + "selection": { + "id": true, + "tests": true + } + } + })) + .await? + .assert_failure(2009, Some("Field 'tests' not found in enclosing type".to_string())); + + Ok(()) + } + + fn schema_self_rel_child() -> String { + let schema = indoc! { + r#"model Child { + #id(id, Int, @id) + + teacherId Int? + teacher Child? @relation("TeacherStudents", fields: [teacherId], references: [id]) + students Child[] @relation("TeacherStudents") + }"# + }; + + schema.to_owned() + } + + #[connector_test(schema(schema_self_rel_child))] + async fn update_many_self_rel_read_fails(runner: Runner) -> TestResult<()> { + runner + .query_json(serde_json::json!({ + "modelName": "Child", + "action": "updateManyAndReturn", + "query": { + "arguments": { "data": { "id": 1 } }, + "selection": { + "id": true, + "students": true + } + } + })) + .await? + .assert_failure(2009, Some("Field 'students' not found in enclosing type".to_string())); + + Ok(()) + } + + async fn query_number_operation(runner: &Runner, field: &str, op: &str, value: &str) -> TestResult { + let res = run_query_json!( + runner, + format!( + r#"mutation {{ + updateManyTestModelAndReturn( + where: {{}} + data: {{ {field}: {{ {op}: {value} }} }} + ){{ + {field} + }} + }}"# + ) + ); + + Ok(res["data"]["updateManyTestModelAndReturn"].to_string()) + } + + async fn create_row(runner: &Runner, data: &str) -> TestResult<()> { + runner + .query(format!("mutation {{ createOneTestModel(data: {data}) {{ id }} }}")) + .await? + .assert_success(); + + Ok(()) + } +} + +#[test_suite(schema(json_opt), capabilities(AdvancedJsonNullability, UpdateReturning))] +mod json_update_many_and_return { + use query_engine_tests::{assert_error, run_query}; + + #[connector_test] + async fn update_json_adv(runner: Runner) -> TestResult<()> { + insta::assert_snapshot!( + run_query!(&runner, r#"mutation { createOneTestModel(data: { id: 1 }) { json }}"#), + @r###"{"data":{"createOneTestModel":{"json":null}}}"### + ); + + insta::assert_snapshot!( + run_query!(&runner, r#"mutation { updateManyTestModelAndReturn(where: { id: 1 }, data: { json: "{}" }) { json }}"#), + @r###"{"data":{"updateManyTestModelAndReturn":[{"json":"{}"}]}}"### + ); + + insta::assert_snapshot!( + run_query!(&runner, r#"mutation { updateManyTestModelAndReturn(where: { id: 1 }, data: { json: JsonNull }) { json }}"#), + @r###"{"data":{"updateManyTestModelAndReturn":[{"json":"null"}]}}"### + ); + + insta::assert_snapshot!( + run_query!(&runner, r#"mutation { updateManyTestModelAndReturn(where: { id: 1 }, data: { json: DbNull }) { json }}"#), + @r###"{"data":{"updateManyTestModelAndReturn":[{"json":null}]}}"### + ); + + Ok(()) + } + + #[connector_test] + async fn update_json_errors(runner: Runner) -> TestResult<()> { + assert_error!( + &runner, + r#"mutation { + updateManyTestModelAndReturn(where: { id: 1 }, data: { json: AnyNull }) { + id + } + }"#, + 2009, + "`AnyNull` is not a valid `NullableJsonNullValueInput`" + ); + + Ok(()) + } +} diff --git a/query-engine/connectors/mongodb-query-connector/src/interface/connection.rs b/query-engine/connectors/mongodb-query-connector/src/interface/connection.rs index aeedfc236099..b5aaea79d9b2 100644 --- a/query-engine/connectors/mongodb-query-connector/src/interface/connection.rs +++ b/query-engine/connectors/mongodb-query-connector/src/interface/connection.rs @@ -114,6 +114,17 @@ impl WriteOperations for MongoDbConnection { .await } + async fn update_records_returning( + &mut self, + _model: &Model, + _record_filter: connector_interface::RecordFilter, + _args: WriteArgs, + _selected_fields: FieldSelection, + _traceparent: Option, + ) -> connector_interface::Result { + unimplemented!() + } + async fn update_record( &mut self, model: &Model, diff --git a/query-engine/connectors/mongodb-query-connector/src/interface/transaction.rs b/query-engine/connectors/mongodb-query-connector/src/interface/transaction.rs index ae6546edb7f1..0ba1350d92f6 100644 --- a/query-engine/connectors/mongodb-query-connector/src/interface/transaction.rs +++ b/query-engine/connectors/mongodb-query-connector/src/interface/transaction.rs @@ -144,6 +144,17 @@ impl WriteOperations for MongoDbTransaction<'_> { .await } + async fn update_records_returning( + &mut self, + _model: &Model, + _record_filter: connector_interface::RecordFilter, + _args: connector_interface::WriteArgs, + _selected_fields: FieldSelection, + _traceparent: Option, + ) -> connector_interface::Result { + unimplemented!() + } + async fn update_record( &mut self, model: &Model, diff --git a/query-engine/connectors/query-connector/src/interface.rs b/query-engine/connectors/query-connector/src/interface.rs index 6c7c003903df..c196c76bc309 100644 --- a/query-engine/connectors/query-connector/src/interface.rs +++ b/query-engine/connectors/query-connector/src/interface.rs @@ -289,6 +289,19 @@ pub trait WriteOperations { traceparent: Option, ) -> crate::Result; + /// Updates many records at once into the database and returns their + /// selected fields. + /// This method should not be used if the connector does not support + /// returning updated rows. + async fn update_records_returning( + &mut self, + model: &Model, + record_filter: RecordFilter, + args: WriteArgs, + selected_fields: FieldSelection, + traceparent: Option, + ) -> crate::Result; + /// Update record in the `Model` with the given `WriteArgs` filtered by the /// `Filter`. async fn update_record( diff --git a/query-engine/connectors/sql-query-connector/src/database/connection.rs b/query-engine/connectors/sql-query-connector/src/database/connection.rs index 924af5ec12e4..af0a15192c1a 100644 --- a/query-engine/connectors/sql-query-connector/src/database/connection.rs +++ b/query-engine/connectors/sql-query-connector/src/database/connection.rs @@ -236,6 +236,22 @@ where .await } + async fn update_records_returning( + &mut self, + model: &Model, + record_filter: RecordFilter, + args: WriteArgs, + selected_fields: FieldSelection, + traceparent: Option, + ) -> connector::Result { + let ctx = Context::new(&self.connection_info, traceparent); + catch( + &self.connection_info, + write::update_records_returning(&self.inner, model, record_filter, args, selected_fields, &ctx), + ) + .await + } + async fn update_record( &mut self, model: &Model, diff --git a/query-engine/connectors/sql-query-connector/src/database/operations/update.rs b/query-engine/connectors/sql-query-connector/src/database/operations/update.rs index 0dc8081f97d2..4773853ea25c 100644 --- a/query-engine/connectors/sql-query-connector/src/database/operations/update.rs +++ b/query-engine/connectors/sql-query-connector/src/database/operations/update.rs @@ -2,12 +2,14 @@ use super::read::get_single_record; use crate::column_metadata::{self, ColumnMetadata}; use crate::filter::FilterBuilder; +use crate::model_extensions::AsColumns; use crate::query_builder::write::{build_update_and_set_query, chunk_update_with_ids}; use crate::row::ToSqlRow; use crate::{Context, QueryExt, Queryable}; use connector_interface::*; use itertools::Itertools; +use quaint::ast::Query; use query_structure::*; /// Performs an update with an explicit selection set. @@ -77,7 +79,11 @@ pub(crate) async fn update_one_without_selection( let id_args = pick_args(&model.primary_identifier().into(), &args); // Perform the update and return the ids on which we've applied the update. // Note: We are _not_ getting back the ids from the update. Either we got some ids passed from the parent operation or we perform a read _before_ doing the update. - let (_, ids) = update_many_from_ids_and_filter(conn, model, record_filter, args, ctx).await?; + let (updates, ids) = update_many_from_ids_and_filter(conn, model, record_filter, args, None, ctx).await?; + for update in updates { + conn.execute(update).await?; + } + // Since we could not get the ids back from the update, we need to apply in-memory transformation to the ids in case they were part of the update. // This is critical to ensure the following operations can operate on the updated ids. let merged_ids = merge_write_args(ids, id_args); @@ -92,53 +98,50 @@ pub(crate) async fn update_one_without_selection( // Generates a query like this: // UPDATE "public"."User" SET "name" = $1 WHERE "public"."User"."age" > $1 -pub(crate) async fn update_many_from_filter( - conn: &dyn Queryable, +pub(super) async fn update_many_from_filter( model: &Model, record_filter: RecordFilter, args: WriteArgs, + selected_fields: Option<&ModelProjection>, ctx: &Context<'_>, -) -> crate::Result { +) -> crate::Result> { let update = build_update_and_set_query(model, args, None, ctx); let filter_condition = FilterBuilder::without_top_level_joins().visit_filter(record_filter.filter, ctx); let update = update.so_that(filter_condition); - let count = conn.execute(update.into()).await?; - - Ok(count as usize) + if let Some(selected_fields) = selected_fields { + Ok(update + .returning(selected_fields.as_columns(ctx).map(|c| c.set_is_selected(true))) + .into()) + } else { + Ok(update.into()) + } } // Generates a query like this: // UPDATE "public"."User" SET "name" = $1 WHERE "public"."User"."id" IN ($2,$3,$4,$5,$6,$7,$8,$9,$10,$11) AND "public"."User"."age" > $1 -pub(crate) async fn update_many_from_ids_and_filter( +pub(super) async fn update_many_from_ids_and_filter( conn: &dyn Queryable, model: &Model, record_filter: RecordFilter, args: WriteArgs, + selected_fields: Option<&ModelProjection>, ctx: &Context<'_>, -) -> crate::Result<(usize, Vec)> { +) -> crate::Result<(Vec>, Vec)> { let filter_condition = FilterBuilder::without_top_level_joins().visit_filter(record_filter.filter.clone(), ctx); let ids: Vec = conn.filter_selectors(model, record_filter, ctx).await?; if ids.is_empty() { - return Ok((0, Vec::new())); + return Ok((vec![], Vec::new())); } let updates = { - let update = build_update_and_set_query(model, args, None, ctx); + let update = build_update_and_set_query(model, args, selected_fields, ctx); let ids: Vec<&SelectionResult> = ids.iter().collect(); chunk_update_with_ids(update, model, &ids, filter_condition, ctx)? }; - let mut count = 0; - - for update in updates { - let update_count = conn.execute(update).await?; - - count += update_count; - } - - Ok((count as usize, ids)) + Ok((updates, ids)) } fn process_result_row( diff --git a/query-engine/connectors/sql-query-connector/src/database/operations/write.rs b/query-engine/connectors/sql-query-connector/src/database/operations/write.rs index 137bff50ca58..b71fc7826713 100644 --- a/query-engine/connectors/sql-query-connector/src/database/operations/write.rs +++ b/query-engine/connectors/sql-query-connector/src/database/operations/write.rs @@ -8,7 +8,7 @@ use crate::{ }; use connector_interface::*; use itertools::Itertools; -use quaint::ast::Insert; +use quaint::ast::{Insert, Query}; use quaint::{ error::ErrorKind, prelude::{native_uuid, uuid_to_bin, uuid_to_bin_swapped, Aliasable, Select, SqlFamily}, @@ -370,6 +370,25 @@ pub(crate) async fn update_record( } } +async fn generate_updates( + conn: &dyn Queryable, + model: &Model, + record_filter: RecordFilter, + args: WriteArgs, + selected_fields: Option<&ModelProjection>, + ctx: &Context<'_>, +) -> crate::Result>> { + if record_filter.has_selectors() { + let (updates, _) = + update_many_from_ids_and_filter(conn, model, record_filter, args, selected_fields, ctx).await?; + Ok(updates) + } else { + Ok(vec![ + update_many_from_filter(model, record_filter, args, selected_fields, ctx).await?, + ]) + } +} + /// Update multiple records in a database defined in `conn` and the records /// defined in `args`, and returning the number of updates /// This works via two ways, when there are ids in record_filter.selectors, it uses that to update @@ -385,15 +404,40 @@ pub(crate) async fn update_records( return Ok(0); } - if record_filter.has_selectors() { - let (count, _) = update_many_from_ids_and_filter(conn, model, record_filter, args, ctx).await?; + let mut count = 0; + for update in generate_updates(conn, model, record_filter, args, None, ctx).await? { + count += conn.execute(update).await?; + } + Ok(count as usize) +} - Ok(count) - } else { - let count = update_many_from_filter(conn, model, record_filter, args, ctx).await?; +/// Update records according to `WriteArgs`. Returns values of fields specified in +/// `selected_fields` for all updated rows. +pub(crate) async fn update_records_returning( + conn: &dyn Queryable, + model: &Model, + record_filter: RecordFilter, + args: WriteArgs, + selected_fields: FieldSelection, + ctx: &Context<'_>, +) -> crate::Result { + let field_names: Vec = selected_fields.db_names().collect(); + let idents = selected_fields.type_identifiers_with_arities(); + let meta = column_metadata::create(&field_names, &idents); + let mut records = ManyRecords::new(field_names.clone()); - Ok(count) + for update in generate_updates(conn, model, record_filter, args, Some(&selected_fields.into()), ctx).await? { + let result_set = conn.query(update).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); + } } + + Ok(records) } /// Delete multiple records in `conn`, defined in the `Filter`. Result is the number of items deleted. diff --git a/query-engine/connectors/sql-query-connector/src/database/transaction.rs b/query-engine/connectors/sql-query-connector/src/database/transaction.rs index eea377ad5a57..9a3429ab6e2c 100644 --- a/query-engine/connectors/sql-query-connector/src/database/transaction.rs +++ b/query-engine/connectors/sql-query-connector/src/database/transaction.rs @@ -230,6 +230,29 @@ impl WriteOperations for SqlConnectorTransaction<'_> { .await } + async fn update_records_returning( + &mut self, + model: &Model, + record_filter: RecordFilter, + args: WriteArgs, + selected_fields: FieldSelection, + traceparent: Option, + ) -> connector::Result { + let ctx = Context::new(&self.connection_info, traceparent); + catch( + &self.connection_info, + write::update_records_returning( + self.inner.as_queryable(), + model, + record_filter, + args, + selected_fields, + &ctx, + ), + ) + .await + } + async fn update_record( &mut self, model: &Model, diff --git a/query-engine/core/src/interpreter/query_interpreters/write.rs b/query-engine/core/src/interpreter/query_interpreters/write.rs index 50096ed93392..065f5e6eebe3 100644 --- a/query-engine/core/src/interpreter/query_interpreters/write.rs +++ b/query-engine/core/src/interpreter/query_interpreters/write.rs @@ -305,11 +305,31 @@ async fn update_many( q: UpdateManyRecords, traceparent: Option, ) -> InterpretationResult { - let res = tx - .update_records(&q.model, q.record_filter, q.args, traceparent) - .await?; + if let Some(selected_fields) = q.selected_fields { + let records = tx + .update_records_returning(&q.model, q.record_filter, q.args, selected_fields.fields, traceparent) + .await?; - Ok(QueryResult::Count(res)) + let nested: Vec = + super::read::process_nested(tx, selected_fields.nested, Some(&records), traceparent).await?; + + let selection = RecordSelection { + name: q.name, + fields: selected_fields.order, + records, + nested, + model: q.model, + virtual_fields: vec![], + }; + + Ok(QueryResult::RecordSelection(Some(Box::new(selection)))) + } else { + let affected_records = tx + .update_records(&q.model, q.record_filter, q.args, traceparent) + .await?; + + Ok(QueryResult::Count(affected_records)) + } } async fn delete_many( diff --git a/query-engine/core/src/query_ast/write.rs b/query-engine/core/src/query_ast/write.rs index 940b36cc57cf..9d30e3ff7b5c 100644 --- a/query-engine/core/src/query_ast/write.rs +++ b/query-engine/core/src/query_ast/write.rs @@ -361,9 +361,20 @@ pub struct UpdateRecordWithoutSelection { #[derive(Debug, Clone)] pub struct UpdateManyRecords { + pub name: String, pub model: Model, pub record_filter: RecordFilter, pub args: WriteArgs, + /// Fields of updated records that client has requested to return. + /// `None` if the connector does not support returning the updated rows. + pub selected_fields: Option, +} + +#[derive(Debug, Clone)] +pub struct UpdateManyRecordsFields { + pub fields: FieldSelection, + pub order: Vec, + pub nested: Vec, } #[derive(Debug, Clone)] diff --git a/query-engine/core/src/query_graph_builder/builder.rs b/query-engine/core/src/query_graph_builder/builder.rs index 8ba2929caccd..1cce2f86edfd 100644 --- a/query-engine/core/src/query_graph_builder/builder.rs +++ b/query-engine/core/src/query_graph_builder/builder.rs @@ -116,7 +116,8 @@ impl<'a> QueryGraphBuilder<'a> { (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::UpdateMany, Some(m)) => QueryGraph::root(|g| write::update_many_records(g, query_schema, m, false, parsed_field)), + (QueryTag::UpdateManyAndReturn, Some(m)) => QueryGraph::root(|g| write::update_many_records(g, query_schema, m, true, parsed_field)), (QueryTag::UpsertOne, Some(m)) => QueryGraph::root(|g| write::upsert_record(g, query_schema, m, parsed_field)), (QueryTag::DeleteOne, Some(m)) => QueryGraph::root(|g| write::delete_record(g, query_schema, m, parsed_field)), (QueryTag::DeleteMany, Some(m)) => QueryGraph::root(|g| write::delete_many_records(g, query_schema, m, parsed_field)), diff --git a/query-engine/core/src/query_graph_builder/write/nested/update_nested.rs b/query-engine/core/src/query_graph_builder/write/nested/update_nested.rs index 735bd0a88c31..4628912a4a46 100644 --- a/query-engine/core/src/query_graph_builder/write/nested/update_nested.rs +++ b/query-engine/core/src/query_graph_builder/write/nested/update_nested.rs @@ -142,8 +142,15 @@ pub fn nested_update_many( let find_child_records_node = utils::insert_find_children_by_parent_node(graph, parent, parent_relation_field, filter)?; - let update_many_node = - update::update_many_record_node(graph, query_schema, Filter::empty(), child_model.clone(), data_map)?; + let update_many_node = update::update_many_record_node( + graph, + query_schema, + Filter::empty(), + child_model.clone(), + None, + None, + data_map, + )?; graph.create_edge( &find_child_records_node, diff --git a/query-engine/core/src/query_graph_builder/write/update.rs b/query-engine/core/src/query_graph_builder/write/update.rs index f112fbaa5d54..2f9024d69320 100644 --- a/query-engine/core/src/query_graph_builder/write/update.rs +++ b/query-engine/core/src/query_graph_builder/write/update.rs @@ -1,5 +1,6 @@ use super::*; use crate::query_graph_builder::write::write_args_parser::*; +use crate::ParsedObject; use crate::{ query_ast::*, query_graph::{Node, NodeRef, QueryGraph, QueryGraphDependency}, @@ -124,6 +125,7 @@ pub fn update_many_records( graph: &mut QueryGraph, query_schema: &QuerySchema, model: Model, + with_field_selection: bool, mut field: ParsedField<'_>, ) -> QueryGraphBuilderResult<()> { graph.flag_transactional(); @@ -139,14 +141,30 @@ pub fn update_many_records( let data_map: ParsedInputMap<'_> = data_argument.value.try_into()?; if query_schema.relation_mode().uses_foreign_keys() { - update_many_record_node(graph, query_schema, filter, model, data_map)?; + update_many_record_node( + graph, + query_schema, + filter, + model, + Some(field.name), + field.nested_fields.filter(|_| with_field_selection), + data_map, + )?; } else { let pre_read_node = graph.create_node(utils::read_ids_infallible( model.clone(), model.primary_identifier(), filter, )); - let update_many_node = update_many_record_node(graph, query_schema, Filter::empty(), model.clone(), data_map)?; + let update_many_node = update_many_record_node( + graph, + query_schema, + Filter::empty(), + model.clone(), + Some(field.name), + field.nested_fields.filter(|_| with_field_selection), + data_map, + )?; utils::insert_emulated_on_update(graph, query_schema, &model, &pre_read_node, &update_many_node)?; @@ -249,6 +267,8 @@ pub fn update_many_record_node( query_schema: &QuerySchema, filter: T, model: Model, + name: Option, + nested_field_selection: Option>, data_map: ParsedInputMap<'_>, ) -> QueryGraphBuilderResult where @@ -263,10 +283,25 @@ where args.update_datetimes(&model); + let selected_fields = if let Some(nested_fields) = nested_field_selection { + let (selected_fields, selection_order, nested_read) = + super::read::utils::extract_selected_fields(nested_fields.fields, &model, query_schema)?; + + Some(UpdateManyRecordsFields { + fields: selected_fields, + order: selection_order, + nested: nested_read, + }) + } else { + None + }; + let update_many = UpdateManyRecords { + name: name.unwrap_or_default(), model, record_filter, args, + selected_fields, }; let update_many_node = graph.create_node(Query::Write(WriteQuery::UpdateManyRecords(update_many))); diff --git a/query-engine/core/src/query_graph_builder/write/utils.rs b/query-engine/core/src/query_graph_builder/write/utils.rs index a931fa1edc30..d9f89feade62 100644 --- a/query-engine/core/src/query_graph_builder/write/utils.rs +++ b/query-engine/core/src/query_graph_builder/write/utils.rs @@ -224,9 +224,11 @@ where let record_filter = filter.into(); let ur = UpdateManyRecords { + name: String::new(), model, record_filter, args, + selected_fields: None, }; graph.create_node(Query::Write(WriteQuery::UpdateManyRecords(ur))) @@ -610,9 +612,11 @@ pub fn emulate_on_delete_set_null( insert_find_children_by_parent_node(graph, node_providing_ids, &parent_relation_field, Filter::empty())?; let set_null_query = WriteQuery::UpdateManyRecords(UpdateManyRecords { + name: String::new(), model: dependent_model.clone(), record_filter: RecordFilter::empty(), args: WriteArgs::new(child_update_args, crate::executor::get_request_now()), + selected_fields: None, }); let set_null_dependents_node = graph.create_node(Query::Write(set_null_query)); @@ -760,9 +764,11 @@ pub fn emulate_on_update_set_null( insert_find_children_by_parent_node(graph, parent_node, &parent_relation_field, Filter::empty())?; let set_null_query = WriteQuery::UpdateManyRecords(UpdateManyRecords { + name: String::new(), model: dependent_model.clone(), record_filter: RecordFilter::empty(), args: WriteArgs::new(child_update_args, crate::executor::get_request_now()), + selected_fields: None, }); let set_null_dependents_node = graph.create_node(Query::Write(set_null_query)); @@ -1080,12 +1086,14 @@ pub fn emulate_on_update_cascade( insert_find_children_by_parent_node(graph, parent_node, &parent_relation_field, Filter::empty())?; let update_query = WriteQuery::UpdateManyRecords(UpdateManyRecords { + name: String::new(), model: dependent_model.clone(), record_filter: RecordFilter::empty(), args: WriteArgs::new( child_update_args.into_iter().collect(), crate::executor::get_request_now(), ), + selected_fields: None, }); let update_dependents_node = graph.create_node(Query::Write(update_query)); diff --git a/query-engine/dmmf/src/tests/test-schemas/snapshots/odoo.snapshot.json.gz b/query-engine/dmmf/src/tests/test-schemas/snapshots/odoo.snapshot.json.gz index 742636528d45..2b24b4c53b7a 100644 Binary files a/query-engine/dmmf/src/tests/test-schemas/snapshots/odoo.snapshot.json.gz and b/query-engine/dmmf/src/tests/test-schemas/snapshots/odoo.snapshot.json.gz differ diff --git a/query-engine/schema/src/build/output_types/mutation_type.rs b/query-engine/schema/src/build/output_types/mutation_type.rs index 65a48efcb0f8..2864ea9a85e0 100644 --- a/query-engine/schema/src/build/output_types/mutation_type.rs +++ b/query-engine/schema/src/build/output_types/mutation_type.rs @@ -34,6 +34,9 @@ pub(crate) fn mutation_fields(ctx: &QuerySchema) -> Vec { field!(update_item_field, model); field!(update_many_field, model); + if ctx.has_capability(ConnectorCapability::UpdateReturning) { + field!(update_many_and_return_field, model); + } field!(delete_many_field, model); } @@ -170,6 +173,23 @@ fn update_many_field(ctx: &QuerySchema, model: Model) -> OutputField<'_> { ) } +/// Builds an update many mutation field (e.g. updateManyUsersAndReturn) for given model. +fn update_many_and_return_field(ctx: &QuerySchema, model: Model) -> OutputField<'_> { + let field_name = format!("updateMany{}AndReturn", model.name()); + let cloned_model = model.clone(); + let object_type = update_many_and_return_output_type(ctx, model.clone()); + + field( + field_name, + move || arguments::update_many_arguments(ctx, cloned_model), + OutputType::list(InnerOutputType::Object(object_type)), + Some(QueryInfo { + model: Some(model.id), + tag: QueryTag::UpdateManyAndReturn, + }), + ) +} + /// Builds an upsert mutation field (e.g. upsertUser) for given model. fn upsert_item_field(ctx: &QuerySchema, model: Model) -> OutputField<'_> { let cloned_model = model.clone(); @@ -184,3 +204,30 @@ fn upsert_item_field(ctx: &QuerySchema, model: Model) -> OutputField<'_> { }), ) } + +fn update_many_and_return_output_type(ctx: &'_ QuerySchema, model: Model) -> ObjectType<'_> { + let model_id = model.id; + let mut obj = ObjectType::new( + Identifier::new_model(IdentifierType::UpdateManyAndReturnOutput(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 +} diff --git a/query-engine/schema/src/identifier_type.rs b/query-engine/schema/src/identifier_type.rs index 0aa77140139c..b67a43442d46 100644 --- a/query-engine/schema/src/identifier_type.rs +++ b/query-engine/schema/src/identifier_type.rs @@ -51,6 +51,7 @@ pub enum IdentifierType { UpdateManyWhereCombinationInput(RelationField), UpdateOneWhereCombinationInput(RelationField), UpdateToOneRelWhereCombinationInput(RelationField), + UpdateManyAndReturnOutput(Model), WhereInput(ParentContainer), WhereUniqueInput(Model), Raw(String), @@ -315,6 +316,9 @@ impl std::fmt::Display for IdentifierType { _ => write!(f, "{}UncheckedUpdateManyInput", model.name()), }, IdentifierType::RelationLoadStrategy => write!(f, "RelationLoadStrategy"), + IdentifierType::UpdateManyAndReturnOutput(model) => { + write!(f, "UpdateMany{}AndReturnOutputType", model.name()) + } } } } diff --git a/query-engine/schema/src/query_schema.rs b/query-engine/schema/src/query_schema.rs index f8fede0ea355..a970329383ec 100644 --- a/query-engine/schema/src/query_schema.rs +++ b/query-engine/schema/src/query_schema.rs @@ -247,6 +247,7 @@ pub enum QueryTag { CreateManyAndReturn, UpdateOne, UpdateMany, + UpdateManyAndReturn, DeleteOne, DeleteMany, UpsertOne, @@ -273,6 +274,7 @@ impl fmt::Display for QueryTag { Self::CreateManyAndReturn => "createManyAndReturn", Self::UpdateOne => "updateOne", Self::UpdateMany => "updateMany", + Self::UpdateManyAndReturn => "updateManyAndReturn", Self::DeleteOne => "deleteOne", Self::DeleteMany => "deleteMany", Self::UpsertOne => "upsertOne", @@ -302,6 +304,7 @@ impl From<&str> for QueryTag { "createManyAndReturn" => Self::CreateManyAndReturn, "updateOne" => Self::UpdateOne, "updateMany" => Self::UpdateMany, + "updateManyAndReturn" => Self::UpdateManyAndReturn, "deleteOne" => Self::DeleteOne, "deleteMany" => Self::DeleteMany, "upsertOne" => Self::UpsertOne,