Skip to content

Commit 05130f2

Browse files
Fix where clause missing in loadSubset on subquery collection (#1097)
* Unit test to show that where clause is missing in loadSubset on subquery * ci: apply automated fixes * Pull up where clauses of subqueries into sourceWhereClauses * Unit tests for orderBy information in loadSubset * Fix where clause pull up * ci: apply automated fixes * Changeset --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
1 parent 5c717e6 commit 05130f2

File tree

4 files changed

+323
-0
lines changed

4 files changed

+323
-0
lines changed

.changeset/open-hairs-crash.md

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'@tanstack/db': patch
3+
---
4+
5+
Fix bug that caused the WHERE clause of a subquery not to be passed to the `loadSubset` function

packages/db/src/query/compiler/index.ts

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,7 @@ export function compileQuery(
148148
queryMapping,
149149
aliasToCollectionId,
150150
aliasRemapping,
151+
sourceWhereClauses,
151152
)
152153
sources[mainSource] = mainInput
153154

@@ -184,6 +185,7 @@ export function compileQuery(
184185
compileQuery,
185186
aliasToCollectionId,
186187
aliasRemapping,
188+
sourceWhereClauses,
187189
)
188190
}
189191

@@ -466,6 +468,7 @@ function processFrom(
466468
queryMapping: QueryMapping,
467469
aliasToCollectionId: Record<string, string>,
468470
aliasRemapping: Record<string, string>,
471+
sourceWhereClauses: Map<string, BasicExpression<boolean>>,
469472
): { alias: string; input: KeyedStream; collectionId: string } {
470473
switch (from.type) {
471474
case `collectionRef`: {
@@ -504,6 +507,28 @@ function processFrom(
504507
Object.assign(aliasToCollectionId, subQueryResult.aliasToCollectionId)
505508
Object.assign(aliasRemapping, subQueryResult.aliasRemapping)
506509

510+
// Pull up source WHERE clauses from subquery to parent scope.
511+
// This enables loadSubset to receive the correct where clauses for subquery collections.
512+
//
513+
// IMPORTANT: Skip pull-up for optimizer-created subqueries. These are detected when:
514+
// 1. The outer alias (from.alias) matches the inner alias (from.query.from.alias)
515+
// 2. The subquery was found in queryMapping (it's a user-defined subquery, not optimizer-created)
516+
//
517+
// For optimizer-created subqueries, the parent already has the sourceWhereClauses
518+
// extracted from the original raw query, so pulling up would be redundant.
519+
// More importantly, pulling up for optimizer-created subqueries can cause issues
520+
// when the optimizer has restructured the query.
521+
const isUserDefinedSubquery = queryMapping.has(from.query)
522+
const subqueryFromAlias = from.query.from.alias
523+
const isOptimizerCreated =
524+
!isUserDefinedSubquery && from.alias === subqueryFromAlias
525+
526+
if (!isOptimizerCreated) {
527+
for (const [alias, whereClause] of subQueryResult.sourceWhereClauses) {
528+
sourceWhereClauses.set(alias, whereClause)
529+
}
530+
}
531+
507532
// Create a FLATTENED remapping from outer alias to innermost alias.
508533
// For nested subqueries, this ensures one-hop lookups (not recursive chains).
509534
//

packages/db/src/query/compiler/joins.ts

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ export function processJoins(
6666
onCompileSubquery: CompileQueryFn,
6767
aliasToCollectionId: Record<string, string>,
6868
aliasRemapping: Record<string, string>,
69+
sourceWhereClauses: Map<string, BasicExpression<boolean>>,
6970
): NamespacedAndKeyedStream {
7071
let resultPipeline = pipeline
7172

@@ -89,6 +90,7 @@ export function processJoins(
8990
onCompileSubquery,
9091
aliasToCollectionId,
9192
aliasRemapping,
93+
sourceWhereClauses,
9294
)
9395
}
9496

@@ -118,6 +120,7 @@ function processJoin(
118120
onCompileSubquery: CompileQueryFn,
119121
aliasToCollectionId: Record<string, string>,
120122
aliasRemapping: Record<string, string>,
123+
sourceWhereClauses: Map<string, BasicExpression<boolean>>,
121124
): NamespacedAndKeyedStream {
122125
const isCollectionRef = joinClause.from.type === `collectionRef`
123126

@@ -140,6 +143,7 @@ function processJoin(
140143
onCompileSubquery,
141144
aliasToCollectionId,
142145
aliasRemapping,
146+
sourceWhereClauses,
143147
)
144148

145149
// Add the joined source to the sources map
@@ -431,6 +435,7 @@ function processJoinSource(
431435
onCompileSubquery: CompileQueryFn,
432436
aliasToCollectionId: Record<string, string>,
433437
aliasRemapping: Record<string, string>,
438+
sourceWhereClauses: Map<string, BasicExpression<boolean>>,
434439
): { alias: string; input: KeyedStream; collectionId: string } {
435440
switch (from.type) {
436441
case `collectionRef`: {
@@ -469,6 +474,26 @@ function processJoinSource(
469474
Object.assign(aliasToCollectionId, subQueryResult.aliasToCollectionId)
470475
Object.assign(aliasRemapping, subQueryResult.aliasRemapping)
471476

477+
// Pull up source WHERE clauses from subquery to parent scope.
478+
// This enables loadSubset to receive the correct where clauses for subquery collections.
479+
//
480+
// IMPORTANT: Skip pull-up for optimizer-created subqueries. These are detected when:
481+
// 1. The outer alias (from.alias) matches the inner alias (from.query.from.alias)
482+
// 2. The subquery was found in queryMapping (it's a user-defined subquery, not optimizer-created)
483+
//
484+
// For optimizer-created subqueries, the parent already has the sourceWhereClauses
485+
// extracted from the original raw query, so pulling up would be redundant.
486+
const isUserDefinedSubquery = queryMapping.has(from.query)
487+
const fromInnerAlias = from.query.from.alias
488+
const isOptimizerCreated =
489+
!isUserDefinedSubquery && from.alias === fromInnerAlias
490+
491+
if (!isOptimizerCreated) {
492+
for (const [alias, whereClause] of subQueryResult.sourceWhereClauses) {
493+
sourceWhereClauses.set(alias, whereClause)
494+
}
495+
}
496+
472497
// Create a flattened remapping from outer alias to innermost alias.
473498
// For nested subqueries, this ensures one-hop lookups (not recursive chains).
474499
//
Lines changed: 268 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,268 @@
1+
import { beforeEach, describe, expect, it, vi } from 'vitest'
2+
import { createCollection } from '../../src/collection/index.js'
3+
import {
4+
and,
5+
createLiveQueryCollection,
6+
eq,
7+
gte,
8+
} from '../../src/query/index.js'
9+
import { PropRef, Value } from '../../src/query/ir.js'
10+
import type { Collection } from '../../src/collection/index.js'
11+
import type {
12+
LoadSubsetOptions,
13+
NonSingleResult,
14+
UtilsRecord,
15+
} from '../../src/types.js'
16+
import type { OrderBy } from '../../src/query/ir.js'
17+
18+
// Sample types for testing
19+
type Order = {
20+
id: number
21+
scheduled_at: string
22+
status: string
23+
address_id: number
24+
}
25+
26+
type Charge = {
27+
id: number
28+
address_id: number
29+
amount: number
30+
}
31+
32+
// Sample data
33+
const sampleOrders: Array<Order> = [
34+
{
35+
id: 1,
36+
scheduled_at: `2024-01-15`,
37+
status: `queued`,
38+
address_id: 1,
39+
},
40+
{
41+
id: 2,
42+
scheduled_at: `2024-01-10`,
43+
status: `queued`,
44+
address_id: 2,
45+
},
46+
{
47+
id: 3,
48+
scheduled_at: `2024-01-20`,
49+
status: `completed`,
50+
address_id: 1,
51+
},
52+
]
53+
54+
const sampleCharges: Array<Charge> = [
55+
{ id: 1, address_id: 1, amount: 100 },
56+
{ id: 2, address_id: 2, amount: 200 },
57+
]
58+
59+
type ChargersCollection = Collection<
60+
Charge,
61+
string | number,
62+
UtilsRecord,
63+
never,
64+
Charge
65+
> &
66+
NonSingleResult
67+
68+
type OrdersCollection = Collection<
69+
Order,
70+
string | number,
71+
UtilsRecord,
72+
never,
73+
Order
74+
> &
75+
NonSingleResult
76+
77+
describe(`loadSubset with subqueries`, () => {
78+
let chargesCollection: ChargersCollection
79+
80+
beforeEach(() => {
81+
// Create charges collection
82+
chargesCollection = createCollection<Charge>({
83+
id: `charges`,
84+
getKey: (charge) => charge.id,
85+
sync: {
86+
sync: ({ begin, write, commit, markReady }) => {
87+
begin()
88+
for (const charge of sampleCharges) {
89+
write({ type: `insert`, value: charge })
90+
}
91+
commit()
92+
markReady()
93+
},
94+
},
95+
})
96+
})
97+
98+
function createOrdersCollectionWithTracking(): {
99+
collection: OrdersCollection
100+
loadSubsetCalls: Array<LoadSubsetOptions>
101+
} {
102+
const loadSubsetCalls: Array<LoadSubsetOptions> = []
103+
104+
const collection = createCollection<Order>({
105+
id: `orders`,
106+
getKey: (order) => order.id,
107+
syncMode: `on-demand`,
108+
sync: {
109+
sync: ({ begin, write, commit, markReady }) => {
110+
begin()
111+
for (const order of sampleOrders) {
112+
write({ type: `insert`, value: order })
113+
}
114+
commit()
115+
markReady()
116+
return {
117+
loadSubset: vi.fn((options: LoadSubsetOptions) => {
118+
loadSubsetCalls.push(options)
119+
return Promise.resolve()
120+
}),
121+
}
122+
},
123+
},
124+
})
125+
126+
return { collection, loadSubsetCalls }
127+
}
128+
129+
it(`should call loadSubset with where clause for direct query`, async () => {
130+
const today = `2024-01-12`
131+
const { collection: ordersCollection, loadSubsetCalls } =
132+
createOrdersCollectionWithTracking()
133+
134+
const directQuery = createLiveQueryCollection((q) =>
135+
q
136+
.from({ order: ordersCollection })
137+
.where(({ order }) => gte(order.scheduled_at, today))
138+
.where(({ order }) => eq(order.status, `queued`)),
139+
)
140+
141+
await directQuery.preload()
142+
143+
// Verify loadSubset was called
144+
expect(loadSubsetCalls.length).toBeGreaterThan(0)
145+
146+
// Verify the last call (or any call) has the where clause
147+
const lastCall = loadSubsetCalls[loadSubsetCalls.length - 1]
148+
expect(lastCall).toBeDefined()
149+
expect(lastCall!.where).toBeDefined()
150+
151+
const expectedWhereClause = and(
152+
gte(new PropRef([`scheduled_at`]), new Value(today)),
153+
eq(new PropRef([`status`]), new Value(`queued`)),
154+
)
155+
156+
expect(lastCall!.where).toEqual(expectedWhereClause)
157+
})
158+
159+
it(`should call loadSubset with where clause for subquery`, async () => {
160+
const today = `2024-01-12`
161+
const { collection: ordersCollection, loadSubsetCalls } =
162+
createOrdersCollectionWithTracking()
163+
164+
const subqueryQuery = createLiveQueryCollection((q) => {
165+
// Build subquery with filters
166+
const prepaidOrderQ = q
167+
.from({ prepaidOrder: ordersCollection })
168+
.where(({ prepaidOrder }) => gte(prepaidOrder.scheduled_at, today))
169+
.where(({ prepaidOrder }) => eq(prepaidOrder.status, `queued`))
170+
171+
// Use subquery in main query
172+
return q
173+
.from({ charge: chargesCollection })
174+
.fullJoin({ prepaidOrder: prepaidOrderQ }, ({ charge, prepaidOrder }) =>
175+
eq(charge.address_id, prepaidOrder.address_id),
176+
)
177+
})
178+
179+
await subqueryQuery.preload()
180+
181+
// Verify loadSubset was called for the orders collection
182+
expect(loadSubsetCalls.length).toBeGreaterThan(0)
183+
184+
// Verify the last call (or any call) has the where clause
185+
const lastCall = loadSubsetCalls[loadSubsetCalls.length - 1]
186+
expect(lastCall).toBeDefined()
187+
expect(lastCall!.where).toBeDefined()
188+
189+
const expectedWhereClause = and(
190+
gte(new PropRef([`scheduled_at`]), new Value(today)),
191+
eq(new PropRef([`status`]), new Value(`queued`)),
192+
)
193+
194+
expect(lastCall!.where).toEqual(expectedWhereClause)
195+
})
196+
197+
it(`should call loadSubset with orderBy clause for direct query`, async () => {
198+
const { collection: ordersCollection, loadSubsetCalls } =
199+
createOrdersCollectionWithTracking()
200+
201+
const directQuery = createLiveQueryCollection((q) =>
202+
q
203+
.from({ order: ordersCollection })
204+
.orderBy(({ order }) => order.scheduled_at, `desc`)
205+
.limit(2),
206+
)
207+
208+
await directQuery.preload()
209+
210+
// Verify loadSubset was called
211+
expect(loadSubsetCalls.length).toBeGreaterThan(0)
212+
213+
// Verify the last call has the orderBy clause and limit
214+
const lastCall = loadSubsetCalls[loadSubsetCalls.length - 1]
215+
expect(lastCall).toBeDefined()
216+
expect(lastCall!.orderBy).toBeDefined()
217+
expect(lastCall!.limit).toBe(2)
218+
219+
const expectedOrderBy: OrderBy = [
220+
{
221+
expression: new PropRef([`scheduled_at`]),
222+
compareOptions: { direction: `desc`, nulls: `first` },
223+
},
224+
]
225+
226+
expect(lastCall!.orderBy).toEqual(expectedOrderBy)
227+
})
228+
229+
it(`should call loadSubset with orderBy clause for subquery`, async () => {
230+
const { collection: ordersCollection, loadSubsetCalls } =
231+
createOrdersCollectionWithTracking()
232+
233+
const subqueryQuery = createLiveQueryCollection((q) => {
234+
// Build subquery with orderBy and limit
235+
const prepaidOrderQ = q
236+
.from({ prepaidOrder: ordersCollection })
237+
.orderBy(({ prepaidOrder }) => prepaidOrder.scheduled_at, `desc`)
238+
.limit(2)
239+
240+
// Use subquery in main query
241+
return q
242+
.from({ charge: chargesCollection })
243+
.fullJoin({ prepaidOrder: prepaidOrderQ }, ({ charge, prepaidOrder }) =>
244+
eq(charge.address_id, prepaidOrder.address_id),
245+
)
246+
})
247+
248+
await subqueryQuery.preload()
249+
250+
// Verify loadSubset was called for the orders collection
251+
expect(loadSubsetCalls.length).toBeGreaterThan(0)
252+
253+
// Verify the last call has the orderBy clause and limit
254+
const lastCall = loadSubsetCalls[loadSubsetCalls.length - 1]
255+
expect(lastCall).toBeDefined()
256+
expect(lastCall!.orderBy).toBeDefined()
257+
expect(lastCall!.limit).toBe(2)
258+
259+
const expectedOrderBy: OrderBy = [
260+
{
261+
expression: new PropRef([`scheduled_at`]),
262+
compareOptions: { direction: `desc`, nulls: `first` },
263+
},
264+
]
265+
266+
expect(lastCall!.orderBy).toEqual(expectedOrderBy)
267+
})
268+
})

0 commit comments

Comments
 (0)