fix: make sure columns are fully qualified in aggregates (#5629)

Restores the original behavior where group by aggregations would have
column names with explicit table names. This behavior cannot be used
with non-group-by aggregates, because they sometimes rely on referring
to columns from a subquery.


[TML-1483](https://linear.app/prisma-company/issue/TML-1483/fix-groupby-regression)

Fixes prisma/prisma#28084

Author: jacek-prisma

query-compiler/query-compiler/src/data_mapper.rs

@@ -298,7 +298,7 @@ fn get_result_node_for_aggregation(
             sel.identifiers().map(move |ident| {
                 let db_alias = ident.db_alias();
                 let type_info = FieldTypeInformation::new(ident.typ, ident.arity, None);
-                (ident.name, sel.aggregation_name(), db_alias, type_info)
+                (ident.field.name(), sel.aggregation_name(), db_alias, type_info)
             })
         })
         .sorted_by_key(|(name, prefix, _, _)| ordered_set.get_index_of(&(*prefix, *name)))

query-compiler/query-compiler/tests/snapshots/queries__queries@group-by.json.snap

@@ -11,8 +11,9 @@ dataMap {
         _all: Int (_count$_all)
     }
 }
-query «SELECT COUNT("id") AS "_count$id", COUNT("title") AS "_count$title",
-       COUNT("userId") AS "_count$userId", COUNT(*) AS "_count$_all" FROM
-       "public"."Post" WHERE "public"."Post"."title"::text LIKE $1 GROUP BY
-       "public"."Post"."title" OFFSET $2»
+query «SELECT COUNT("public"."Post"."id") AS "_count$id",
+       COUNT("public"."Post"."title") AS "_count$title",
+       COUNT("public"."Post"."userId") AS "_count$userId", COUNT(*) AS
+       "_count$_all" FROM "public"."Post" WHERE "public"."Post"."title"::text
+       LIKE $1 GROUP BY "public"."Post"."title" OFFSET $2»
 params [const(String("%something%")), const(BigInt(0))]

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

@@ -577,6 +577,48 @@ mod aggregation_group_by {
         Ok(())
     }
 
+    fn schema_28084() -> String {
+        let schema = indoc! {
+            r#"model Teacher {
+              #id(id, String, @id)
+              name      String
+              age       Int
+              assistant Assistant?
+            }
+
+            model Assistant {
+              #id(id, String, @id)
+              name      String
+              age       Int?
+              teacher   Teacher @relation(fields: [teacherId], references: [id])
+              teacherId String  @unique
+            }
+            "#
+        };
+
+        schema.to_owned()
+    }
+
+    // regression test for https://github.com/prisma/prisma/issues/28084
+    #[connector_test(schema(schema_28084))]
+    async fn regression_28084(runner: Runner) -> TestResult<()> {
+        run_query!(
+            &runner,
+            r#"mutation { createOneTeacher(data: {
+                id: "a1",
+                name: "Alice",
+                age: 30,
+                assistant: { create: { id: "j1", name: "John" } } }
+            ) { id } }"#
+        );
+
+        insta::assert_snapshot!(
+          run_query!(&runner, r#"{ groupByTeacher(by: [name], where: { assistant: { name: { contains: "John" } } }) { _sum { age } } }"#),
+          @r###"{"data":{"groupByTeacher":[{"_sum":{"age":30}}]}}"###
+        );
+
+        Ok(())
+    }
     /// Error cases
 
     #[connector_test]

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

@@ -118,7 +118,7 @@ pub fn from_aggregation_selection(selection: &AggregationSelection) -> OutputMet
 
     for ident in selection.identifiers() {
         map.insert(
-            ident.field_db_name.into(),
+            ident.field.db_name().into(),
             OutputMeta::Scalar(ScalarOutputMeta {
                 ident: ident.typ.id,
                 default: None,

query-engine/query-builders/sql-query-builder/src/model_extensions/column.rs

@@ -50,6 +50,7 @@ impl AsColumns for ModelProjection {
 
 pub trait AsColumn {
     fn as_column(&self, ctx: &Context<'_>) -> Column<'static>;
+    fn as_column_with_style(&self, ctx: &Context<'_>, style: ColumnStyle) -> Column<'static>;
 }
 
 impl AsColumns for Field {
@@ -90,15 +91,33 @@ where
 
 impl AsColumn for ScalarField {
     fn as_column(&self, ctx: &Context<'_>) -> Column<'static> {
-        // Unwrap is safe: SQL connectors do not anything other than models as field containers.
-        let full_table_name = super::table::db_name_with_schema(&self.container().as_model().unwrap(), ctx);
-        let col = self.db_name().to_string();
+        self.as_column_with_style(ctx, ColumnStyle::ExplicitTable)
+    }
+
+    fn as_column_with_style(&self, ctx: &Context<'_>, style: ColumnStyle) -> Column<'static> {
+        let col = match style {
+            ColumnStyle::ExplicitTable => {
+                // Unwrap is safe: SQL connectors do not anything other than models as field containers.
+                let full_table_name = super::table::db_name_with_schema(&self.container().as_model().unwrap(), ctx);
+                let col = self.db_name().to_string();
+                Column::from((full_table_name, col))
+            }
+            ColumnStyle::ImplicitTable => Column::new(self.db_name().to_owned()),
+        };
 
-        Column::from((full_table_name, col))
-            .type_family(self.type_family())
+        col.type_family(self.type_family())
             .native_column_type(self.native_type().map(|nt| NativeColumnType::from(nt.name())))
             .set_is_enum(self.type_identifier().is_enum())
             .set_is_list(self.is_list())
             .default(quaint::ast::DefaultValue::Generated)
     }
 }
+
+/// The style of column rendering.
+#[derive(Debug, Clone, Copy)]
+pub enum ColumnStyle {
+    /// The column is rendered with an explicit table name, e.g. `User.id`.
+    ExplicitTable,
+    /// The column is rendered without a table name, e.g. `id`.
+    ImplicitTable,
+}

query-engine/query-builders/sql-query-builder/src/read.rs

@@ -6,7 +6,7 @@ use crate::{
     context::Context,
     cursor_condition,
     filter::FilterBuilder,
-    model_extensions::{AsColumn, AsColumns, AsTable},
+    model_extensions::{AsColumn, AsColumns, AsTable, ColumnStyle},
     nested_aggregations,
     ordering::OrderByBuilder,
     sql_trace::SqlTraceComment,
@@ -184,7 +184,7 @@ pub fn aggregate(
     let sub_table = Table::from(sub_query).alias("sub");
     selections.iter().fold(
         Select::from_table(sub_table).add_traceparent(ctx.traceparent),
-        |acc, sel| apply_aggregate_selections(acc, sel, ctx),
+        |acc, sel| apply_aggregate_selections(acc, sel, ctx, ColumnStyle::ImplicitTable),
     )
 }
 
@@ -197,9 +197,9 @@ pub fn group_by_aggregate(
     ctx: &Context<'_>,
 ) -> Select<'static> {
     let (base_query, _) = args.into_select(model, &[], ctx);
-    let select_query = selections
-        .iter()
-        .fold(base_query, |acc, sel| apply_aggregate_selections(acc, sel, ctx));
+    let select_query = selections.iter().fold(base_query, |acc, sel| {
+        apply_aggregate_selections(acc, sel, ctx, ColumnStyle::ExplicitTable)
+    });
 
     let grouped = group_by
         .into_iter()
@@ -221,32 +221,50 @@ fn apply_aggregate_selections(
     select: Select<'static>,
     selection: &AggregationSelection,
     ctx: &Context,
+    col_style: ColumnStyle,
 ) -> Select<'static> {
     match selection {
-        AggregationSelection::Field(field) => select.column(field.as_column(ctx).set_is_selected(true)),
+        AggregationSelection::Field(field) => {
+            select.column(field.as_column_with_style(ctx, col_style).set_is_selected(true))
+        }
 
-        AggregationSelection::Count { all, .. } => selection.identifiers().fold(select, |select, next_field| {
-            let expr = if all.is_some() && next_field.name == "_all" {
-                asterisk()
-            } else {
-                Column::from(next_field.field_db_name.to_owned()).into()
+        AggregationSelection::Count { .. } => selection.identifiers().fold(select, |select, next_field| {
+            let expr = match next_field.field {
+                SelectionField::All => asterisk(),
+                SelectionField::Scalar(field) => field.as_column_with_style(ctx, col_style).into(),
             };
             select.value(count(expr).alias(next_field.db_alias().into_owned()))
         }),
 
         AggregationSelection::Average(_) => selection.identifiers().fold(select, |select, next_field| {
-            select
-                .value(avg(Column::from(next_field.field_db_name.to_owned())).alias(next_field.db_alias().into_owned()))
+            select.value(
+                avg(next_field
+                    .field
+                    .as_scalar()
+                    .unwrap()
+                    .as_column_with_style(ctx, col_style))
+                .alias(next_field.db_alias().into_owned()),
+            )
         }),
 
         AggregationSelection::Sum(_) => selection.identifiers().fold(select, |select, next_field| {
-            select
-                .value(sum(Column::from(next_field.field_db_name.to_owned())).alias(next_field.db_alias().into_owned()))
+            select.value(
+                sum(next_field
+                    .field
+                    .as_scalar()
+                    .unwrap()
+                    .as_column_with_style(ctx, col_style))
+                .alias(next_field.db_alias().into_owned()),
+            )
         }),
 
         AggregationSelection::Min(_) => selection.identifiers().fold(select, |select, next_field| {
             select.value(
-                min(Column::from(next_field.field_db_name.to_owned())
+                min(next_field
+                    .field
+                    .as_scalar()
+                    .unwrap()
+                    .as_column_with_style(ctx, col_style)
                     .set_is_enum(next_field.typ.id.is_enum())
                     .set_is_selected(true))
                 .alias(next_field.db_alias().into_owned()),
@@ -255,7 +273,11 @@ fn apply_aggregate_selections(
 
         AggregationSelection::Max(_) => selection.identifiers().fold(select, |select, next_field| {
             select.value(
-                max(Column::from(next_field.field_db_name.to_owned())
+                max(next_field
+                    .field
+                    .as_scalar()
+                    .unwrap()
+                    .as_column_with_style(ctx, col_style)
                     .set_is_enum(next_field.typ.id.is_enum())
                     .set_is_selected(true))
                 .alias(next_field.db_alias().into_owned()),

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

@@ -62,9 +62,8 @@ impl AggregationSelection {
             AggregationSelection::Count { all, fields } => Either::Right(
                 self.map_field_types(fields, |_| TypeIdentifier::Int, |_| FieldArity::Required)
                     .chain(all.iter().map(|all| SelectionIdentifier {
-                        name: all.name(),
+                        field: SelectionField::All,
                         aggregation_name: self.aggregation_name(),
-                        field_db_name: all.db_name(),
                         typ: all.r#type(),
                         arity: all.arity(),
                     })),
@@ -90,12 +89,11 @@ impl AggregationSelection {
         arity_mapper: fn(FieldArity) -> FieldArity,
     ) -> impl Iterator<Item = SelectionIdentifier<'a>> {
         let aggregation_name = self.aggregation_name();
-        fields.iter().map(move |f| SelectionIdentifier {
-            name: f.name(),
+        fields.iter().map(move |field| SelectionIdentifier {
+            field: SelectionField::Scalar(field),
             aggregation_name,
-            field_db_name: f.db_name(),
-            typ: f.dm.clone().zip(type_mapper(f.type_identifier())),
-            arity: arity_mapper(f.arity()),
+            typ: field.dm.clone().zip(type_mapper(field.type_identifier())),
+            arity: arity_mapper(field.arity()),
         })
     }
 }
@@ -110,16 +108,6 @@ impl CountAllAggregationSelection {
         CountAllAggregationSelection { dm }
     }
 
-    #[inline]
-    pub fn name(&self) -> &str {
-        "_all"
-    }
-
-    #[inline]
-    pub fn db_name(&self) -> &str {
-        "_all"
-    }
-
     #[inline]
     pub fn type_identifier(&self) -> TypeIdentifier {
         TypeIdentifier::Int
@@ -142,8 +130,7 @@ impl<I> From<&'_ Zipper<I>> for CountAllAggregationSelection {
 }
 
 pub struct SelectionIdentifier<'a> {
-    pub name: &'a str,
-    pub field_db_name: &'a str,
+    pub field: SelectionField<'a>,
     pub aggregation_name: Option<&'static str>,
     pub typ: Type,
     pub arity: FieldArity,
@@ -153,8 +140,36 @@ impl<'a> SelectionIdentifier<'a> {
     pub fn db_alias(&self) -> Cow<'a, str> {
         const FIELD_SEPARATOR: &str = "$";
         self.aggregation_name
-            .map_or(Cow::Borrowed(self.field_db_name), |aggregation| {
-                Cow::Owned(format!("{aggregation}{FIELD_SEPARATOR}{}", self.field_db_name))
+            .map_or(Cow::Borrowed(self.field.db_name()), |aggregation| {
+                Cow::Owned(format!("{aggregation}{FIELD_SEPARATOR}{}", self.field.db_name()))
             })
     }
 }
+
+pub enum SelectionField<'a> {
+    Scalar(&'a ScalarFieldRef),
+    All,
+}
+
+impl<'a> SelectionField<'a> {
+    pub fn name(&self) -> &'a str {
+        match self {
+            Self::Scalar(f) => f.name(),
+            Self::All => "_all",
+        }
+    }
+
+    pub fn db_name(&self) -> &'a str {
+        match self {
+            Self::Scalar(f) => f.db_name(),
+            Self::All => "_all",
+        }
+    }
+
+    pub fn as_scalar(&self) -> Option<&'a ScalarFieldRef> {
+        match self {
+            Self::Scalar(f) => Some(f),
+            Self::All => None,
+        }
+    }
+}