diff --git a/query-engine/connectors/mongodb-query-connector/src/error.rs b/query-engine/connectors/mongodb-query-connector/src/error.rs index 6c39d2c033b1..a93350168387 100644 --- a/query-engine/connectors/mongodb-query-connector/src/error.rs +++ b/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 DecorateErrorWithFieldInformationExtension for crate::Result { @@ -286,6 +287,7 @@ impl DecorateErrorWithFieldInformationExtension for crate::Result { 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 DecorateErrorWithFieldInformationExtension for crate::Result { 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())) + } } 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 c4929790ecd3..09cb46eae6fa 100644 --- a/query-engine/connectors/mongodb-query-connector/src/interface/connection.rs +++ b/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, ) -> connector_interface::Result> { @@ -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, ) -> connector_interface::Result { @@ -264,7 +260,6 @@ impl ReadOperations for MongoDbConnection { model, query_arguments, selected_fields, - aggregation_selections, )) .await } 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 220e31a25adb..4c3b1dfec68f 100644 --- a/query-engine/connectors/mongodb-query-connector/src/interface/transaction.rs +++ b/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, ) -> connector_interface::Result> { @@ -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, ) -> connector_interface::Result { @@ -294,7 +289,6 @@ impl<'conn> ReadOperations for MongoDbTransaction<'conn> { model, query_arguments, selected_fields, - aggregation_selections, )) .await } diff --git a/query-engine/connectors/mongodb-query-connector/src/output_meta.rs b/query-engine/connectors/mongodb-query-connector/src/output_meta.rs index 1100c9d436ad..d1937bf7fee7 100644 --- a/query-engine/connectors/mongodb-query-connector/src/output_meta.rs +++ b/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 for OutputMeta { fn from(s: ScalarOutputMeta) -> Self { Self::Scalar(s) diff --git a/query-engine/connectors/mongodb-query-connector/src/projection.rs b/query-engine/connectors/mongodb-query-connector/src/projection.rs index cbb16e0a097c..96948fc0e027 100644 --- a/query-engine/connectors/mongodb-query-connector/src/projection.rs +++ b/query-engine/connectors/mongodb-query-connector/src/projection.rs @@ -26,7 +26,10 @@ fn path_prefixed_selection(doc: &mut Document, parent_paths: Vec, 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(_) => {} } } } diff --git a/query-engine/connectors/mongodb-query-connector/src/query_builder/read_query_builder.rs b/query-engine/connectors/mongodb-query-connector/src/query_builder/read_query_builder.rs index e8cdf26caba0..27185de5c917 100644 --- a/query-engine/connectors/mongodb-query-connector/src/query_builder/read_query_builder.rs +++ b/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, ) -> crate::Result { - 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())) diff --git a/query-engine/connectors/mongodb-query-connector/src/root_queries/read.rs b/query-engine/connectors/mongodb-query-connector/src/root_queries/read.rs index c4bceeaab7f9..9fc322312265 100644 --- a/query-engine/connectors/mongodb-query-connector/src/root_queries/read.rs +++ b/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> { let coll = database.collection(model.db_name()); @@ -25,11 +23,11 @@ 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?; @@ -37,10 +35,7 @@ pub async fn get_single_record<'conn>( 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 { 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?; diff --git a/query-engine/connectors/mongodb-query-connector/src/root_queries/write.rs b/query-engine/connectors/mongodb-query-connector/src/root_queries/write.rs index 4527748d7c4c..f418a007da4c 100644 --- a/query-engine/connectors/mongodb-query-connector/src/root_queries/write.rs +++ b/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 }) diff --git a/query-engine/connectors/mongodb-query-connector/src/value.rs b/query-engine/connectors/mongodb-query-connector/src/value.rs index 9faecaa13f4c..b0d4946f23cf 100644 --- a/query-engine/connectors/mongodb-query-connector/src/value.rs +++ b/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!(), } } } diff --git a/query-engine/connectors/query-connector/src/interface.rs b/query-engine/connectors/query-connector/src/interface.rs index ea2cec9283ef..d42d6f0524b7 100644 --- a/query-engine/connectors/query-connector/src/interface.rs +++ b/query-engine/connectors/query-connector/src/interface.rs @@ -1,4 +1,4 @@ -use crate::{coerce_null_to_zero_value, NativeUpsert, WriteArgs}; +use crate::{NativeUpsert, WriteArgs}; use async_trait::async_trait; use prisma_value::PrismaValue; use query_structure::{ast::FieldArity, *}; @@ -176,47 +176,6 @@ pub enum AggregationResult { Max(ScalarFieldRef, PrismaValue), } -#[derive(Debug, Clone)] -pub enum RelAggregationSelection { - // Always a count(*) for now - Count(RelationFieldRef, Option), -} - -pub type RelAggregationRow = Vec; - -#[derive(Debug, Clone)] -pub enum RelAggregationResult { - Count(RelationFieldRef, PrismaValue), -} - -impl RelAggregationSelection { - pub fn db_alias(&self) -> String { - match self { - RelAggregationSelection::Count(rf, _) => { - format!("_aggr_count_{}", rf.name()) - } - } - } - - pub fn field_name(&self) -> &str { - match self { - RelAggregationSelection::Count(rf, _) => rf.name(), - } - } - - pub fn type_identifier_with_arity(&self) -> (TypeIdentifier, FieldArity) { - match self { - RelAggregationSelection::Count(_, _) => (TypeIdentifier::Int, FieldArity::Required), - } - } - - pub fn into_result(self, val: PrismaValue) -> RelAggregationResult { - match self { - RelAggregationSelection::Count(rf, _) => RelAggregationResult::Count(rf, coerce_null_to_zero_value(val)), - } - } -} - #[async_trait] pub trait ReadOperations { /// Gets a single record or `None` back from the database. @@ -230,7 +189,6 @@ pub trait ReadOperations { model: &Model, filter: &Filter, selected_fields: &FieldSelection, - aggregation_selections: &[RelAggregationSelection], relation_load_strategy: RelationLoadStrategy, trace_id: Option, ) -> crate::Result>; @@ -246,7 +204,6 @@ pub trait ReadOperations { model: &Model, query_arguments: QueryArguments, selected_fields: &FieldSelection, - aggregation_selections: &[RelAggregationSelection], relation_load_strategy: RelationLoadStrategy, trace_id: Option, ) -> crate::Result; diff --git a/query-engine/connectors/query-connector/src/write_args.rs b/query-engine/connectors/query-connector/src/write_args.rs index c89f4e51514f..c881ee9c2cfa 100644 --- a/query-engine/connectors/query-connector/src/write_args.rs +++ b/query-engine/connectors/query-connector/src/write_args.rs @@ -328,6 +328,7 @@ impl From<(&SelectedField, PrismaValue)> for WriteOperation { SelectedField::Scalar(sf) => (sf, pv).into(), SelectedField::Composite(cs) => (&cs.field, pv).into(), SelectedField::Relation(_) => todo!(), + SelectedField::Virtual(_) => todo!(), } } } @@ -462,7 +463,7 @@ pub fn merge_write_args(loaded_ids: Vec, incoming_args: WriteAr .pairs .iter() .enumerate() - .filter_map(|(i, (selection, _))| incoming_args.get_field_value(selection.db_name()).map(|val| (i, val))) + .filter_map(|(i, (selection, _))| incoming_args.get_field_value(&selection.db_name()).map(|val| (i, val))) .collect(); loaded_ids 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 eca2372afb22..457fb6136b52 100644 --- a/query-engine/connectors/sql-query-connector/src/database/connection.rs +++ b/query-engine/connectors/sql-query-connector/src/database/connection.rs @@ -3,7 +3,7 @@ use super::{catch, transaction::SqlConnectorTransaction}; use crate::{database::operations::*, Context, SqlError}; use async_trait::async_trait; -use connector::{ConnectionLike, RelAggregationSelection}; +use connector::ConnectionLike; use connector_interface::{ self as connector, AggregationRow, AggregationSelection, Connection, ReadOperations, RecordFilter, Transaction, WriteArgs, WriteOperations, @@ -86,7 +86,6 @@ where model: &Model, filter: &Filter, selected_fields: &FieldSelection, - aggr_selections: &[RelAggregationSelection], relation_load_strategy: RelationLoadStrategy, trace_id: Option, ) -> connector::Result> { @@ -99,7 +98,6 @@ where model, filter, selected_fields, - aggr_selections, relation_load_strategy, &ctx, ), @@ -112,7 +110,6 @@ where model: &Model, query_arguments: QueryArguments, selected_fields: &FieldSelection, - aggr_selections: &[RelAggregationSelection], relation_load_strategy: RelationLoadStrategy, trace_id: Option, ) -> connector::Result { @@ -124,7 +121,6 @@ where model, query_arguments, selected_fields, - aggr_selections, relation_load_strategy, &ctx, ), diff --git a/query-engine/connectors/sql-query-connector/src/database/operations/read.rs b/query-engine/connectors/sql-query-connector/src/database/operations/read.rs index 2d6de2472979..32e9dba67b79 100644 --- a/query-engine/connectors/sql-query-connector/src/database/operations/read.rs +++ b/query-engine/connectors/sql-query-connector/src/database/operations/read.rs @@ -17,23 +17,12 @@ pub(crate) async fn get_single_record( model: &Model, filter: &Filter, selected_fields: &FieldSelection, - aggr_selections: &[RelAggregationSelection], relation_load_strategy: RelationLoadStrategy, ctx: &Context<'_>, ) -> crate::Result> { match relation_load_strategy { RelationLoadStrategy::Join => get_single_record_joins(conn, model, filter, selected_fields, ctx).await, - RelationLoadStrategy::Query => { - get_single_record_wo_joins( - conn, - model, - filter, - &ModelProjection::from(selected_fields), - aggr_selections, - ctx, - ) - .await - } + RelationLoadStrategy::Query => get_single_record_wo_joins(conn, model, filter, selected_fields, ctx).await, } } @@ -67,30 +56,24 @@ pub(crate) async fn get_single_record_wo_joins( conn: &dyn Queryable, model: &Model, filter: &Filter, - selected_fields: &ModelProjection, - aggr_selections: &[RelAggregationSelection], + selected_fields: &FieldSelection, ctx: &Context<'_>, ) -> crate::Result> { + let selected_fields = selected_fields.without_relations().into_virtuals_last(); + let query = read::get_records( model, - selected_fields.as_columns(ctx).mark_all_selected(), - aggr_selections, + ModelProjection::from(&selected_fields) + .as_columns(ctx) + .mark_all_selected(), + selected_fields.virtuals(), filter, ctx, ); - let mut field_names: Vec<_> = selected_fields.db_names().collect(); - let mut aggr_field_names: Vec<_> = aggr_selections.iter().map(|aggr_sel| aggr_sel.db_alias()).collect(); - - field_names.append(&mut aggr_field_names); - - let mut idents = selected_fields.type_identifiers_with_arities(); - let mut aggr_idents = aggr_selections - .iter() - .map(|aggr_sel| aggr_sel.type_identifier_with_arity()) - .collect(); + let field_names: Vec<_> = selected_fields.db_names().collect(); - idents.append(&mut aggr_idents); + let idents = selected_fields.type_identifiers_with_arities(); let record = execute_find_one(conn, query, &idents, &field_names, ctx) .await? @@ -124,24 +107,13 @@ pub(crate) async fn get_many_records( model: &Model, query_arguments: QueryArguments, selected_fields: &FieldSelection, - aggr_selections: &[RelAggregationSelection], relation_load_strategy: RelationLoadStrategy, ctx: &Context<'_>, ) -> crate::Result { match relation_load_strategy { - RelationLoadStrategy::Join => { - get_many_records_joins(conn, model, query_arguments, selected_fields, aggr_selections, ctx).await - } + RelationLoadStrategy::Join => get_many_records_joins(conn, model, query_arguments, selected_fields, ctx).await, RelationLoadStrategy::Query => { - get_many_records_wo_joins( - conn, - model, - query_arguments, - &ModelProjection::from(selected_fields), - aggr_selections, - ctx, - ) - .await + get_many_records_wo_joins(conn, model, query_arguments, selected_fields, ctx).await } } } @@ -151,7 +123,6 @@ pub(crate) async fn get_many_records_joins( _model: &Model, query_arguments: QueryArguments, selected_fields: &FieldSelection, - _aggr_selections: &[RelAggregationSelection], ctx: &Context<'_>, ) -> crate::Result { let field_names: Vec<_> = selected_fields.db_names().collect(); @@ -197,25 +168,14 @@ pub(crate) async fn get_many_records_wo_joins( conn: &dyn Queryable, model: &Model, mut query_arguments: QueryArguments, - selected_fields: &ModelProjection, - aggr_selections: &[RelAggregationSelection], + selected_fields: &FieldSelection, ctx: &Context<'_>, ) -> crate::Result { + let selected_fields = selected_fields.without_relations().into_virtuals_last(); let reversed = query_arguments.needs_reversed_order(); - let mut field_names: Vec<_> = selected_fields.db_names().collect(); - let mut aggr_field_names: Vec<_> = aggr_selections.iter().map(|aggr_sel| aggr_sel.db_alias()).collect(); - - field_names.append(&mut aggr_field_names); - - let mut aggr_idents = aggr_selections - .iter() - .map(|aggr_sel| aggr_sel.type_identifier_with_arity()) - .collect(); - - let mut idents = selected_fields.type_identifiers_with_arities(); - - idents.append(&mut aggr_idents); + let field_names: Vec<_> = selected_fields.db_names().collect(); + let idents = selected_fields.type_identifiers_with_arities(); let meta = column_metadata::create(field_names.as_slice(), idents.as_slice()); let mut records = ManyRecords::new(field_names.clone()); @@ -252,8 +212,10 @@ pub(crate) async fn get_many_records_wo_joins( for args in batches.into_iter() { let query = read::get_records( model, - selected_fields.as_columns(ctx).mark_all_selected(), - aggr_selections, + ModelProjection::from(&selected_fields) + .as_columns(ctx) + .mark_all_selected(), + selected_fields.virtuals(), args, ctx, ); @@ -274,8 +236,10 @@ pub(crate) async fn get_many_records_wo_joins( _ => { let query = read::get_records( model, - selected_fields.as_columns(ctx).mark_all_selected(), - aggr_selections, + ModelProjection::from(&selected_fields) + .as_columns(ctx) + .mark_all_selected(), + selected_fields.virtuals(), query_arguments, ctx, ); 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 40ca5ce84fc0..54e04651d2f4 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 @@ -25,17 +25,7 @@ pub(crate) async fn update_one_with_selection( // TODO(perf): Technically, if the selectors are fulfilling the field selection, there's no need to perform an additional read. if args.args.is_empty() { let filter = build_update_one_filter(record_filter); - - return get_single_record( - conn, - model, - &filter, - &selected_fields, - &[], - RelationLoadStrategy::Query, - ctx, - ) - .await; + return get_single_record(conn, model, &filter, &selected_fields, RelationLoadStrategy::Query, ctx).await; } let selected_fields = ModelProjection::from(selected_fields); 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 e9afa966b7a5..d5c067851864 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 @@ -13,6 +13,7 @@ use quaint::{ prelude::{native_uuid, uuid_to_bin, uuid_to_bin_swapped, Aliasable, Select, SqlFamily}, }; use query_structure::*; +use std::borrow::Cow; use std::{ collections::{HashMap, HashSet}, ops::Deref, @@ -173,7 +174,7 @@ pub(crate) async fn create_record( // All values provided in the write args (Some(identifier), _, _) if !identifier.misses_autogen_value() => { - let field_names = identifier.db_names().map(ToOwned::to_owned).collect(); + let field_names = identifier.db_names().map(Cow::into_owned).collect(); let record = Record::from(identifier); Ok(SingleRecord { record, field_names }) @@ -183,7 +184,7 @@ pub(crate) async fn create_record( (Some(mut identifier), _, Some(num)) if identifier.misses_autogen_value() => { identifier.add_autogen_value(num as i64); - let field_names = identifier.db_names().map(ToOwned::to_owned).collect(); + let field_names = identifier.db_names().map(Cow::into_owned).collect(); let record = Record::from(identifier); Ok(SingleRecord { record, field_names }) 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 62aae519fc42..c85185c16466 100644 --- a/query-engine/connectors/sql-query-connector/src/database/transaction.rs +++ b/query-engine/connectors/sql-query-connector/src/database/transaction.rs @@ -1,7 +1,7 @@ use super::catch; use crate::{database::operations::*, Context, SqlError}; use async_trait::async_trait; -use connector::{ConnectionLike, RelAggregationSelection}; +use connector::ConnectionLike; use connector_interface::{ self as connector, AggregationRow, AggregationSelection, ReadOperations, RecordFilter, Transaction, WriteArgs, WriteOperations, @@ -68,7 +68,6 @@ impl<'tx> ReadOperations for SqlConnectorTransaction<'tx> { model: &Model, filter: &Filter, selected_fields: &FieldSelection, - aggr_selections: &[RelAggregationSelection], relation_load_strategy: RelationLoadStrategy, trace_id: Option, ) -> connector::Result> { @@ -80,7 +79,6 @@ impl<'tx> ReadOperations for SqlConnectorTransaction<'tx> { model, filter, selected_fields, - aggr_selections, relation_load_strategy, &ctx, ), @@ -93,7 +91,6 @@ impl<'tx> ReadOperations for SqlConnectorTransaction<'tx> { model: &Model, query_arguments: QueryArguments, selected_fields: &FieldSelection, - aggr_selections: &[RelAggregationSelection], relation_load_strategy: RelationLoadStrategy, trace_id: Option, ) -> connector::Result { @@ -105,7 +102,6 @@ impl<'tx> ReadOperations for SqlConnectorTransaction<'tx> { model, query_arguments, selected_fields, - aggr_selections, relation_load_strategy, &ctx, ), diff --git a/query-engine/connectors/sql-query-connector/src/model_extensions/selection_result.rs b/query-engine/connectors/sql-query-connector/src/model_extensions/selection_result.rs index 9cf94e06d423..21d6aac3dbe2 100644 --- a/query-engine/connectors/sql-query-connector/src/model_extensions/selection_result.rs +++ b/query-engine/connectors/sql-query-connector/src/model_extensions/selection_result.rs @@ -38,6 +38,7 @@ impl SelectionResultExt for SelectionResult { SelectedField::Scalar(sf) => Some(sf.value(v.clone(), ctx)), SelectedField::Composite(_) => None, SelectedField::Relation(_) => None, + SelectedField::Virtual(_) => None, }) .collect() } diff --git a/query-engine/connectors/sql-query-connector/src/nested_aggregations.rs b/query-engine/connectors/sql-query-connector/src/nested_aggregations.rs index 91236e77024a..9a8312153e1c 100644 --- a/query-engine/connectors/sql-query-connector/src/nested_aggregations.rs +++ b/query-engine/connectors/sql-query-connector/src/nested_aggregations.rs @@ -2,8 +2,8 @@ use crate::{ join_utils::{compute_aggr_join, AggregationType, AliasedJoin}, Context, }; -use connector_interface::RelAggregationSelection; use quaint::prelude::*; +use query_structure::VirtualSelection; #[derive(Debug)] pub(crate) struct RelAggregationJoins { @@ -13,13 +13,16 @@ pub(crate) struct RelAggregationJoins { pub(crate) columns: Vec>, } -pub(crate) fn build(aggr_selections: &[RelAggregationSelection], ctx: &Context<'_>) -> RelAggregationJoins { +pub(crate) fn build<'a>( + virtual_selections: impl IntoIterator, + ctx: &Context<'_>, +) -> RelAggregationJoins { let mut joins = vec![]; let mut columns: Vec> = vec![]; - for (index, selection) in aggr_selections.iter().enumerate() { + for (index, selection) in virtual_selections.into_iter().enumerate() { match selection { - RelAggregationSelection::Count(rf, filter) => { + VirtualSelection::RelationCount(rf, filter) => { let join_alias = format!("aggr_selection_{index}"); let aggregator_alias = selection.db_alias(); let join = compute_aggr_join( diff --git a/query-engine/connectors/sql-query-connector/src/query_builder/read.rs b/query-engine/connectors/sql-query-connector/src/query_builder/read.rs index b3fbb548b64c..6a0572ecc0da 100644 --- a/query-engine/connectors/sql-query-connector/src/query_builder/read.rs +++ b/query-engine/connectors/sql-query-connector/src/query_builder/read.rs @@ -2,49 +2,49 @@ use crate::{ cursor_condition, filter::FilterBuilder, model_extensions::*, nested_aggregations, ordering::OrderByBuilder, sql_trace::SqlTraceComment, Context, }; -use connector_interface::{AggregationSelection, RelAggregationSelection}; +use connector_interface::AggregationSelection; use itertools::Itertools; use quaint::ast::*; use query_structure::*; use tracing::Span; pub(crate) trait SelectDefinition { - fn into_select( + fn into_select<'a>( self, _: &Model, - aggr_selections: &[RelAggregationSelection], + virtual_selections: impl IntoIterator, ctx: &Context<'_>, ) -> (Select<'static>, Vec>); } impl SelectDefinition for Filter { - fn into_select( + fn into_select<'a>( self, model: &Model, - aggr_selections: &[RelAggregationSelection], + virtual_selections: impl IntoIterator, ctx: &Context<'_>, ) -> (Select<'static>, Vec>) { let args = QueryArguments::from((model.clone(), self)); - args.into_select(model, aggr_selections, ctx) + args.into_select(model, virtual_selections, ctx) } } impl SelectDefinition for &Filter { - fn into_select( + fn into_select<'a>( self, model: &Model, - aggr_selections: &[RelAggregationSelection], + virtual_selections: impl IntoIterator, ctx: &Context<'_>, ) -> (Select<'static>, Vec>) { - self.clone().into_select(model, aggr_selections, ctx) + self.clone().into_select(model, virtual_selections, ctx) } } impl SelectDefinition for Select<'static> { - fn into_select( + fn into_select<'a>( self, _: &Model, - _: &[RelAggregationSelection], + _: impl IntoIterator, _ctx: &Context<'_>, ) -> (Select<'static>, Vec>) { (self, vec![]) @@ -52,15 +52,15 @@ impl SelectDefinition for Select<'static> { } impl SelectDefinition for QueryArguments { - fn into_select( + fn into_select<'a>( self, model: &Model, - aggr_selections: &[RelAggregationSelection], + virtual_selections: impl IntoIterator, ctx: &Context<'_>, ) -> (Select<'static>, Vec>) { let order_by_definitions = OrderByBuilder::default().build(&self, ctx); let cursor_condition = cursor_condition::build(&self, model, &order_by_definitions, ctx); - let aggregation_joins = nested_aggregations::build(aggr_selections, ctx); + let aggregation_joins = nested_aggregations::build(virtual_selections, ctx); let limit = if self.ignore_take { None } else { self.take_abs() }; let skip = if self.ignore_skip { 0 } else { self.skip.unwrap_or(0) }; @@ -124,17 +124,17 @@ impl SelectDefinition for QueryArguments { } } -pub(crate) fn get_records( +pub(crate) fn get_records<'a, T>( model: &Model, columns: impl Iterator>, - aggr_selections: &[RelAggregationSelection], + virtual_selections: impl IntoIterator, query: T, ctx: &Context<'_>, ) -> Select<'static> where T: SelectDefinition, { - let (select, additional_selection_set) = query.into_select(model, aggr_selections, ctx); + let (select, additional_selection_set) = query.into_select(model, virtual_selections, ctx); let select = columns.fold(select, |acc, col| acc.column(col)); let select = select.append_trace(&Span::current()).add_trace_id(ctx.trace_id); diff --git a/query-engine/core/src/interpreter/interpreter_impl.rs b/query-engine/core/src/interpreter/interpreter_impl.rs index f1011b13f8f5..cc10f0749063 100644 --- a/query-engine/core/src/interpreter/interpreter_impl.rs +++ b/query-engine/core/src/interpreter/interpreter_impl.rs @@ -66,7 +66,7 @@ impl ExpressionResult { // We always select IDs, the unwraps are safe. QueryResult::RecordSelection(Some(rs)) => Some( - rs.scalars + rs.records .extract_selection_results(field_selection) .expect("Expected record selection to contain required model ID fields.") .into_iter() diff --git a/query-engine/core/src/interpreter/query_interpreters/nested_read.rs b/query-engine/core/src/interpreter/query_interpreters/nested_read.rs index fd2e28ad5558..1238a35f1a24 100644 --- a/query-engine/core/src/interpreter/query_interpreters/nested_read.rs +++ b/query-engine/core/src/interpreter/query_interpreters/nested_read.rs @@ -1,6 +1,6 @@ -use super::{inmemory_record_processor::InMemoryRecordProcessor, read}; +use super::inmemory_record_processor::InMemoryRecordProcessor; use crate::{interpreter::InterpretationResult, query_ast::*}; -use connector::{self, ConnectionLike, RelAggregationRow, RelAggregationSelection}; +use connector::ConnectionLike; use query_structure::*; use std::collections::HashMap; @@ -9,7 +9,7 @@ pub(crate) async fn m2m( query: &mut RelatedRecordsQuery, parent_result: Option<&ManyRecords>, trace_id: Option, -) -> InterpretationResult<(ManyRecords, Option>)> { +) -> InterpretationResult { let processor = InMemoryRecordProcessor::new_from_query_args(&mut query.args); let parent_field = &query.parent_field; @@ -27,14 +27,14 @@ pub(crate) async fn m2m( }; if parent_ids.is_empty() { - return Ok((ManyRecords::empty(&query.selected_fields), None)); + return Ok(ManyRecords::empty(&query.selected_fields)); } let ids = tx .get_related_m2m_record_ids(&query.parent_field, &parent_ids, trace_id.clone()) .await?; if ids.is_empty() { - return Ok((ManyRecords::empty(&query.selected_fields), None)); + return Ok(ManyRecords::empty(&query.selected_fields)); } let child_model_id = query.parent_field.related_model().primary_identifier(); @@ -49,32 +49,32 @@ pub(crate) async fn m2m( // a roundtrip can be avoided if: // - there is no additional filter - // - there is no aggregation selection + // - there is no virtual fields selection (relation aggregation) // - the selection set is the child_link_id - let mut scalars = - if query.args.do_nothing() && query.aggregation_selections.is_empty() && child_link_id == query.selected_fields - { - ManyRecords::from((child_ids, &query.selected_fields)).with_unique_records() - } else { - let mut args = query.args.clone(); - let filter = child_link_id.is_in(ConditionListValue::list(child_ids)); - - args.filter = match args.filter { - Some(existing_filter) => Some(Filter::and(vec![existing_filter, filter])), - None => Some(filter), - }; - - tx.get_many_records( - &query.parent_field.related_model(), - args, - &query.selected_fields, - &query.aggregation_selections, - RelationLoadStrategy::Query, - trace_id.clone(), - ) - .await? + let mut scalars = if query.args.do_nothing() + && !query.selected_fields.has_virtual_fields() + && child_link_id == query.selected_fields + { + ManyRecords::from((child_ids, &query.selected_fields)).with_unique_records() + } else { + let mut args = query.args.clone(); + let filter = child_link_id.is_in(ConditionListValue::list(child_ids)); + + args.filter = match args.filter { + Some(existing_filter) => Some(Filter::and(vec![existing_filter, filter])), + None => Some(filter), }; + tx.get_many_records( + &query.parent_field.related_model(), + args, + &query.selected_fields, + RelationLoadStrategy::Query, + trace_id.clone(), + ) + .await? + }; + // Child id to parent ids let mut id_map: HashMap> = HashMap::new(); @@ -124,10 +124,8 @@ pub(crate) async fn m2m( } let scalars = processor.apply(scalars); - let (scalars, aggregation_rows) = - read::extract_aggregation_rows_from_scalars(scalars, query.aggregation_selections.clone()); - Ok((scalars, aggregation_rows)) + Ok(scalars) } // [DTODO] This is implemented in an inefficient fashion, e.g. too much Arc cloning going on. @@ -139,9 +137,8 @@ pub async fn one2m( parent_result: Option<&ManyRecords>, mut query_args: QueryArguments, selected_fields: &FieldSelection, - aggr_selections: Vec, trace_id: Option, -) -> InterpretationResult<(ManyRecords, Option>)> { +) -> InterpretationResult { let parent_model_id = parent_field.model().primary_identifier(); let parent_link_id = parent_field.linking_fields(); let child_link_id = parent_field.related_field().linking_fields(); @@ -185,7 +182,7 @@ pub async fn one2m( .collect(); if uniq_selections.is_empty() { - return Ok((ManyRecords::empty(selected_fields), None)); + return Ok(ManyRecords::empty(selected_fields)); } // If we're fetching related records from a single parent, then we can apply normal pagination instead of in-memory processing. @@ -210,7 +207,6 @@ pub async fn one2m( &parent_field.related_model(), args, selected_fields, - &aggr_selections, RelationLoadStrategy::Query, trace_id, ) @@ -269,7 +265,6 @@ pub async fn one2m( } else { scalars }; - let (scalars, aggregation_rows) = read::extract_aggregation_rows_from_scalars(scalars, aggr_selections); - Ok((scalars, aggregation_rows)) + Ok(scalars) } diff --git a/query-engine/core/src/interpreter/query_interpreters/read.rs b/query-engine/core/src/interpreter/query_interpreters/read.rs index 53a4f6dbe18a..883246dc84e2 100644 --- a/query-engine/core/src/interpreter/query_interpreters/read.rs +++ b/query-engine/core/src/interpreter/query_interpreters/read.rs @@ -1,9 +1,8 @@ use super::{inmemory_record_processor::InMemoryRecordProcessor, *}; use crate::{interpreter::InterpretationResult, query_ast::*, result_ast::*}; -use connector::{self, error::ConnectorError, ConnectionLike, RelAggregationRow, RelAggregationSelection}; +use connector::{error::ConnectorError, ConnectionLike}; use futures::future::{BoxFuture, FutureExt}; use query_structure::{ManyRecords, RelationLoadStrategy, RelationSelection}; -use std::collections::HashMap; use user_facing_errors::KnownError; pub(crate) fn execute<'conn>( @@ -33,31 +32,28 @@ fn read_one( let fut = async move { let model = query.model; let filter = query.filter.expect("Expected filter to be set for ReadOne query."); - let scalars = tx + let record = tx .get_single_record( &model, &filter, &query.selected_fields, - &query.aggregation_selections, query.relation_load_strategy, trace_id, ) .await?; - match scalars { + match record { Some(record) if query.relation_load_strategy.is_query() => { - let scalars: ManyRecords = record.into(); - let (scalars, aggregation_rows) = - extract_aggregation_rows_from_scalars(scalars, query.aggregation_selections); - let nested: Vec = process_nested(tx, query.nested, Some(&scalars)).await?; + let records = record.into(); + let nested = process_nested(tx, query.nested, Some(&records)).await?; Ok(RecordSelection { name: query.name, fields: query.selection_order, - scalars, + records, nested, model, - aggregation_rows, + virtual_fields: query.selected_fields.virtuals_owned(), } .into()) } @@ -79,10 +75,10 @@ fn read_one( None => Ok(QueryResult::RecordSelection(Some(Box::new(RecordSelection { name: query.name, fields: query.selection_order, - scalars: ManyRecords::default(), + records: ManyRecords::default(), nested: vec![], model, - aggregation_rows: None, + virtual_fields: query.selected_fields.virtuals_owned(), })))), } }; @@ -119,36 +115,33 @@ fn read_many_by_queries( }; let fut = async move { - let scalars = tx + let records = tx .get_many_records( &query.model, query.args.clone(), &query.selected_fields, - &query.aggregation_selections, query.relation_load_strategy, trace_id, ) .await?; - let scalars = if let Some(p) = processor { - p.apply(scalars) + let records = if let Some(p) = processor { + p.apply(records) } else { - scalars + records }; - let (scalars, aggregation_rows) = extract_aggregation_rows_from_scalars(scalars, query.aggregation_selections); - - if scalars.records.is_empty() && query.options.contains(QueryOption::ThrowOnEmpty) { + if records.records.is_empty() && query.options.contains(QueryOption::ThrowOnEmpty) { record_not_found() } else { - let nested: Vec = process_nested(tx, query.nested, Some(&scalars)).await?; + let nested: Vec = process_nested(tx, query.nested, Some(&records)).await?; Ok(RecordSelection { name: query.name, fields: query.selection_order, - scalars, + records, nested, model: query.model, - aggregation_rows, + virtual_fields: query.selected_fields.virtuals_owned(), } .into()) } @@ -168,7 +161,6 @@ fn read_many_by_joins( &query.model, query.args.clone(), &query.selected_fields, - &query.aggregation_selections, query.relation_load_strategy, trace_id, ) @@ -214,7 +206,7 @@ fn read_related<'conn>( let fut = async move { let relation = query.parent_field.relation(); - let (scalars, aggregation_rows) = if relation.is_many_to_many() { + let records = if relation.is_many_to_many() { nested_read::m2m(tx, &mut query, parent_result, trace_id).await? } else { nested_read::one2m( @@ -224,21 +216,20 @@ fn read_related<'conn>( parent_result, query.args.clone(), &query.selected_fields, - query.aggregation_selections, trace_id, ) .await? }; let model = query.parent_field.related_model(); - let nested: Vec = process_nested(tx, query.nested, Some(&scalars)).await?; + let nested: Vec = process_nested(tx, query.nested, Some(&records)).await?; Ok(RecordSelection { name: query.name, fields: query.selection_order, - scalars, + records, nested, model, - aggregation_rows, + virtual_fields: query.selected_fields.virtuals_owned(), } .into()) }; @@ -297,58 +288,6 @@ fn process_nested<'conn>( fut.boxed() } -/// Removes the relation aggregation data from the database result and collect it into some RelAggregationRow -/// Explanation: Relation aggregations on a findMany are selected from an output object type. eg: -/// 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 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(crate) fn extract_aggregation_rows_from_scalars( - mut scalars: ManyRecords, - aggr_selections: Vec, -) -> (ManyRecords, Option>) { - if aggr_selections.is_empty() { - return (scalars, None); - } - - let aggr_field_names: HashMap = aggr_selections - .iter() - .map(|aggr_sel| (aggr_sel.db_alias(), aggr_sel)) - .collect(); - - let indexes_to_remove: Vec<_> = scalars - .field_names - .iter() - .enumerate() - .filter_map(|(i, field_name)| aggr_field_names.get(field_name).map(|aggr_sel| (i, *aggr_sel))) - .collect(); - - let mut aggregation_rows: Vec = vec![]; - - for (n_record_removed, (index_to_remove, aggr_sel)) in indexes_to_remove.into_iter().enumerate() { - let index_to_remove = index_to_remove - n_record_removed; - - // Remove all aggr field names - scalars.field_names.remove(index_to_remove); - - // Remove and collect all aggr prisma values - for (r_index, record) in scalars.records.iter_mut().enumerate() { - let val = record.values.remove(index_to_remove); - let aggr_result = aggr_sel.clone().into_result(val); - - // Group the aggregation results by record - match aggregation_rows.get_mut(r_index) { - Some(inner_vec) => inner_vec.push(aggr_result), - None => aggregation_rows.push(vec![aggr_result]), - } - } - } - - (scalars, Some(aggregation_rows)) -} - // Custom error built for findXOrThrow queries, when a record is not found and it needs to throw an error #[inline] fn record_not_found() -> InterpretationResult { diff --git a/query-engine/core/src/interpreter/query_interpreters/write.rs b/query-engine/core/src/interpreter/query_interpreters/write.rs index d143e4a9f63f..39203763ed4d 100644 --- a/query-engine/core/src/interpreter/query_interpreters/write.rs +++ b/query-engine/core/src/interpreter/query_interpreters/write.rs @@ -48,10 +48,10 @@ async fn create_one( Ok(QueryResult::RecordSelection(Some(Box::new(RecordSelection { name: q.name, fields: q.selection_order, - aggregation_rows: None, model: q.model, - scalars: res.into(), + records: res.into(), nested: vec![], + virtual_fields: vec![], })))) } @@ -86,10 +86,10 @@ async fn update_one( .map(|res| RecordSelection { name: q.name, fields: q.selection_order, - scalars: res.into(), + records: res.into(), nested: vec![], model: q.model, - aggregation_rows: None, + virtual_fields: vec![], }) .map(Box::new); @@ -115,10 +115,10 @@ async fn native_upsert( Ok(RecordSelection { name: query.name().to_string(), fields: query.selection_order().to_owned(), - scalars: scalars.into(), + records: scalars.into(), nested: Vec::new(), model: query.model().clone(), - aggregation_rows: None, + virtual_fields: vec![], } .into()) } @@ -144,10 +144,10 @@ async fn delete_one( let selection = RecordSelection { name: q.name, fields: selected_fields.order, - scalars: record.into(), + records: record.into(), nested: vec![], model: q.model, - aggregation_rows: None, + virtual_fields: vec![], }; Ok(QueryResult::RecordSelection(Some(Box::new(selection)))) diff --git a/query-engine/core/src/query_ast/read.rs b/query-engine/core/src/query_ast/read.rs index 1edd9a074f8d..64a2440c0f52 100644 --- a/query-engine/core/src/query_ast/read.rs +++ b/query-engine/core/src/query_ast/read.rs @@ -1,10 +1,10 @@ //! Prisma read query AST use super::FilteredQuery; use crate::ToGraphviz; -use connector::{AggregationSelection, RelAggregationSelection}; +use connector::AggregationSelection; use enumflags2::BitFlags; use query_structure::{prelude::*, Filter, QueryArguments, RelationLoadStrategy}; -use std::{fmt::Display, mem}; +use std::fmt::Display; #[allow(clippy::enum_variant_names)] #[derive(Debug, Clone)] @@ -35,13 +35,13 @@ impl ReadQuery { pub fn satisfy_dependency(&mut self, field_selection: FieldSelection) { match self { ReadQuery::RecordQuery(x) => { - x.selected_fields = mem::take(&mut x.selected_fields).merge(field_selection); + x.selected_fields.merge_in_place(field_selection); } ReadQuery::ManyRecordsQuery(x) => { - x.selected_fields = mem::take(&mut x.selected_fields).merge(field_selection); + x.selected_fields.merge_in_place(field_selection); } ReadQuery::RelatedRecordsQuery(x) => { - x.selected_fields = mem::take(&mut x.selected_fields).merge(field_selection); + x.selected_fields.merge_in_place(field_selection); } ReadQuery::AggregateRecordsQuery(_) => (), } @@ -74,15 +74,15 @@ impl ReadQuery { } } - pub(crate) fn has_aggregation_selections(&self) -> bool { - fn has_aggregations(selections: &[RelAggregationSelection], nested: &[ReadQuery]) -> bool { - !selections.is_empty() || nested.iter().any(|q| q.has_aggregation_selections()) + pub(crate) fn has_virtual_selections(&self) -> bool { + fn has_virtuals(selection: &FieldSelection, nested: &[ReadQuery]) -> bool { + selection.has_virtual_fields() || nested.iter().any(|q| q.has_virtual_selections()) } match self { - ReadQuery::RecordQuery(q) => has_aggregations(&q.aggregation_selections, &q.nested), - ReadQuery::ManyRecordsQuery(q) => has_aggregations(&q.aggregation_selections, &q.nested), - ReadQuery::RelatedRecordsQuery(q) => has_aggregations(&q.aggregation_selections, &q.nested), + ReadQuery::RecordQuery(q) => has_virtuals(&q.selected_fields, &q.nested), + ReadQuery::ManyRecordsQuery(q) => has_virtuals(&q.selected_fields, &q.nested), + ReadQuery::RelatedRecordsQuery(q) => has_virtuals(&q.selected_fields, &q.nested), ReadQuery::AggregateRecordsQuery(_) => false, } } @@ -198,10 +198,10 @@ pub struct RecordQuery { pub alias: Option, pub model: Model, pub filter: Option, + // TODO: split into `user_selection` and `full_selection` and get rid of `selection_order` pub selected_fields: FieldSelection, pub(crate) nested: Vec, pub selection_order: Vec, - pub aggregation_selections: Vec, pub options: QueryOptions, pub relation_load_strategy: RelationLoadStrategy, } @@ -212,10 +212,10 @@ pub struct ManyRecordsQuery { pub alias: Option, pub model: Model, pub args: QueryArguments, + // TODO: split into `user_selection` and `full_selection` and get rid of `selection_order` pub selected_fields: FieldSelection, pub(crate) nested: Vec, pub selection_order: Vec, - pub aggregation_selections: Vec, pub options: QueryOptions, pub relation_load_strategy: RelationLoadStrategy, } @@ -226,11 +226,10 @@ pub struct RelatedRecordsQuery { pub alias: Option, pub parent_field: RelationFieldRef, pub args: QueryArguments, + // TODO: split into `user_selection` and `full_selection` and get rid of `selection_order` pub selected_fields: FieldSelection, pub nested: Vec, pub selection_order: Vec, - pub aggregation_selections: Vec, - /// Fields and values of the parent to satisfy the relation query without /// relying on the parent result passed by the interpreter. pub parent_results: Option>, @@ -245,8 +244,8 @@ impl RelatedRecordsQuery { self.args.distinct.is_some() || self.nested.iter().any(|q| q.has_distinct()) } - pub fn has_aggregation_selections(&self) -> bool { - !self.aggregation_selections.is_empty() || self.nested.iter().any(|q| q.has_aggregation_selections()) + pub fn has_virtual_selections(&self) -> bool { + self.selected_fields.has_virtual_fields() || self.nested.iter().any(|q| q.has_virtual_selections()) } } diff --git a/query-engine/core/src/query_ast/write.rs b/query-engine/core/src/query_ast/write.rs index fac1692fa5f8..b3c09aabd200 100644 --- a/query-engine/core/src/query_ast/write.rs +++ b/query-engine/core/src/query_ast/write.rs @@ -37,7 +37,7 @@ impl WriteQuery { for (selected_field, value) in result { args.insert( - DatasourceFieldName(selected_field.db_name().to_owned()), + DatasourceFieldName(selected_field.db_name().into_owned()), (&selected_field, value), ) } @@ -257,7 +257,7 @@ impl CreateManyRecords { for (selected_field, value) in result { for args in self.args.iter_mut() { args.insert( - DatasourceFieldName(selected_field.db_name().to_owned()), + DatasourceFieldName(selected_field.db_name().into_owned()), (&selected_field, value.clone()), ) } diff --git a/query-engine/core/src/query_graph/mod.rs b/query-engine/core/src/query_graph/mod.rs index 6a99f1462ccf..458be8280a3a 100644 --- a/query-engine/core/src/query_graph/mod.rs +++ b/query-engine/core/src/query_graph/mod.rs @@ -795,7 +795,6 @@ impl QueryGraph { selected_fields: identifiers.merge(primary_model_id.clone()), nested: vec![], selection_order: vec![], - aggregation_selections: vec![], options: QueryOptions::none(), relation_load_strategy: query_structure::RelationLoadStrategy::Query, }); diff --git a/query-engine/core/src/query_graph_builder/extractors/rel_aggregations.rs b/query-engine/core/src/query_graph_builder/extractors/rel_aggregations.rs index 953048cbb9a3..6a6ab715b069 100644 --- a/query-engine/core/src/query_graph_builder/extractors/rel_aggregations.rs +++ b/query-engine/core/src/query_graph_builder/extractors/rel_aggregations.rs @@ -1,12 +1,6 @@ use super::*; use schema::constants::aggregations::*; -pub(crate) fn extract_nested_rel_aggr_selections( - field_pairs: Vec>, -) -> (Vec>, Vec>) { - field_pairs.into_iter().partition(is_aggr_selection) -} - pub(crate) fn is_aggr_selection(pair: &FieldPair<'_>) -> bool { matches!(pair.parsed_field.name.as_str(), UNDERSCORE_COUNT) } diff --git a/query-engine/core/src/query_graph_builder/read/many.rs b/query-engine/core/src/query_graph_builder/read/many.rs index 3a462588f957..29eb769f74d0 100644 --- a/query-engine/core/src/query_graph_builder/read/many.rs +++ b/query-engine/core/src/query_graph_builder/read/many.rs @@ -30,12 +30,9 @@ fn find_many_with_options( let name = field.name; let alias = field.alias; let nested_fields = field.nested_fields.unwrap().fields; - let (aggr_fields_pairs, nested_fields) = extractors::extract_nested_rel_aggr_selections(nested_fields); - let aggregation_selections = utils::collect_relation_aggr_selections(aggr_fields_pairs, &model)?; let selection_order: Vec = utils::collect_selection_order(&nested_fields); let selected_fields = utils::collect_selected_fields(&nested_fields, args.distinct.clone(), &model, query_schema)?; let nested = utils::collect_nested_queries(nested_fields, &model, query_schema)?; - let model = model; let selected_fields = utils::merge_relation_selections(selected_fields, None, &nested); let selected_fields = utils::merge_cursor_fields(selected_fields, &args.cursor); @@ -45,7 +42,7 @@ fn find_many_with_options( args.cursor.as_ref(), args.distinct.as_ref(), &nested, - &aggregation_selections, + &selected_fields, query_schema, ); @@ -57,7 +54,6 @@ fn find_many_with_options( selected_fields, nested, selection_order, - aggregation_selections, options, relation_load_strategy, })) diff --git a/query-engine/core/src/query_graph_builder/read/one.rs b/query-engine/core/src/query_graph_builder/read/one.rs index 97b2e13e3f93..afc07ed0e89e 100644 --- a/query-engine/core/src/query_graph_builder/read/one.rs +++ b/query-engine/core/src/query_graph_builder/read/one.rs @@ -44,11 +44,8 @@ fn find_unique_with_options( let name = field.name; let alias = field.alias; - let model = model; let nested_fields = field.nested_fields.unwrap().fields; - let (aggr_fields_pairs, nested_fields) = extractors::extract_nested_rel_aggr_selections(nested_fields); - let aggregation_selections = utils::collect_relation_aggr_selections(aggr_fields_pairs, &model)?; - let selection_order: Vec = utils::collect_selection_order(&nested_fields); + let selection_order = utils::collect_selection_order(&nested_fields); let selected_fields = utils::collect_selected_fields(&nested_fields, None, &model, query_schema)?; let nested = utils::collect_nested_queries(nested_fields, &model, query_schema)?; let selected_fields = utils::merge_relation_selections(selected_fields, None, &nested); @@ -58,7 +55,7 @@ fn find_unique_with_options( None, None, &nested, - &aggregation_selections, + &selected_fields, query_schema, ); @@ -70,7 +67,6 @@ fn find_unique_with_options( selected_fields, nested, selection_order, - aggregation_selections, options, relation_load_strategy, })) diff --git a/query-engine/core/src/query_graph_builder/read/related.rs b/query-engine/core/src/query_graph_builder/read/related.rs index 7ebed8a7a06c..84eb13bac551 100644 --- a/query-engine/core/src/query_graph_builder/read/related.rs +++ b/query-engine/core/src/query_graph_builder/read/related.rs @@ -13,8 +13,6 @@ pub(crate) fn find_related( let name = field.name; let alias = field.alias; let sub_selections = field.nested_fields.unwrap().fields; - let (aggr_fields_pairs, sub_selections) = extractors::extract_nested_rel_aggr_selections(sub_selections); - let aggregation_selections = utils::collect_relation_aggr_selections(aggr_fields_pairs, &model)?; let selection_order: Vec = utils::collect_selection_order(&sub_selections); let selected_fields = utils::collect_selected_fields(&sub_selections, args.distinct.clone(), &model, query_schema)?; let nested = utils::collect_nested_queries(sub_selections, &model, query_schema)?; @@ -31,7 +29,6 @@ pub(crate) fn find_related( selected_fields, nested, selection_order, - aggregation_selections, parent_results: None, })) } diff --git a/query-engine/core/src/query_graph_builder/read/utils.rs b/query-engine/core/src/query_graph_builder/read/utils.rs index c63b299fcc27..69fe60d95f39 100644 --- a/query-engine/core/src/query_graph_builder/read/utils.rs +++ b/query-engine/core/src/query_graph_builder/read/utils.rs @@ -1,6 +1,5 @@ use super::*; use crate::{ArgumentListLookup, FieldPair, ParsedField, ReadQuery}; -use connector::RelAggregationSelection; use psl::{datamodel_connector::ConnectorCapability, PreviewFeature}; use query_structure::{prelude::*, RelationLoadStrategy}; use schema::{ @@ -78,32 +77,45 @@ where let parent = parent.into(); - let selected_fields = pairs - .iter() - .filter_map(|pair| { - parent - .find_field(&pair.parsed_field.name) - .map(|field| (pair.parsed_field.clone(), field)) - }) - .flat_map(|field| match field { - (pf, Field::Relation(rf)) => { - let mut fields: Vec> = rf - .scalar_fields() - .into_iter() - .map(SelectedField::from) - .map(Ok) - .collect(); + let mut selected_fields = Vec::new(); + + for pair in pairs { + let field = parent.find_field(&pair.parsed_field.name); + + match (pair.parsed_field.clone(), field) { + (pf, Some(Field::Relation(rf))) => { + let fields = rf.scalar_fields().into_iter().map(SelectedField::from); + + selected_fields.extend(fields); if should_collect_relation_selection { - fields.push(extract_relation_selection(pf, rf, query_schema)); + selected_fields.push(extract_relation_selection(pf, rf, query_schema)?); } + } - fields + (_, Some(Field::Scalar(sf))) => { + selected_fields.push(sf.into()); } - (_, Field::Scalar(sf)) => vec![Ok(sf.into())], - (pf, Field::Composite(cf)) => vec![extract_composite_selection(pf, cf, query_schema)], - }) - .collect::, _>>()?; + + (pf, Some(Field::Composite(cf))) => { + selected_fields.push(extract_composite_selection(pf, cf, query_schema)?); + } + + (pf, None) if pf.name == UNDERSCORE_COUNT => match parent { + ParentContainer::Model(ref model) => { + selected_fields.extend(extract_relation_count_selections(pf, model)?); + } + ParentContainer::CompositeType(_) => { + unreachable!("Unexpected relation aggregation selection inside a composite type query") + } + }, + + (pf, None) => unreachable!( + "Field '{}' does not exist on enclosing type and is not a known virtual field", + pf.name + ), + } + } Ok(selected_fields) } @@ -144,6 +156,35 @@ fn extract_relation_selection( })) } +fn extract_relation_count_selections( + pf: ParsedField<'_>, + model: &Model, +) -> QueryGraphBuilderResult> { + let object = pf + .nested_fields + .expect("Invalid query shape: relation aggregation virtual field selected without relations to aggregate."); + + object + .fields + .into_iter() + .map(|mut nested_pair| -> QueryGraphBuilderResult<_> { + let rf = model + .fields() + .find_from_relation_fields(&nested_pair.parsed_field.name) + .expect("Selected relation in relation aggregation virtual field must exist on the model"); + + let filter = nested_pair + .parsed_field + .arguments + .lookup(args::WHERE) + .map(|where_arg| extract_filter(where_arg.value.try_into()?, rf.related_model())) + .transpose()?; + + Ok(SelectedField::Virtual(VirtualSelection::RelationCount(rf, filter))) + }) + .collect() +} + pub(crate) fn collect_nested_queries( from: Vec>, model: &Model, @@ -211,52 +252,21 @@ pub fn merge_cursor_fields(selected_fields: FieldSelection, cursor: &Option>, - model: &Model, -) -> QueryGraphBuilderResult> { - let mut selections = vec![]; - - for pair in from { - match pair.parsed_field.name.as_str() { - UNDERSCORE_COUNT => { - let nested_fields = pair.parsed_field.nested_fields.unwrap(); - - for mut nested_pair in nested_fields.fields { - let rf = model - .fields() - .find_from_relation_fields(&nested_pair.parsed_field.name) - .unwrap(); - let filter = match nested_pair.parsed_field.arguments.lookup(args::WHERE) { - Some(where_arg) => Some(extract_filter(where_arg.value.try_into()?, rf.related_model())?), - _ => None, - }; - - selections.push(RelAggregationSelection::Count(rf, filter)); - } - } - field_name => panic!("Unknown field name \"{field_name}\" for a relation aggregation"), - } - } - - Ok(selections) -} - pub(crate) fn get_relation_load_strategy( requested_strategy: Option, cursor: Option<&SelectionResult>, distinct: Option<&FieldSelection>, nested_queries: &[ReadQuery], - aggregation_selections: &[RelAggregationSelection], + selected_fields: &FieldSelection, query_schema: &QuerySchema, ) -> RelationLoadStrategy { if query_schema.has_feature(PreviewFeature::RelationJoins) && query_schema.has_capability(ConnectorCapability::LateralJoin) && cursor.is_none() && distinct.is_none() - && aggregation_selections.is_empty() + && !selected_fields.has_virtual_fields() && !nested_queries.iter().any(|q| match q { - ReadQuery::RelatedRecordsQuery(q) => q.has_cursor() || q.has_distinct() || q.has_aggregation_selections(), + ReadQuery::RelatedRecordsQuery(q) => q.has_cursor() || q.has_distinct() || q.has_virtual_selections(), _ => false, }) && requested_strategy != Some(RelationLoadStrategy::Query) 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 7e3094da8000..a931fa1edc30 100644 --- a/query-engine/core/src/query_graph_builder/write/utils.rs +++ b/query-engine/core/src/query_graph_builder/write/utils.rs @@ -42,7 +42,6 @@ where selected_fields, nested: vec![], selection_order: vec![], - aggregation_selections: vec![], options: QueryOptions::none(), relation_load_strategy: query_structure::RelationLoadStrategy::Query, }); @@ -113,7 +112,6 @@ where parent_results: None, args: (child_model, filter).into(), selected_fields, - aggregation_selections: vec![], nested: vec![], selection_order: vec![], }))); @@ -839,7 +837,7 @@ pub fn emulate_on_update_restrict( let linking_fields_updated = linking_fields .into_iter() - .any(|parent_pk| parent_update_args.get_field_value(parent_pk.db_name()).is_some()); + .any(|parent_pk| parent_update_args.get_field_value(&parent_pk.db_name()).is_some()); graph.create_edge( &read_node, diff --git a/query-engine/core/src/response_ir/internal.rs b/query-engine/core/src/response_ir/internal.rs index a75d69f34573..121525dad15c 100644 --- a/query-engine/core/src/response_ir/internal.rs +++ b/query-engine/core/src/response_ir/internal.rs @@ -5,10 +5,9 @@ use crate::{ result_ast::{RecordSelectionWithRelations, RelationRecordSelection}, CoreError, QueryResult, RecordAggregations, RecordSelection, }; -use connector::{AggregationResult, RelAggregationResult, RelAggregationRow}; +use connector::AggregationResult; use indexmap::IndexMap; -use itertools::Itertools; -use query_structure::{CompositeFieldRef, Field, PrismaValue, SelectionResult}; +use query_structure::{CompositeFieldRef, Field, PrismaValue, SelectionResult, VirtualSelection}; use schema::{ constants::{aggregations::*, output_fields::*}, *, @@ -175,24 +174,6 @@ fn serialize_aggregations( Ok(envelope) } -fn write_rel_aggregation_row(row: &RelAggregationRow, map: &mut HashMap) { - for result in row.iter() { - match result { - RelAggregationResult::Count(rf, count) => match map.get_mut(UNDERSCORE_COUNT) { - Some(item) => match item { - Item::Map(inner_map) => inner_map.insert(rf.name().to_owned(), Item::Value(count.clone())), - _ => unreachable!(), - }, - None => { - let mut inner_map: Map = Map::new(); - inner_map.insert(rf.name().to_owned(), Item::Value(count.clone())); - map.insert(UNDERSCORE_COUNT.to_owned(), Item::Map(inner_map)) - } - }, - }; - } -} - fn extract_aggregate_object_type<'a, 'b>(output_type: &'b OutputType<'a>) -> &'b ObjectType<'a> { match &output_type.inner { InnerOutputType::Object(obj) => obj, @@ -392,7 +373,7 @@ fn serialize_objects_with_relation( } } - let map = reorder_object_with_selection_order(result.fields.clone(), object); + let map = reorder_object_with_selection_order(&result.fields, object); let result = Item::Map(map); @@ -459,6 +440,11 @@ fn serialize_relation_selection( Ok(Item::Map(map)) } +enum SerializedField<'a, 'b> { + Model(Field, &'a OutputField<'b>), + Virtual(&'a VirtualSelection), +} + /// Serializes the given result into objects of given type. /// Doesn't validate the shape of the result set ("unchecked" result). /// Returns a vector of serialized objects (as Item::Map), grouped into a map by parent, if present. @@ -479,18 +465,36 @@ fn serialize_objects( // to prevent expensive copying during serialization). // Finally, serialize the objects based on the selected fields. - let mut object_mapping = UncheckedItemsWithParents::with_capacity(result.scalars.records.len()); - let db_field_names = result.scalars.field_names; + let mut object_mapping = UncheckedItemsWithParents::with_capacity(result.records.records.len()); + let db_field_names = result.records.field_names; let model = result.model; let fields: Vec<_> = db_field_names .iter() - .filter_map(|f| model.fields().find_from_non_virtual_by_db_name(f).ok()) + .map(|name| { + model + .fields() + .find_from_non_virtual_by_db_name(name) + .ok() + .and_then(|field| { + typ.find_field(field.name()) + .map(|out_field| SerializedField::Model(field, out_field)) + }) + .or_else(|| { + result + .virtual_fields + .iter() + .find(|f| f.db_alias() == *name) + .map(SerializedField::Virtual) + }) + // Shouldn't happen, implies that the query returned unknown fields. + .expect("Field must be a known scalar or virtual") + }) .collect(); // Write all fields, nested and list fields unordered into a map, afterwards order all into the final order. // If nothing is written to the object, write null instead. - for (r_index, record) in result.scalars.records.into_iter().enumerate() { + for record in result.records.records { let record_id = Some(record.extract_selection_result(&db_field_names, &model.primary_identifier())?); if !object_mapping.contains_key(&record.parent_id) { @@ -502,44 +506,33 @@ fn serialize_objects( let mut object = HashMap::with_capacity(values.len()); for (val, field) in values.into_iter().zip(fields.iter()) { - let out_field = typ.find_field(field.name()).unwrap(); - match field { - Field::Composite(cf) => { - object.insert(field.name().to_owned(), serialize_composite(cf, out_field, val)?); + SerializedField::Model(field, out_field) => { + if let Field::Composite(cf) = field { + object.insert(field.name().to_owned(), serialize_composite(cf, out_field, val)?); + } else if !out_field.field_type().is_object() { + object.insert(field.name().to_owned(), serialize_scalar(out_field, val)?); + } } - _ if !out_field.field_type().is_object() => { - object.insert(field.name().to_owned(), serialize_scalar(out_field, val)?); - } + SerializedField::Virtual(vs) => { + let (virtual_obj_name, nested_field_name) = vs.serialized_name(); + + let virtual_obj = object + .entry(virtual_obj_name.into()) + .or_insert(Item::Map(Map::new())) + .as_map_mut() + .expect("Virtual and scalar fields must not collide"); - _ => (), + virtual_obj.insert(nested_field_name.into(), Item::Value(vs.coerce_value(val)?)); + } } } // Write nested results write_nested_items(&record_id, &mut nested_mapping, &mut object, typ)?; - let aggr_row = result.aggregation_rows.as_ref().map(|rows| rows.get(r_index).unwrap()); - if let Some(aggr_row) = aggr_row { - write_rel_aggregation_row(aggr_row, &mut object); - } - - let mut aggr_fields = aggr_row - .map(|row| { - row.iter() - .map(|aggr_result| match aggr_result { - RelAggregationResult::Count(_, _) => UNDERSCORE_COUNT.to_owned(), - }) - .unique() - .collect() - }) - .unwrap_or_default(); - - let mut all_fields = result.fields.clone(); - all_fields.append(&mut aggr_fields); - - let map = reorder_object_with_selection_order(all_fields, object); + let map = reorder_object_with_selection_order(&result.fields, object); object_mapping.get_mut(&record.parent_id).unwrap().push(Item::Map(map)); } @@ -548,7 +541,7 @@ fn serialize_objects( } fn reorder_object_with_selection_order( - selection_order: Vec, + selection_order: &[String], mut object: HashMap, ) -> IndexMap { selection_order diff --git a/query-engine/core/src/response_ir/mod.rs b/query-engine/core/src/response_ir/mod.rs index e9a4eeb0c9a4..535a409474df 100644 --- a/query-engine/core/src/response_ir/mod.rs +++ b/query-engine/core/src/response_ir/mod.rs @@ -126,6 +126,16 @@ impl Item { } } + /// Returns a mutable reference to the underlying map, if the element is a map and the map is + /// owned. Unlike [`Item::as_map`], it doesn't allow obtaining a reference to a shared map + /// referenced via [`ItemRef`]. + pub fn as_map_mut(&mut self) -> Option<&mut Map> { + match self { + Self::Map(m) => Some(m), + _ => None, + } + } + pub fn into_map(self) -> Option { match self { Self::Map(m) => Some(m), diff --git a/query-engine/core/src/result_ast/mod.rs b/query-engine/core/src/result_ast/mod.rs index a54f333c90a2..e86c39ddf392 100644 --- a/query-engine/core/src/result_ast/mod.rs +++ b/query-engine/core/src/result_ast/mod.rs @@ -1,5 +1,5 @@ -use connector::{AggregationRow, RelAggregationRow}; -use query_structure::{ManyRecords, Model, SelectionResult}; +use connector::AggregationRow; +use query_structure::{ManyRecords, Model, SelectionResult, VirtualSelection}; #[derive(Debug, Clone)] pub(crate) enum QueryResult { @@ -55,8 +55,8 @@ pub struct RecordSelection { /// Holds an ordered list of selected field names for each contained record. pub(crate) fields: Vec, - /// Scalar field results - pub(crate) scalars: ManyRecords, + /// Selection results (includes scalar and virtual fields) + pub(crate) records: ManyRecords, /// Nested query results // Todo this is only here because reads are still resolved in one go @@ -65,8 +65,10 @@ pub struct RecordSelection { /// The model of the contained records. pub(crate) model: Model, - /// Holds an ordered list of aggregation selections results for each contained record - pub(crate) aggregation_rows: Option>, + /// The list of virtual selections included in the query result. + /// TODO: in the future it should be covered by [`RecordSelection::fields`] by storing ordered + /// `Vec` or `FieldSelection` instead of `Vec`. + pub(crate) virtual_fields: Vec, } impl From for QueryResult { diff --git a/query-engine/query-structure/src/field_selection.rs b/query-engine/query-structure/src/field_selection.rs index 8a9811e3e232..20b037f4e571 100644 --- a/query-engine/query-structure/src/field_selection.rs +++ b/query-engine/query-structure/src/field_selection.rs @@ -1,12 +1,12 @@ use crate::{ parent_container::ParentContainer, prisma_value_ext::PrismaValueExtensions, CompositeFieldRef, DomainError, Field, - Model, ModelProjection, QueryArguments, RelationField, ScalarField, ScalarFieldRef, SelectionResult, - TypeIdentifier, + Filter, Model, ModelProjection, QueryArguments, RelationField, RelationFieldRef, ScalarField, ScalarFieldRef, + SelectionResult, TypeIdentifier, }; use itertools::Itertools; use prisma_value::PrismaValue; use psl::schema_ast::ast::FieldArity; -use std::fmt::Display; +use std::{borrow::Cow, fmt::Display}; /// A selection of fields from a model. #[derive(Debug, Clone, PartialEq, Default, Hash, Eq)] @@ -35,6 +35,7 @@ impl FieldSelection { .unwrap_or(false), // TODO: Relation selections are ignored for now to prevent breaking the existing query-based strategy to resolve relations. SelectedField::Relation(_) => true, + SelectedField::Virtual(vs) => self.contains(&vs.db_alias()), }) } @@ -42,16 +43,44 @@ impl FieldSelection { self.selections.iter() } + pub fn virtuals(&self) -> impl Iterator { + self.selections().filter_map(|field| match field { + SelectedField::Virtual(ref vs) => Some(vs), + _ => None, + }) + } + + pub fn virtuals_owned(&self) -> Vec { + self.virtuals().cloned().collect() + } + + pub fn without_relations(&self) -> Self { + FieldSelection::new( + self.selections() + .filter(|field| !matches!(field, SelectedField::Relation(_))) + .cloned() + .collect(), + ) + } + + pub fn into_virtuals_last(self) -> Self { + let (virtuals, non_virtuals): (Vec<_>, Vec<_>) = self + .into_iter() + .partition(|field| matches!(field, SelectedField::Virtual(_))); + + FieldSelection::new(non_virtuals.into_iter().chain(virtuals).collect()) + } + /// Returns all Prisma (e.g. schema model field) names of contained fields. /// Does _not_ recurse into composite selections and only iterates top level fields. pub fn prisma_names(&self) -> impl Iterator + '_ { - self.selections.iter().map(|f| f.prisma_name().to_owned()) + self.selections.iter().map(|f| f.prisma_name().into_owned()) } /// Returns all database (e.g. column or document field) names of contained fields. /// Does _not_ recurse into composite selections and only iterates level fields. pub fn db_names(&self) -> impl Iterator + '_ { - self.selections.iter().map(|f| f.db_name().to_owned()) + self.selections.iter().map(|f| f.db_name().into_owned()) } /// Checked if a field of prisma name `name` is present in this `FieldSelection`. @@ -69,6 +98,7 @@ impl FieldSelection { SelectedField::Scalar(sf) => sf.clone().into(), SelectedField::Composite(cf) => cf.field.clone().into(), SelectedField::Relation(rs) => rs.field.clone().into(), + SelectedField::Virtual(vs) => vs.field(), }) .collect() } @@ -82,6 +112,7 @@ impl FieldSelection { SelectedField::Scalar(sf) => Some(sf.clone()), SelectedField::Composite(_) => None, SelectedField::Relation(_) => None, + SelectedField::Virtual(_) => None, }) .collect::>(); @@ -158,6 +189,7 @@ impl FieldSelection { SelectedField::Relation(rf) if rf.field.is_list() => Some((TypeIdentifier::Json, FieldArity::Required)), SelectedField::Relation(rf) => Some((TypeIdentifier::Json, rf.field.arity())), SelectedField::Composite(_) => None, + SelectedField::Virtual(vs) => Some(vs.type_identifier_with_arity()), }) .collect() } @@ -172,15 +204,20 @@ impl FieldSelection { pub fn into_projection(self) -> ModelProjection { self.into() } + + pub fn has_virtual_fields(&self) -> bool { + self.selections() + .any(|field| matches!(field, SelectedField::Virtual(_))) + } } /// A selected field. Can be contained on a model or composite type. -// Todo: Think about virtual selections like aggregations. #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum SelectedField { Scalar(ScalarFieldRef), Composite(CompositeSelection), Relation(RelationSelection), + Virtual(VirtualSelection), } #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -213,20 +250,76 @@ impl RelationSelection { } } +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum VirtualSelection { + RelationCount(RelationFieldRef, Option), +} + +impl VirtualSelection { + pub fn db_alias(&self) -> String { + match self { + Self::RelationCount(rf, _) => format!("_aggr_count_{}", rf.name()), + } + } + + pub fn serialized_name(&self) -> (&'static str, &str) { + match self { + // TODO: we can't use UNDERSCORE_COUNT here because it would require a circular + // dependency between `schema` and `query-structure` crates. + Self::RelationCount(rf, _) => ("_count", rf.name()), + } + } + + pub fn model(&self) -> Model { + match self { + Self::RelationCount(rf, _) => rf.model(), + } + } + + pub fn coerce_value(&self, value: PrismaValue) -> crate::Result { + match self { + Self::RelationCount(_, _) => match value { + PrismaValue::Null => Ok(PrismaValue::Int(0)), + _ => value.coerce(TypeIdentifier::Int), + }, + } + } + + pub fn field(&self) -> Field { + match self { + Self::RelationCount(rf, _) => rf.clone().into(), + } + } + + pub fn type_identifier_with_arity(&self) -> (TypeIdentifier, FieldArity) { + match self { + Self::RelationCount(_, _) => (TypeIdentifier::Int, FieldArity::Required), + } + } +} + +impl Display for VirtualSelection { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.db_alias()) + } +} + impl SelectedField { - pub fn prisma_name(&self) -> &str { + pub fn prisma_name(&self) -> Cow<'_, str> { match self { - SelectedField::Scalar(sf) => sf.name(), - SelectedField::Composite(cf) => cf.field.name(), - SelectedField::Relation(rs) => rs.field.name(), + SelectedField::Scalar(sf) => sf.name().into(), + SelectedField::Composite(cf) => cf.field.name().into(), + SelectedField::Relation(rs) => rs.field.name().into(), + SelectedField::Virtual(vs) => vs.db_alias().into(), } } - pub fn db_name(&self) -> &str { + pub fn db_name(&self) -> Cow<'_, str> { match self { - SelectedField::Scalar(sf) => sf.db_name(), - SelectedField::Composite(cs) => cs.field.db_name(), - SelectedField::Relation(rs) => rs.field.name(), + SelectedField::Scalar(sf) => sf.db_name().into(), + SelectedField::Composite(cs) => cs.field.db_name().into(), + SelectedField::Relation(rs) => rs.field.name().into(), + SelectedField::Virtual(vs) => vs.db_alias().into(), } } @@ -242,15 +335,17 @@ impl SelectedField { SelectedField::Scalar(sf) => sf.container(), SelectedField::Composite(cs) => cs.field.container(), SelectedField::Relation(rs) => ParentContainer::from(rs.field.model()), + SelectedField::Virtual(vs) => ParentContainer::from(vs.model()), } } /// Coerces a value to fit the selection. If the conversion is not possible, an error will be thrown. pub(crate) fn coerce_value(&self, value: PrismaValue) -> crate::Result { match self { - SelectedField::Scalar(sf) => value.coerce(&sf.type_identifier()), + SelectedField::Scalar(sf) => value.coerce(sf.type_identifier()), SelectedField::Composite(cs) => cs.coerce_value(value), SelectedField::Relation(_) => todo!(), + SelectedField::Virtual(vs) => vs.coerce_value(value), } } @@ -277,6 +372,7 @@ impl CompositeSelection { .map(|cs| cs.is_superset_of(other_cs)) .unwrap_or(false), SelectedField::Relation(_) => true, // A composite selection cannot hold relations. + SelectedField::Virtual(vs) => self.contains(&vs.db_alias()), }) } @@ -358,6 +454,7 @@ impl Display for SelectedField { rs.field, rs.selections.iter().map(|selection| format!("{selection}")).join(", ") ), + SelectedField::Virtual(vs) => write!(f, "{vs}"), } } } diff --git a/query-engine/query-structure/src/filter/into_filter.rs b/query-engine/query-structure/src/filter/into_filter.rs index eaf4711628fe..660429999961 100644 --- a/query-engine/query-structure/src/filter/into_filter.rs +++ b/query-engine/query-structure/src/filter/into_filter.rs @@ -16,6 +16,7 @@ impl IntoFilter for SelectionResult { SelectedField::Scalar(sf) => sf.equals(value), SelectedField::Composite(_) => unreachable!(), // [Composites] todo SelectedField::Relation(_) => unreachable!(), + SelectedField::Virtual(_) => unreachable!(), }) .collect(); diff --git a/query-engine/query-structure/src/prisma_value_ext.rs b/query-engine/query-structure/src/prisma_value_ext.rs index 09e052ea844b..49a35d61b277 100644 --- a/query-engine/query-structure/src/prisma_value_ext.rs +++ b/query-engine/query-structure/src/prisma_value_ext.rs @@ -3,12 +3,12 @@ use crate::DomainError; use bigdecimal::ToPrimitive; pub(crate) trait PrismaValueExtensions { - fn coerce(self, to_type: &TypeIdentifier) -> crate::Result; + fn coerce(self, to_type: TypeIdentifier) -> crate::Result; } impl PrismaValueExtensions for PrismaValue { // Todo this is not exhaustive for now. - fn coerce(self, to_type: &TypeIdentifier) -> crate::Result { + fn coerce(self, to_type: TypeIdentifier) -> crate::Result { let coerced = match (self, to_type) { // Trivial cases (PrismaValue::Null, _) => PrismaValue::Null, diff --git a/query-engine/query-structure/src/projections/model_projection.rs b/query-engine/query-structure/src/projections/model_projection.rs index 0d1a8f4b5171..de9a689cac34 100644 --- a/query-engine/query-structure/src/projections/model_projection.rs +++ b/query-engine/query-structure/src/projections/model_projection.rs @@ -31,6 +31,7 @@ impl From<&FieldSelection> for ModelProjection { SelectedField::Scalar(sf) => Some(sf.clone().into()), SelectedField::Composite(_cf) => None, SelectedField::Relation(_) => None, + SelectedField::Virtual(_) => None, }) .collect(), } diff --git a/query-engine/query-structure/src/record.rs b/query-engine/query-structure/src/record.rs index a9c3328262bc..15841d856ba7 100644 --- a/query-engine/query-structure/src/record.rs +++ b/query-engine/query-structure/src/record.rs @@ -49,7 +49,7 @@ impl ManyRecords { pub fn empty(selected_fields: &FieldSelection) -> Self { Self { records: Vec::new(), - field_names: selected_fields.prisma_names().collect(), + field_names: selected_fields.db_names().collect(), } } @@ -172,7 +172,7 @@ impl Record { let pairs: Vec<_> = extraction_selection .selections() .map(|selection| { - self.get_field_value(field_names, selection.db_name()) + self.get_field_value(field_names, &selection.db_name()) .and_then(|val| Ok((selection.clone(), selection.coerce_value(val.clone())?))) }) .collect::>>()?; diff --git a/query-engine/query-structure/src/selection_result.rs b/query-engine/query-structure/src/selection_result.rs index 6f87ec74f6ce..2c6b379ee2b1 100644 --- a/query-engine/query-structure/src/selection_result.rs +++ b/query-engine/query-structure/src/selection_result.rs @@ -1,6 +1,6 @@ use crate::{DomainError, FieldSelection, PrismaValue, ScalarFieldRef, SelectedField}; use itertools::Itertools; -use std::convert::TryFrom; +use std::{borrow::Cow, convert::TryFrom}; /// Represents a set of results. #[derive(Default, Clone, PartialEq, Eq, Hash)] @@ -61,7 +61,7 @@ impl SelectionResult { self.len() == 0 } - pub fn db_names(&self) -> impl Iterator + '_ { + pub fn db_names(&self) -> impl Iterator> + '_ { self.pairs.iter().map(|(field, _)| field.db_name()) } @@ -95,6 +95,7 @@ impl SelectionResult { SelectedField::Scalar(sf) => Some(sf.clone()), SelectedField::Composite(_) => None, SelectedField::Relation(_) => None, + SelectedField::Virtual(_) => None, }) .collect();