qe: fix mismatch between selection indexes for joins (#4705)

Both the old query builder and the new query builder append virtual
selections after all other selections in the query. This means we have
to take it into account instead of relying on the fields in the result
set to be in the same order as in `FieldSelection`.

The old code path converted the selection into virtuals-last form but
the new one didn't, which resulted in the cached indexes for the
relations and virtuals to be mixed up when selecting relation
aggregations before the relations in the query.

Now the new code path does the same transformation.

Fixes: prisma/team-orm#927

query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/mod.rs

@@ -30,3 +30,4 @@ mod prisma_7072;
 mod prisma_7434;
 mod prisma_8265;
 mod prisma_engines_4286;
+mod team_orm_927;

query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/team_orm_927.rs

@@ -0,0 +1,90 @@
+//! Regression test for https://github.com/prisma/team-orm/issues/927
+
+use query_engine_tests::*;
+
+#[test_suite(schema(schema))]
+mod count_before_relation {
+    fn schema() -> String {
+        indoc! {
+            r#"
+            model Parent {
+                #id(id, Int, @id)
+                children Child[]
+            }
+
+            model Child {
+                #id(id, Int, @id)
+                parentId Int
+                parent   Parent @relation(fields: [parentId], references: [id])
+            }
+            "#
+        }
+        .to_owned()
+    }
+
+    #[connector_test]
+    async fn find_unique(runner: Runner) -> TestResult<()> {
+        seed(&runner).await?;
+
+        insta::assert_snapshot!(
+            run_query!(
+                runner,
+                r#"
+                query {
+                    findUniqueParent(
+                        where: { id: 1 }
+                    ) {
+                        _count { children }
+                        children { id }
+                    }
+                }
+                "#
+            ),
+            @r###"{"data":{"findUniqueParent":{"_count":{"children":1},"children":[{"id":1}]}}}"###
+        );
+
+        Ok(())
+    }
+
+    #[connector_test]
+    async fn find_many(runner: Runner) -> TestResult<()> {
+        seed(&runner).await?;
+
+        insta::assert_snapshot!(
+            run_query!(
+                runner,
+                r#"
+                query {
+                    findManyParent {
+                        _count { children }
+                        children { id }
+                    }
+                }
+                "#
+            ),
+            @r###"{"data":{"findManyParent":[{"_count":{"children":1},"children":[{"id":1}]}]}}"###
+        );
+
+        Ok(())
+    }
+
+    async fn seed(runner: &Runner) -> TestResult<()> {
+        run_query!(
+            runner,
+            r#"
+            mutation {
+                createOneParent(
+                    data: {
+                        id: 1,
+                        children: {
+                            create: { id: 1 }
+                        }
+                    }
+                ) { id }
+            }
+            "#
+        );
+
+        Ok(())
+    }
+}

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

@@ -33,6 +33,7 @@ pub(crate) async fn get_single_record_joins(
     selected_fields: &FieldSelection,
     ctx: &Context<'_>,
 ) -> crate::Result<Option<SingleRecord>> {
+    let selected_fields = selected_fields.to_virtuals_last();
     let field_names: Vec<_> = selected_fields.db_names_grouping_virtuals().collect();
     let idents = selected_fields.type_identifiers_with_arities_grouping_virtuals();
 
@@ -44,7 +45,7 @@ pub(crate) async fn get_single_record_joins(
 
     let query = query_builder::select::SelectBuilder::build(
         QueryArguments::from((model.clone(), filter.clone())),
-        selected_fields,
+        &selected_fields,
         ctx,
     );
 
@@ -130,6 +131,7 @@ pub(crate) async fn get_many_records_joins(
     selected_fields: &FieldSelection,
     ctx: &Context<'_>,
 ) -> crate::Result<ManyRecords> {
+    let selected_fields = selected_fields.to_virtuals_last();
     let field_names: Vec<_> = selected_fields.db_names_grouping_virtuals().collect();
     let idents = selected_fields.type_identifiers_with_arities_grouping_virtuals();
     let meta = column_metadata::create(field_names.as_slice(), idents.as_slice());
@@ -155,7 +157,7 @@ pub(crate) async fn get_many_records_joins(
         _ => (),
     };
 
-    let query = query_builder::select::SelectBuilder::build(query_arguments.clone(), selected_fields, ctx);
+    let query = query_builder::select::SelectBuilder::build(query_arguments.clone(), &selected_fields, ctx);
 
     for item in conn.filter(query.into(), meta.as_slice(), ctx).await?.into_iter() {
         let mut record = Record::from(item);

query-engine/query-structure/src/field_selection.rs

@@ -68,6 +68,10 @@ impl FieldSelection {
         FieldSelection::new(non_virtuals.into_iter().chain(virtuals).collect())
     }
 
+    pub fn to_virtuals_last(&self) -> Self {
+        self.clone().into_virtuals_last()
+    }
+
     /// Returns the selections, grouping the virtual fields that are wrapped into objects in the
     /// query (like `_count`) and returning only the first virtual field in each of those groups.
     /// This is useful when we want to treat the group as a whole but we don't need the information