qe: support _count in joined queries (#4678)

Implement relation aggregations support for the `join` strategy:

* Build relation aggregation queries in the new query builder

* Add new methods to the abstractions that represent selections to account for
  differences in representation in queries that use or don't use JSON objects

* Implement coercion and serialization logic

Next step: prisma/team-orm#903

Part of: prisma/team-orm#700
Closes: prisma/team-orm#902

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

@@ -5,18 +5,26 @@ use std::{io, str::FromStr};
 
 use crate::{query_arguments_ext::QueryArgumentsExt, SqlError};
 
+pub(crate) enum IndexedSelection<'a> {
+    Relation(&'a RelationSelection),
+    Virtual(&'a str),
+}
+
 /// Coerces relations resolved as JSON to PrismaValues.
 /// Note: Some in-memory processing is baked into this function too for performance reasons.
 pub(crate) fn coerce_record_with_json_relation(
     record: &mut Record,
-    rs_indexes: Vec<(usize, &RelationSelection)>,
+    indexes: &[(usize, IndexedSelection<'_>)],
 ) -> crate::Result<()> {
-    for (val_idx, rs) in rs_indexes {
-        let val = record.values.get_mut(val_idx).unwrap();
+    for (val_idx, kind) in indexes {
+        let val = record.values.get_mut(*val_idx).unwrap();
         // TODO(perf): Find ways to avoid serializing and deserializing multiple times.
         let json_val: serde_json::Value = serde_json::from_str(val.as_json().unwrap()).unwrap();
 
-        *val = coerce_json_relation_to_pv(json_val, rs)?;
+        *val = match kind {
+            IndexedSelection::Relation(rs) => coerce_json_relation_to_pv(json_val, rs)?,
+            IndexedSelection::Virtual(name) => coerce_json_virtual_field_to_pv(name, json_val)?,
+        };
     }
 
     Ok(())
@@ -57,7 +65,7 @@ fn coerce_json_relation_to_pv(value: serde_json::Value, rs: &RelationSelection)
                 .map(|value| coerce_json_relation_to_pv(value, rs));
 
             // TODO(HACK): We probably want to update the sql builder instead to not aggregate to-one relations as array
-            // If the arary is empty, it means there's no relations, so we coerce it to
+            // If the array is empty, it means there's no relations, so we coerce it to
             if let Some(val) = coerced {
                 val
             } else {
@@ -69,16 +77,20 @@ fn coerce_json_relation_to_pv(value: serde_json::Value, rs: &RelationSelection)
             let related_model = rs.field.related_model();
 
             for (key, value) in obj {
-                match related_model.fields().all().find(|f| f.db_name() == key).unwrap() {
-                    Field::Scalar(sf) => {
+                match related_model.fields().all().find(|f| f.db_name() == key) {
+                    Some(Field::Scalar(sf)) => {
                         map.push((key, coerce_json_scalar_to_pv(value, &sf)?));
                     }
-                    Field::Relation(rf) => {
+                    Some(Field::Relation(rf)) => {
                         // TODO: optimize this
                         if let Some(nested_selection) = relations.iter().find(|rs| rs.field == rf) {
                             map.push((key, coerce_json_relation_to_pv(value, nested_selection)?));
                         }
                     }
+                    None => {
+                        let coerced_value = coerce_json_virtual_field_to_pv(&key, value)?;
+                        map.push((key, coerced_value));
+                    }
                     _ => (),
                 }
             }
@@ -191,27 +203,51 @@ pub(crate) fn coerce_json_scalar_to_pv(value: serde_json::Value, sf: &ScalarFiel
     }
 }
 
+fn coerce_json_virtual_field_to_pv(key: &str, value: serde_json::Value) -> crate::Result<PrismaValue> {
+    match value {
+        serde_json::Value::Object(obj) => {
+            let values: crate::Result<Vec<_>> = obj
+                .into_iter()
+                .map(|(key, value)| coerce_json_virtual_field_to_pv(&key, value).map(|value| (key, value)))
+                .collect();
+            Ok(PrismaValue::Object(values?))
+        }
+
+        serde_json::Value::Number(num) => num
+            .as_i64()
+            .ok_or_else(|| {
+                build_generic_conversion_error(format!(
+                    "Unexpected numeric value {num} for virtual field '{key}': only integers are supported"
+                ))
+            })
+            .map(PrismaValue::Int),
+
+        _ => Err(build_generic_conversion_error(format!(
+            "Field '{key}' is not a model field and doesn't have a supported type for a virtual field"
+        ))),
+    }
+}
+
 fn build_conversion_error(sf: &ScalarField, from: &str, to: &str) -> SqlError {
     let container_name = sf.container().name();
     let field_name = sf.name();
 
-    let error = io::Error::new(
-        io::ErrorKind::InvalidData,
-        format!("Unexpected conversion failure for field {container_name}.{field_name} from {from} to {to}."),
-    );
-
-    SqlError::ConversionError(error.into())
+    build_generic_conversion_error(format!(
+        "Unexpected conversion failure for field {container_name}.{field_name} from {from} to {to}."
+    ))
 }
 
 fn build_conversion_error_with_reason(sf: &ScalarField, from: &str, to: &str, reason: &str) -> SqlError {
     let container_name = sf.container().name();
     let field_name = sf.name();
 
-    let error = io::Error::new(
-        io::ErrorKind::InvalidData,
-        format!("Unexpected conversion failure for field {container_name}.{field_name} from {from} to {to}. Reason: ${reason}"),
-    );
+    build_generic_conversion_error(format!(
+        "Unexpected conversion failure for field {container_name}.{field_name} from {from} to {to}. Reason: {reason}"
+    ))
+}
 
+fn build_generic_conversion_error(message: String) -> SqlError {
+    let error = io::Error::new(io::ErrorKind::InvalidData, message);
     SqlError::ConversionError(error.into())
 }
 

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

@@ -1,4 +1,4 @@
-use super::coerce::coerce_record_with_json_relation;
+use super::coerce::{coerce_record_with_json_relation, IndexedSelection};
 use crate::{
     column_metadata,
     model_extensions::*,
@@ -33,9 +33,14 @@ pub(crate) async fn get_single_record_joins(
     selected_fields: &FieldSelection,
     ctx: &Context<'_>,
 ) -> crate::Result<Option<SingleRecord>> {
-    let field_names: Vec<_> = selected_fields.db_names().collect();
-    let idents = selected_fields.type_identifiers_with_arities();
-    let rs_indexes = get_relation_selection_indexes(selected_fields.relations().collect(), &field_names);
+    let field_names: Vec<_> = selected_fields.db_names_grouping_virtuals().collect();
+    let idents = selected_fields.type_identifiers_with_arities_grouping_virtuals();
+
+    let indexes = get_selection_indexes(
+        selected_fields.relations().collect(),
+        selected_fields.virtuals().collect(),
+        &field_names,
+    );
 
     let query = query_builder::select::SelectBuilder::default().build(
         QueryArguments::from((model.clone(), filter.clone())),
@@ -46,7 +51,7 @@ pub(crate) async fn get_single_record_joins(
     let mut record = execute_find_one(conn, query, &idents, &field_names, ctx).await?;
 
     if let Some(record) = record.as_mut() {
-        coerce_record_with_json_relation(record, rs_indexes)?;
+        coerce_record_with_json_relation(record, &indexes)?;
     };
 
     Ok(record.map(|record| SingleRecord { record, field_names }))
@@ -125,10 +130,15 @@ pub(crate) async fn get_many_records_joins(
     selected_fields: &FieldSelection,
     ctx: &Context<'_>,
 ) -> crate::Result<ManyRecords> {
-    let field_names: Vec<_> = selected_fields.db_names().collect();
-    let idents = selected_fields.type_identifiers_with_arities();
+    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());
-    let rs_indexes = get_relation_selection_indexes(selected_fields.relations().collect(), &field_names);
+
+    let indexes = get_selection_indexes(
+        selected_fields.relations().collect(),
+        selected_fields.virtuals().collect(),
+        &field_names,
+    );
 
     let mut records = ManyRecords::new(field_names.clone());
 
@@ -151,7 +161,7 @@ pub(crate) async fn get_many_records_joins(
         let mut record = Record::from(item);
 
         // Coerces json values to prisma values
-        coerce_record_with_json_relation(&mut record, rs_indexes.clone())?;
+        coerce_record_with_json_relation(&mut record, &indexes)?;
 
         records.push(record)
     }
@@ -395,18 +405,27 @@ async fn group_by_aggregate(
         .collect())
 }
 
-/// Find the indexes of the relation records to traverse a set of records faster when coercing JSON values
-fn get_relation_selection_indexes<'a>(
-    selections: Vec<&'a RelationSelection>,
-    field_names: &[String],
-) -> Vec<(usize, &'a RelationSelection)> {
-    let mut output: Vec<(usize, &RelationSelection)> = Vec::new();
-
-    for (idx, field_name) in field_names.iter().enumerate() {
-        if let Some(rs) = selections.iter().find(|rq| rq.field.name() == *field_name) {
-            output.push((idx, rs));
-        }
-    }
-
-    output
+/// Find the indexes of the relation records and the virtual selection objects to traverse a set of
+/// records faster when coercing JSON values.
+fn get_selection_indexes<'a>(
+    relations: Vec<&'a RelationSelection>,
+    virtuals: Vec<&'a VirtualSelection>,
+    field_names: &'a [String],
+) -> Vec<(usize, IndexedSelection<'a>)> {
+    field_names
+        .iter()
+        .enumerate()
+        .filter_map(|(idx, field_name)| {
+            relations
+                .iter()
+                .find_map(|rs| (rs.field.name() == field_name).then_some(IndexedSelection::Relation(rs)))
+                .or_else(|| {
+                    virtuals.iter().find_map(|vs| {
+                        let obj_name = vs.serialized_name().0;
+                        (obj_name == field_name).then_some(IndexedSelection::Virtual(obj_name))
+                    })
+                })
+                .map(|indexed_selection| (idx, indexed_selection))
+        })
+        .collect()
 }