fix: extract aggregation after in-memory processing

query-engine/connector-test-kit-rs/query-engine-tests/tests/queries/aggregation/many_count_relation.rs

@@ -1,6 +1,6 @@
 use query_engine_tests::*;
 
-#[test_suite(schema(schema))]
+#[test_suite(schema(schema), exclude(MongoDb))]
 mod many_count_rel {
     use indoc::indoc;
     use query_engine_tests::run_query;
@@ -30,7 +30,7 @@ mod many_count_rel {
     }
 
     // "Counting with no records in the database" should "return 0"
-    #[connector_test(exclude(MongoDb))] // TODO(dom): Not working on mongo
+    #[connector_test] // TODO(dom): Not working on mongo
     async fn no_rel_records(runner: Runner) -> TestResult<()> {
         create_row(&runner, r#"{ id: 1, title: "a" }"#).await?;
 
@@ -47,7 +47,7 @@ mod many_count_rel {
     }
 
     //"Counting one2m and m2m records" should "work"
-    #[connector_test(exclude(MongoDb))] // TODO(dom): Not working on mongo
+    #[connector_test] // TODO(dom): Not working on mongo
     async fn count_one2m_m2m(runner: Runner) -> TestResult<()> {
         // 1 comment / 2 categories
         create_row(
@@ -85,7 +85,7 @@ mod many_count_rel {
     }
 
     // "Counting with some records and filters" should "not affect the count"
-    #[connector_test(exclude(MongoDb))] // TODO(dom): Not working on mongo
+    #[connector_test] // TODO(dom): Not working on mongo
     async fn count_with_filters(runner: Runner) -> TestResult<()> {
         // 4 comment / 4 categories
         create_row(
@@ -151,7 +151,7 @@ mod many_count_rel {
 
     // Counting nested one2m and m2m should work
     // TODO(dom): Not working on mongo
-    #[connector_test(schema(schema_nested), exclude(MongoDb))]
+    #[connector_test(schema(schema_nested))]
     async fn nested_count_one2m_m2m(runner: Runner) -> TestResult<()> {
         run_query!(
             &runner,
@@ -215,6 +215,54 @@ mod many_count_rel {
         Ok(())
     }
 
+    fn schema_inmemory_process() -> String {
+        let schema = indoc! {
+            r#"model Post {
+              #id(id, Int, @id, @default(autoincrement()))
+              comments  Comment[]
+              createdAt DateTime  @default(now())
+              updatedAt DateTime  @updatedAt
+            }
+
+            model Comment {
+              #id(id, Int, @id, @default(autoincrement()))
+              post      Post     @relation(fields: [postId], references: [id])
+              postId    Int
+              createdAt DateTime @default(now())
+              updatedAt DateTime @updatedAt
+            }"#
+        };
+
+        schema.to_owned()
+    }
+
+    // Regression test for https://github.com/prisma/prisma/issues/8050
+    // Ensures aggregation rows are properly extracted even when in-memory processing is applied to the records
+    #[connector_test(schema(schema_inmemory_process))]
+    async fn works_with_inmemory_args_processing(runner: Runner) -> TestResult<()> {
+        create_row(&runner, r#"{ comments: { create: [{}, {}] } }"#).await?;
+        create_row(&runner, r#"{ comments: { create: [{}, {}, {}, {}] } }"#).await?;
+        create_row(&runner, r#"{}"#).await?;
+        create_row(&runner, r#"{}"#).await?;
+
+        insta::assert_snapshot!(
+          run_query!(&runner, r#"query { findManyPost(
+            orderBy: { createdAt: "desc" }
+            cursor: { id: 4 }
+            skip: 1
+            take: 6
+          ) {
+            id
+            _count {
+              comments
+            }
+          } }"#),
+          @r###"{"data":{"findManyPost":[{"id":3,"_count":{"comments":0}},{"id":2,"_count":{"comments":4}},{"id":1,"_count":{"comments":2}}]}}"###
+        );
+
+        Ok(())
+    }
+
     async fn create_row(runner: &Runner, data: &str) -> TestResult<()> {
         runner
             .query(format!("mutation {{ createOnePost(data: {}) {{ id }} }}", data))

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

@@ -88,10 +88,11 @@ fn read_many(
                     &query.aggregation_selections,
                 )
                 .await?;
+            let scalars = processor.apply(scalars);
             let (scalars, aggregation_rows) =
                 extract_aggregation_rows_from_scalars(scalars.clone(), query.aggregation_selections);
 
-            (processor.apply(scalars), aggregation_rows)
+            (scalars, aggregation_rows)
         } else {
             let scalars = tx
                 .get_many_records(
@@ -103,7 +104,6 @@ fn read_many(
                 .await?;
             let (scalars, aggregation_rows) =
                 extract_aggregation_rows_from_scalars(scalars.clone(), query.aggregation_selections);
-
             (scalars, aggregation_rows)
         };
 
@@ -215,7 +215,7 @@ fn process_nested<'conn>(
 /// findManyX { _count { rel_1, rel 2 } }
 /// Output object types are typically used for selecting relations, so they're are queried in a different request
 /// In the case of relation aggregations though, we query that data along side the request sent for the base model ("X" in the query above)
-/// This means the SQL result we get back from the database contains additional aggregation data that needs to be remapped according to the shema
+/// This means the SQL result we get back from the database contains additional aggregation data that needs to be remapped according to the schema
 /// This function takes care of removing the aggregation data from the database result and collects it separately
 /// so that it can be serialized separately later according to the schema
 pub fn extract_aggregation_rows_from_scalars(