Skip to content

Commit 1c4c6a0

Browse files
jacek-prismajacek-prisma
authored
fix(qc): fix nested aggregations (#5572)
[ORM-1170](https://linear.app/prisma-company/issue/ORM-1170/fix-failing-filter-count-relations-test) Refactors the code to re-use `into_select` which automatically handles nested aggregations. It also adds cursor filters that we were previously missing and relied on finding the cursor in memory. It changed a lot of snapshots because I needed to swap the two sides of the `JOIN` in order to re-use existing code. Also fixes a test on the client: prisma/prisma#27792
1 parent 3dbd971 commit 1c4c6a0

10 files changed

Lines changed: 165 additions & 52 deletions

‎query-compiler/query-compiler/tests/data/aggregate-nested-m2m.json‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
{
2+
"modelName": "User",
3+
"action": "findFirst",
4+
"query": {
5+
"arguments": {
6+
"where": {
7+
"email": {
8+
"contains": "prisma.io"
9+
}
10+
},
11+
"relationLoadStrategy": "query"
12+
},
13+
"selection": {
14+
"posts": {
15+
"arguments": {
16+
"orderBy": {
17+
"id": "asc"
18+
}
19+
},
20+
"selection": {
21+
"_count": {
22+
"arguments": {},
23+
"selection": {
24+
"categories": {
25+
"arguments": {
26+
"where": {
27+
"name": {
28+
"contains": "tech"
29+
}
30+
}
31+
},
32+
"selection": {}
33+
}
34+
}
35+
}
36+
}
37+
}
38+
}
39+
}
40+
}

‎query-compiler/query-compiler/tests/snapshots/queries__queries@aggregate-nested-m2m.json.snap‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
---
2+
source: query-compiler/query-compiler/tests/queries.rs
3+
expression: pretty
4+
input_file: query-compiler/query-compiler/tests/data/aggregate-nested-m2m.json
5+
---
6+
dataMap {
7+
posts (from @nested$posts): {
8+
_count (inlined): {
9+
categories: Int (_aggr_count_categories)
10+
}
11+
}
12+
}
13+
unique (let @parent = query «SELECT "public"."User"."id" FROM "public"."User"
14+
WHERE "public"."User"."email"::text LIKE $1 LIMIT
15+
$2 OFFSET $3»
16+
params [const(String("%prisma.io%")), const(BigInt(1)),
17+
const(BigInt(0))]
18+
in let @parent$id = mapField id (get @parent)
19+
in join (get @parent)
20+
with (query «SELECT "public"."Post"."id", "public"."Post"."userId",
21+
COALESCE("aggr_selection_0_Category"."_aggr_count_categories",
22+
0) AS "_aggr_count_categories" FROM "public"."Post" LEFT JOIN
23+
(SELECT "public"."_CategoryToPost"."B",
24+
COUNT("public"."_CategoryToPost"."B") AS
25+
"_aggr_count_categories" FROM "public"."Category" LEFT JOIN
26+
"public"."_CategoryToPost" ON ("public"."Category"."id" =
27+
"public"."_CategoryToPost"."A") WHERE
28+
"public"."Category"."name"::text LIKE $1 GROUP BY
29+
"public"."_CategoryToPost"."B") AS
30+
"aggr_selection_0_Category" ON ("public"."Post"."id" =
31+
"aggr_selection_0_Category"."B") WHERE
32+
"public"."Post"."userId" IN [$2] ORDER BY
33+
"public"."Post"."id" ASC OFFSET $3»
34+
params [const(String("%tech%")), var(@parent$id as Int),
35+
const(BigInt(0))]) on left.(id) = right.(userId) as @nested$posts)

‎query-compiler/query-compiler/tests/snapshots/queries__queries@create-nested-connectOrCreate-mixed.json.snap‎

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -113,15 +113,20 @@ transaction
113113
OFFSET $2»
114114
params [var(@parent$userId as Int),
115115
const(BigInt(0))]) on unique left.(userId) = right.(id) as @nested$user,
116-
(query «SELECT "t1"."id",
117-
"t1"."name", "t0"."B" AS
116+
(query «SELECT
117+
"public"."Category"."id",
118+
"public"."Category"."name",
119+
"t0"."B" AS
118120
"CategoryToPost@Post" FROM
121+
"public"."Category" INNER
122+
JOIN
119123
"public"."_CategoryToPost"
120-
AS "t0" INNER JOIN
121-
"public"."Category" AS "t1"
122-
ON "t0"."A" = "t1"."id"
123-
WHERE "t0"."B" = $1»
124-
params [var(@parent$id as Int)]) on left.(id) = right.(CategoryToPost@Post) as @nested$categories
124+
AS "t0" ON "t0"."A" =
125+
"public"."Category"."id"
126+
WHERE (1=1 AND "t0"."B" =
127+
$1) OFFSET $2»
128+
params [var(@parent$id as Int),
129+
const(BigInt(0))]) on left.(id) = right.(CategoryToPost@Post) as @nested$categories
125130
in get 11
126131
in get 11
127132
in get 0

‎query-compiler/query-compiler/tests/snapshots/queries__queries@create-nested-connectOrCreate-one2m.json.snap‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -98,14 +98,16 @@ transaction
9898
$2»
9999
params [var(@parent$userId as Int),
100100
const(BigInt(0))]) on unique left.(userId) = right.(id) as @nested$user,
101-
(query «SELECT "t1"."id", "t1"."name",
101+
(query «SELECT "public"."Category"."id",
102+
"public"."Category"."name",
102103
"t0"."B" AS "CategoryToPost@Post"
103-
FROM "public"."_CategoryToPost" AS
104-
"t0" INNER JOIN
105-
"public"."Category" AS "t1" ON
106-
"t0"."A" = "t1"."id" WHERE
107-
"t0"."B" = $1»
108-
params [var(@parent$id as Int)]) on left.(id) = right.(CategoryToPost@Post) as @nested$categories
104+
FROM "public"."Category" INNER
105+
JOIN "public"."_CategoryToPost" AS
106+
"t0" ON "t0"."A" =
107+
"public"."Category"."id" WHERE
108+
(1=1 AND "t0"."B" = $1) OFFSET $2»
109+
params [var(@parent$id as Int),
110+
const(BigInt(0))]) on left.(id) = right.(CategoryToPost@Post) as @nested$categories
109111
in get 7
110112
in get 7
111113
in get 0

‎query-compiler/query-compiler/tests/snapshots/queries__queries@query-m2m.json.snap‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,10 @@ let @parent = unique (query «SELECT "public"."Post"."id",
2020
const(BigInt(0))])
2121
in let @parent$id = mapField id (get @parent)
2222
in join (get @parent)
23-
with (query «SELECT "t1"."id", "t1"."name", "t0"."B" AS
24-
"CategoryToPost@Post" FROM "public"."_CategoryToPost" AS "t0"
25-
INNER JOIN "public"."Category" AS "t1" ON "t0"."A" =
26-
"t1"."id" WHERE "t0"."B" = $1»
27-
params [var(@parent$id as Int)]) on left.(id) = right.(CategoryToPost@Post) as @nested$categories
23+
with (query «SELECT "public"."Category"."id", "public"."Category"."name",
24+
"t0"."B" AS "CategoryToPost@Post" FROM "public"."Category"
25+
INNER JOIN "public"."_CategoryToPost" AS "t0" ON "t0"."A" =
26+
"public"."Category"."id" WHERE (1=1 AND "t0"."B" = $1) OFFSET
27+
$2»
28+
params [var(@parent$id as Int),
29+
const(BigInt(0))]) on left.(id) = right.(CategoryToPost@Post) as @nested$categories

‎query-compiler/query-compiler/tests/snapshots/queries__queries@query-many-m2m.json.snap‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,10 @@ let @parent = query «SELECT "public"."Post"."id", "public"."Post"."title",
1818
params [const(BigInt(0))]
1919
in let @parent$id = mapField id (get @parent)
2020
in join (get @parent)
21-
with (query «SELECT "t1"."id", "t1"."name", "t0"."B" AS
22-
"CategoryToPost@Post" FROM "public"."_CategoryToPost" AS "t0"
23-
INNER JOIN "public"."Category" AS "t1" ON "t0"."A" =
24-
"t1"."id" WHERE "t0"."B" IN [$1]»
25-
params [var(@parent$id as Int)]) on left.(id) = right.(CategoryToPost@Post) as @nested$categories
21+
with (query «SELECT "public"."Category"."id", "public"."Category"."name",
22+
"t0"."B" AS "CategoryToPost@Post" FROM "public"."Category"
23+
INNER JOIN "public"."_CategoryToPost" AS "t0" ON "t0"."A" =
24+
"public"."Category"."id" WHERE (1=1 AND "t0"."B" IN [$1])
25+
OFFSET $2»
26+
params [var(@parent$id as Int),
27+
const(BigInt(0))]) on left.(id) = right.(CategoryToPost@Post) as @nested$categories

‎query-compiler/query-compiler/tests/snapshots/queries__queries@query-one2m-pagination.json.snap‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,12 @@ in let @parent$id = mapField id (get @parent)
2222
skip 1
2323
take 1
2424
distinct by (id, CategoryToPost@Post)
25-
(query «SELECT "t1"."id", "t1"."name", "t0"."B" AS
26-
"CategoryToPost@Post" FROM "public"."_CategoryToPost" AS "t0"
27-
INNER JOIN "public"."Category" AS "t1" ON "t0"."A" =
28-
"t1"."id" WHERE "t0"."B" IN [$1] ORDER BY "t1"."id" ASC»
29-
params [var(@parent$id as Int)])) on left.(id) = right.(CategoryToPost@Post) as @nested$categories
25+
(query «SELECT "public"."Category"."id", "public"."Category"."name",
26+
"t0"."B" AS "CategoryToPost@Post" FROM "public"."Category"
27+
INNER JOIN "public"."_CategoryToPost" AS "t0" ON "t0"."A" =
28+
"public"."Category"."id" WHERE ("public"."Category"."id" >=
29+
(SELECT "public"."Category"."id" FROM "public"."Category"
30+
WHERE ("public"."Category"."id") = ($1)) AND "t0"."B" IN
31+
[$2]) ORDER BY "public"."Category"."id" ASC OFFSET $3»
32+
params [const(BigInt(1)), var(@parent$id as Int),
33+
const(BigInt(0))])) on left.(id) = right.(CategoryToPost@Post) as @nested$categories

‎query-compiler/query-compiler/tests/snapshots/queries__queries@test.json.snap‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
---
2+
source: query-compiler/query-compiler/tests/queries.rs
3+
expression: pretty
4+
input_file: query-compiler/query-compiler/tests/data/test.json
5+
---
6+
dataMap {
7+
posts (from posts): {
8+
_count (from _count): {
9+
categories: Int (categories)
10+
}
11+
}
12+
}
13+
unique (query «SELECT "t0"."id", "User_posts"."__prisma_data__" AS "posts" FROM
14+
"public"."User" AS "t0" LEFT JOIN LATERAL (SELECT
15+
COALESCE(JSONB_AGG("__prisma_data__"), '[]') AS "__prisma_data__"
16+
FROM (SELECT "t3"."__prisma_data__" FROM (SELECT
17+
JSONB_BUILD_OBJECT('_count', JSONB_BUILD_OBJECT('categories',
18+
COALESCE("t5"."_aggr_count_categories", 0))) AS
19+
"__prisma_data__", "t2"."id" FROM (SELECT "t1".* FROM
20+
"public"."Post" AS "t1" WHERE "t0"."id" = "t1"."userId" /* root
21+
select */) AS "t2" LEFT JOIN LATERAL (SELECT COUNT(*) AS
22+
"_aggr_count_categories" FROM "public"."Category" AS "t6" LEFT
23+
JOIN "public"."_CategoryToPost" AS "t7" ON "t7"."A" = "t6"."id"
24+
WHERE ("t7"."B" = "t2"."id" AND "t6"."name"::text LIKE $1)) AS
25+
"t5" ON true /* inner select */) AS "t3" ORDER BY "t3"."id" ASC
26+
/* middle select */) AS "t4" /* outer select */) AS "User_posts"
27+
ON true WHERE "t0"."email" = $2 LIMIT $3»
28+
params [const(String("%tech%")), const(String("john@doe.com")),
29+
const(BigInt(1))])

‎query-compiler/query-compiler/tests/snapshots/queries__queries@update-m2m-disconnect.json.snap‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,12 @@ transaction
1616
params [const(BigInt(1)), const(BigInt(1)),
1717
const(BigInt(0))])
1818
in let 0$id = mapField id (get 0)
19-
in let 1 = query «SELECT "t1"."id", "t0"."B" AS "CategoryToPost@Post" FROM
20-
"public"."_CategoryToPost" AS "t0" INNER JOIN
21-
"public"."Category" AS "t1" ON "t0"."A" = "t1"."id"
22-
WHERE ("t1"."id" = $1 AND 1=1 AND "t0"."B" IN [$2])»
23-
params [const(BigInt(1)), var(0$id as Int)]
19+
in let 1 = query «SELECT "public"."Category"."id", "t0"."B" AS
20+
"CategoryToPost@Post" FROM "public"."Category" INNER
21+
JOIN "public"."_CategoryToPost" AS "t0" ON "t0"."A" =
22+
"public"."Category"."id" WHERE ("public"."Category"."id"
23+
= $1 AND 1=1 AND "t0"."B" IN [$2]) OFFSET $3»
24+
params [const(BigInt(1)), var(0$id as Int), const(BigInt(0))]
2425
in let 0 = validate (get 0)
2526
[ rowCountNeq 0
2627
] orRaise "MISSING_RELATED_RECORD";

‎query-engine/query-builders/sql-query-builder/src/lib.rs‎

Lines changed: 11 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -151,10 +151,11 @@ impl<'a, V: Visitor<'a>> QueryBuilder for SqlQueryBuilder<'a, V> {
151151
) -> Result<DbQuery, Box<dyn std::error::Error + Send + Sync>> {
152152
use std::slice;
153153

154+
use crate::read::SelectDefinition;
154155
use filter::default_scalar_filter;
155156
use itertools::Itertools;
156-
use quaint::ast::{Aliasable, Joinable, Select};
157-
use select::{JoinConditionExt, SelectBuilderExt};
157+
use quaint::ast::{Aliasable, Joinable};
158+
use select::JoinConditionExt;
158159

159160
let chunkable = Chunkable::from(&query_arguments);
160161
let link_alias = linkage.to_string();
@@ -163,12 +164,6 @@ impl<'a, V: Visitor<'a>> QueryBuilder for SqlQueryBuilder<'a, V> {
163164
let m2m_alias = self.context.next_table_alias();
164165
let m2m_table = rf.as_table(&self.context).alias(m2m_alias.to_string());
165166

166-
let related_alias = self.context.next_table_alias();
167-
let related_table = rf
168-
.related_model()
169-
.as_table(&self.context)
170-
.alias(related_alias.to_string());
171-
172167
let m2m_col = rf
173168
.related_field()
174169
.m2m_column(&self.context)
@@ -199,26 +194,24 @@ impl<'a, V: Visitor<'a>> QueryBuilder for SqlQueryBuilder<'a, V> {
199194

200195
let columns = ModelProjection::from(selected_fields)
201196
.as_columns(&self.context)
202-
.map(|col| col.table(related_alias.to_string()))
197+
.map(|col| col.table(rf.related_model().as_table(&self.context)))
203198
// Add an m2m column with an alias to make it possible to join it outside of this
204199
// function.
205200
.chain([m2m_col.alias(link_alias)]);
206201

207-
let join_condition = rf.m2m_join_conditions(Some(m2m_alias), Some(related_alias), &self.context);
208-
209-
let select = Select::from_table(m2m_table)
210-
.columns(columns)
211-
.inner_join(related_table.on(join_condition))
212-
.with_distinct(&query_arguments, related_alias)
213-
.with_ordering(&query_arguments, Some(related_alias.to_string()), &self.context)
214-
.with_pagination(&query_arguments, None)
215-
.with_filters(query_arguments.filter, Some(related_alias), &self.context);
202+
let join_condition = rf.m2m_join_conditions(Some(m2m_alias), None, &self.context);
216203

204+
let (select, additional_selection_set) =
205+
query_arguments.into_select(&rf.related_model(), selected_fields.virtuals(), &self.context);
206+
let select = select.columns(columns).inner_join(m2m_table.on(join_condition.clone()));
217207
let select = if let Some(filter) = filter {
218208
select.and_where(filter)
219209
} else {
220210
select
221211
};
212+
let select = additional_selection_set
213+
.into_iter()
214+
.fold(select, |acc, val| acc.value(val));
222215

223216
self.convert_query(select, chunkable)
224217
}

0 commit comments

Comments
 (0)