qe: partially refactor relation aggregation selection

query-engine/connectors/mongodb-query-connector/src/error.rs

@@ -4,7 +4,7 @@ use mongodb::{
     bson::{self, extjson},
     error::{CommandError, Error as DriverError, TRANSIENT_TRANSACTION_ERROR},
 };
-use query_structure::{CompositeFieldRef, Field, ScalarFieldRef, SelectedField};
+use query_structure::{CompositeFieldRef, Field, ScalarFieldRef, SelectedField, VirtualSelection};
 use regex::Regex;
 use thiserror::Error;
 use user_facing_errors::query_engine::DatabaseConstraint;
@@ -274,6 +274,7 @@ pub trait DecorateErrorWithFieldInformationExtension {
     fn decorate_with_scalar_field_info(self, sf: &ScalarFieldRef) -> Self;
     fn decorate_with_field_name(self, field_name: &str) -> Self;
     fn decorate_with_composite_field_info(self, cf: &CompositeFieldRef) -> Self;
+    fn decorate_with_virtual_field_info(self, vs: &VirtualSelection) -> Self;
 }
 
 impl<T> DecorateErrorWithFieldInformationExtension for crate::Result<T> {
@@ -286,6 +287,7 @@ impl<T> DecorateErrorWithFieldInformationExtension for crate::Result<T> {
             SelectedField::Scalar(sf) => self.decorate_with_scalar_field_info(sf),
             SelectedField::Composite(composite_sel) => self.decorate_with_composite_field_info(&composite_sel.field),
             SelectedField::Relation(_) => unreachable!(),
+            SelectedField::Virtual(vs) => self.decorate_with_virtual_field_info(vs),
         }
     }
 
@@ -300,4 +302,8 @@ impl<T> DecorateErrorWithFieldInformationExtension for crate::Result<T> {
     fn decorate_with_composite_field_info(self, cf: &CompositeFieldRef) -> Self {
         self.map_err(|err| err.decorate_with_field_name(cf.name()))
     }
+
+    fn decorate_with_virtual_field_info(self, vs: &VirtualSelection) -> Self {
+        self.map_err(|err| err.decorate_with_field_name(&vs.db_alias()))
+    }
 }

query-engine/connectors/mongodb-query-connector/src/interface/connection.rs

@@ -6,8 +6,7 @@ use crate::{
 };
 use async_trait::async_trait;
 use connector_interface::{
-    Connection, ConnectionLike, ReadOperations, RelAggregationSelection, Transaction, UpdateType, WriteArgs,
-    WriteOperations,
+    Connection, ConnectionLike, ReadOperations, Transaction, UpdateType, WriteArgs, WriteOperations,
 };
 use mongodb::{ClientSession, Database};
 use query_structure::{prelude::*, RelationLoadStrategy, SelectionResult};
@@ -234,7 +233,6 @@ impl ReadOperations for MongoDbConnection {
         model: &Model,
         filter: &query_structure::Filter,
         selected_fields: &FieldSelection,
-        aggr_selections: &[RelAggregationSelection],
         _relation_load_strategy: RelationLoadStrategy,
         _trace_id: Option<String>,
     ) -> connector_interface::Result<Option<SingleRecord>> {
@@ -244,7 +242,6 @@ impl ReadOperations for MongoDbConnection {
             model,
             filter,
             selected_fields,
-            aggr_selections,
         ))
         .await
     }
@@ -254,7 +251,6 @@ impl ReadOperations for MongoDbConnection {
         model: &Model,
         query_arguments: query_structure::QueryArguments,
         selected_fields: &FieldSelection,
-        aggregation_selections: &[RelAggregationSelection],
         _relation_load_strategy: RelationLoadStrategy,
         _trace_id: Option<String>,
     ) -> connector_interface::Result<ManyRecords> {
@@ -264,7 +260,6 @@ impl ReadOperations for MongoDbConnection {
             model,
             query_arguments,
             selected_fields,
-            aggregation_selections,
         ))
         .await
     }

query-engine/connectors/mongodb-query-connector/src/interface/transaction.rs

@@ -3,9 +3,7 @@ use crate::{
     error::MongoError,
     root_queries::{aggregate, read, write},
 };
-use connector_interface::{
-    ConnectionLike, ReadOperations, RelAggregationSelection, Transaction, UpdateType, WriteOperations,
-};
+use connector_interface::{ConnectionLike, ReadOperations, Transaction, UpdateType, WriteOperations};
 use mongodb::options::{Acknowledgment, ReadConcern, TransactionOptions, WriteConcern};
 use query_engine_metrics::{decrement_gauge, increment_gauge, metrics, PRISMA_CLIENT_QUERIES_ACTIVE};
 use query_structure::{RelationLoadStrategy, SelectionResult};
@@ -264,7 +262,6 @@ impl<'conn> ReadOperations for MongoDbTransaction<'conn> {
         model: &Model,
         filter: &query_structure::Filter,
         selected_fields: &FieldSelection,
-        aggr_selections: &[RelAggregationSelection],
         _relation_load_strategy: RelationLoadStrategy,
         _trace_id: Option<String>,
     ) -> connector_interface::Result<Option<SingleRecord>> {
@@ -274,7 +271,6 @@ impl<'conn> ReadOperations for MongoDbTransaction<'conn> {
             model,
             filter,
             selected_fields,
-            aggr_selections,
         ))
         .await
     }
@@ -284,7 +280,6 @@ impl<'conn> ReadOperations for MongoDbTransaction<'conn> {
         model: &Model,
         query_arguments: query_structure::QueryArguments,
         selected_fields: &FieldSelection,
-        aggregation_selections: &[RelAggregationSelection],
         _relation_load_strategy: RelationLoadStrategy,
         _trace_id: Option<String>,
     ) -> connector_interface::Result<ManyRecords> {
@@ -294,7 +289,6 @@ impl<'conn> ReadOperations for MongoDbTransaction<'conn> {
             model,
             query_arguments,
             selected_fields,
-            aggregation_selections,
         ))
         .await
     }

query-engine/connectors/mongodb-query-connector/src/output_meta.rs

@@ -1,4 +1,4 @@
-use connector_interface::{AggregationSelection, RelAggregationSelection};
+use connector_interface::AggregationSelection;
 use indexmap::IndexMap;
 use query_structure::{
     ast::FieldArity, DefaultKind, FieldSelection, PrismaValue, ScalarFieldRef, SelectedField, TypeIdentifier,
@@ -48,18 +48,12 @@ impl CompositeOutputMeta {
     }
 }
 
-pub fn from_selected_fields(
-    selected_fields: &FieldSelection,
-    aggregation_selections: &[RelAggregationSelection],
-) -> OutputMetaMapping {
-    let selections: Vec<&SelectedField> = selected_fields.selections().collect();
-    from_selections(&selections, aggregation_selections)
+pub fn from_selected_fields(selected_fields: &FieldSelection) -> OutputMetaMapping {
+    let selections: Vec<_> = selected_fields.selections().collect();
+    from_selections(&selections)
 }
 
-pub fn from_selections(
-    selected_fields: &[&SelectedField],
-    aggregation_selections: &[RelAggregationSelection],
-) -> OutputMetaMapping {
+pub fn from_selections(selected_fields: &[&SelectedField]) -> OutputMetaMapping {
     let mut map = OutputMetaMapping::new();
 
     for selection in selected_fields {
@@ -70,7 +64,7 @@ pub fn from_selections(
 
             SelectedField::Composite(cs) => {
                 let selections: Vec<&SelectedField> = cs.selections.iter().collect();
-                let inner = from_selections(&selections, &[]);
+                let inner = from_selections(&selections);
 
                 map.insert(
                     cs.field.db_name().to_owned(),
@@ -80,12 +74,22 @@ pub fn from_selections(
                     }),
                 );
             }
+
             SelectedField::Relation(_) => unreachable!(),
-        }
-    }
 
-    for selection in aggregation_selections {
-        map.insert(selection.db_alias(), from_rel_aggregation_selection(selection));
+            SelectedField::Virtual(vs) => {
+                let (ident, arity) = vs.type_identifier_with_arity();
+
+                map.insert(
+                    vs.db_alias(),
+                    OutputMeta::Scalar(ScalarOutputMeta {
+                        ident,
+                        default: None,
+                        list: matches!(arity, FieldArity::List),
+                    }),
+                );
+            }
+        }
     }
 
     map
@@ -126,18 +130,6 @@ pub fn from_aggregation_selection(selection: &AggregationSelection) -> OutputMet
     map
 }
 
-/// Mapping for one specific relation aggregation selection.
-/// DB alias -> OutputMeta
-pub fn from_rel_aggregation_selection(aggr_selection: &RelAggregationSelection) -> OutputMeta {
-    let (ident, arity) = aggr_selection.type_identifier_with_arity();
-
-    OutputMeta::Scalar(ScalarOutputMeta {
-        ident,
-        default: None,
-        list: matches!(arity, FieldArity::List),
-    })
-}
-
 impl From<ScalarOutputMeta> for OutputMeta {
     fn from(s: ScalarOutputMeta) -> Self {
         Self::Scalar(s)

query-engine/connectors/mongodb-query-connector/src/projection.rs

@@ -26,7 +26,10 @@ fn path_prefixed_selection(doc: &mut Document, parent_paths: Vec<String>, select
                 parent_paths.push(cs.field.db_name().to_owned());
                 path_prefixed_selection(doc, parent_paths, cs.selections);
             }
+
             query_structure::SelectedField::Relation(_) => unreachable!(),
+
+            query_structure::SelectedField::Virtual(_) => {}
         }
     }
 }

query-engine/connectors/mongodb-query-connector/src/query_builder/read_query_builder.rs

@@ -10,14 +10,14 @@ use crate::{
     root_queries::observing,
     vacuum_cursor, BsonTransform, IntoBson,
 };
-use connector_interface::{AggregationSelection, RelAggregationSelection};
+use connector_interface::AggregationSelection;
 use itertools::Itertools;
 use mongodb::{
     bson::{doc, Document},
     options::AggregateOptions,
     ClientSession, Collection,
 };
-use query_structure::{FieldSelection, Filter, Model, QueryArguments, ScalarFieldRef};
+use query_structure::{FieldSelection, Filter, Model, QueryArguments, ScalarFieldRef, VirtualSelection};
 use std::convert::TryFrom;
 
 // Mongo Driver broke usage of the simple API, can't be used by us anymore.
@@ -355,13 +355,13 @@ impl MongoReadQueryBuilder {
     }
 
     /// Adds the necessary joins and the associated selections to the projection
-    pub fn with_aggregation_selections(
+    pub fn with_virtual_fields<'a>(
         mut self,
-        aggregation_selections: &[RelAggregationSelection],
+        virtual_selections: impl Iterator<Item = &'a VirtualSelection>,
     ) -> crate::Result<Self> {
-        for aggr in aggregation_selections {
+        for aggr in virtual_selections {
             let join = match aggr {
-                RelAggregationSelection::Count(rf, filter) => {
+                VirtualSelection::RelationCount(rf, filter) => {
                     let filter = filter
                         .as_ref()
                         .map(|f| MongoFilterVisitor::new(FilterPrefix::default(), false).visit(f.clone()))

query-engine/connectors/mongodb-query-connector/src/root_queries/read.rs

@@ -3,7 +3,6 @@ use crate::{
     error::DecorateErrorWithFieldInformationExtension, output_meta, query_builder::MongoReadQueryBuilder,
     query_strings::Find, vacuum_cursor, IntoBson,
 };
-use connector_interface::RelAggregationSelection;
 use mongodb::{bson::doc, options::FindOptions, ClientSession, Database};
 use query_structure::*;
 use tracing::{info_span, Instrument};
@@ -15,7 +14,6 @@ pub async fn get_single_record<'conn>(
     model: &Model,
     filter: &Filter,
     selected_fields: &FieldSelection,
-    aggregation_selections: &[RelAggregationSelection],
 ) -> crate::Result<Option<SingleRecord>> {
     let coll = database.collection(model.db_name());
 
@@ -25,22 +23,19 @@ pub async fn get_single_record<'conn>(
         "db.statement" = &format_args!("db.{}.findOne(*)", coll.name())
     );
 
-    let meta_mapping = output_meta::from_selected_fields(selected_fields, aggregation_selections);
+    let meta_mapping = output_meta::from_selected_fields(selected_fields);
     let query_arguments: QueryArguments = (model.clone(), filter.clone()).into();
     let query = MongoReadQueryBuilder::from_args(query_arguments)?
         .with_model_projection(selected_fields.clone())?
-        .with_aggregation_selections(aggregation_selections)?
+        .with_virtual_fields(selected_fields.virtuals())?
         .build()?;
 
     let docs = query.execute(coll, session).instrument(span).await?;
 
     if docs.is_empty() {
         Ok(None)
     } else {
-        let field_names: Vec<_> = selected_fields
-            .db_names()
-            .chain(aggregation_selections.iter().map(|aggr_sel| aggr_sel.db_alias()))
-            .collect();
+        let field_names: Vec<_> = selected_fields.db_names().collect();
         let doc = docs.into_iter().next().unwrap();
         let record = document_to_record(doc, &field_names, &meta_mapping)?;
 
@@ -61,7 +56,6 @@ pub async fn get_many_records<'conn>(
     model: &Model,
     query_arguments: QueryArguments,
     selected_fields: &FieldSelection,
-    aggregation_selections: &[RelAggregationSelection],
 ) -> crate::Result<ManyRecords> {
     let coll = database.collection(model.db_name());
 
@@ -72,12 +66,9 @@ pub async fn get_many_records<'conn>(
     );
 
     let reverse_order = query_arguments.take.map(|t| t < 0).unwrap_or(false);
-    let field_names: Vec<_> = selected_fields
-        .db_names()
-        .chain(aggregation_selections.iter().map(|aggr_sel| aggr_sel.db_alias()))
-        .collect();
+    let field_names: Vec<_> = selected_fields.db_names().collect();
 
-    let meta_mapping = output_meta::from_selected_fields(selected_fields, aggregation_selections);
+    let meta_mapping = output_meta::from_selected_fields(selected_fields);
     let mut records = ManyRecords::new(field_names.clone());
 
     if let Some(0) = query_arguments.take {
@@ -86,7 +77,7 @@ pub async fn get_many_records<'conn>(
 
     let query = MongoReadQueryBuilder::from_args(query_arguments)?
         .with_model_projection(selected_fields.clone())?
-        .with_aggregation_selections(aggregation_selections)?
+        .with_virtual_fields(selected_fields.virtuals())?
         .build()?;
 
     let docs = query.execute(coll, session).instrument(span).await?;

query-engine/connectors/mongodb-query-connector/src/root_queries/write.rs

@@ -312,7 +312,7 @@ pub async fn delete_record<'conn>(
         cause: "Record to delete does not exist.".to_owned(),
     })?;
 
-    let meta_mapping = output_meta::from_selected_fields(&selected_fields, &[]);
+    let meta_mapping = output_meta::from_selected_fields(&selected_fields);
     let field_names: Vec<_> = selected_fields.db_names().collect();
     let record = document_to_record(document, &field_names, &meta_mapping)?;
     Ok(SingleRecord { record, field_names })

query-engine/connectors/mongodb-query-connector/src/value.rs

@@ -24,6 +24,7 @@ impl IntoBson for (&SelectedField, PrismaValue) {
             SelectedField::Scalar(sf) => (sf, value).into_bson(),
             SelectedField::Composite(_) => todo!(), // [Composites] todo
             SelectedField::Relation(_) => unreachable!(),
+            SelectedField::Virtual(_) => unreachable!(),
         }
     }
 }