diff --git a/packages/db/src/query/compiler/index.ts b/packages/db/src/query/compiler/index.ts index 7d5995871..c4a5ae758 100644 --- a/packages/db/src/query/compiler/index.ts +++ b/packages/db/src/query/compiler/index.ts @@ -148,6 +148,7 @@ export function compileQuery( queryMapping, aliasToCollectionId, aliasRemapping, + sourceWhereClauses, ) sources[mainSource] = mainInput @@ -184,6 +185,7 @@ export function compileQuery( compileQuery, aliasToCollectionId, aliasRemapping, + sourceWhereClauses, ) } @@ -466,6 +468,7 @@ function processFrom( queryMapping: QueryMapping, aliasToCollectionId: Record, aliasRemapping: Record, + sourceWhereClauses: Map>, ): { alias: string; input: KeyedStream; collectionId: string } { switch (from.type) { case `collectionRef`: { @@ -504,6 +507,28 @@ function processFrom( Object.assign(aliasToCollectionId, subQueryResult.aliasToCollectionId) Object.assign(aliasRemapping, subQueryResult.aliasRemapping) + // Pull up source WHERE clauses from subquery to parent scope. + // This enables loadSubset to receive the correct where clauses for subquery collections. + // + // IMPORTANT: Skip pull-up for optimizer-created subqueries. These are detected when: + // 1. The outer alias (from.alias) matches the inner alias (from.query.from.alias) + // 2. The subquery was found in queryMapping (it's a user-defined subquery, not optimizer-created) + // + // For optimizer-created subqueries, the parent already has the sourceWhereClauses + // extracted from the original raw query, so pulling up would be redundant. + // More importantly, pulling up for optimizer-created subqueries can cause issues + // when the optimizer has restructured the query. + const isUserDefinedSubquery = queryMapping.has(from.query) + const subqueryFromAlias = from.query.from.alias + const isOptimizerCreated = + !isUserDefinedSubquery && from.alias === subqueryFromAlias + + if (!isOptimizerCreated) { + for (const [alias, whereClause] of subQueryResult.sourceWhereClauses) { + sourceWhereClauses.set(alias, whereClause) + } + } + // Create a FLATTENED remapping from outer alias to innermost alias. // For nested subqueries, this ensures one-hop lookups (not recursive chains). // diff --git a/packages/db/src/query/compiler/joins.ts b/packages/db/src/query/compiler/joins.ts index afa5276b1..5dcb7ed39 100644 --- a/packages/db/src/query/compiler/joins.ts +++ b/packages/db/src/query/compiler/joins.ts @@ -66,6 +66,7 @@ export function processJoins( onCompileSubquery: CompileQueryFn, aliasToCollectionId: Record, aliasRemapping: Record, + sourceWhereClauses: Map>, ): NamespacedAndKeyedStream { let resultPipeline = pipeline @@ -89,6 +90,7 @@ export function processJoins( onCompileSubquery, aliasToCollectionId, aliasRemapping, + sourceWhereClauses, ) } @@ -118,6 +120,7 @@ function processJoin( onCompileSubquery: CompileQueryFn, aliasToCollectionId: Record, aliasRemapping: Record, + sourceWhereClauses: Map>, ): NamespacedAndKeyedStream { const isCollectionRef = joinClause.from.type === `collectionRef` @@ -140,6 +143,7 @@ function processJoin( onCompileSubquery, aliasToCollectionId, aliasRemapping, + sourceWhereClauses, ) // Add the joined source to the sources map @@ -431,6 +435,7 @@ function processJoinSource( onCompileSubquery: CompileQueryFn, aliasToCollectionId: Record, aliasRemapping: Record, + sourceWhereClauses: Map>, ): { alias: string; input: KeyedStream; collectionId: string } { switch (from.type) { case `collectionRef`: { @@ -469,6 +474,26 @@ function processJoinSource( Object.assign(aliasToCollectionId, subQueryResult.aliasToCollectionId) Object.assign(aliasRemapping, subQueryResult.aliasRemapping) + // Pull up source WHERE clauses from subquery to parent scope. + // This enables loadSubset to receive the correct where clauses for subquery collections. + // + // IMPORTANT: Skip pull-up for optimizer-created subqueries. These are detected when: + // 1. The outer alias (from.alias) matches the inner alias (from.query.from.alias) + // 2. The subquery was found in queryMapping (it's a user-defined subquery, not optimizer-created) + // + // For optimizer-created subqueries, the parent already has the sourceWhereClauses + // extracted from the original raw query, so pulling up would be redundant. + const isUserDefinedSubquery = queryMapping.has(from.query) + const fromInnerAlias = from.query.from.alias + const isOptimizerCreated = + !isUserDefinedSubquery && from.alias === fromInnerAlias + + if (!isOptimizerCreated) { + for (const [alias, whereClause] of subQueryResult.sourceWhereClauses) { + sourceWhereClauses.set(alias, whereClause) + } + } + // Create a flattened remapping from outer alias to innermost alias. // For nested subqueries, this ensures one-hop lookups (not recursive chains). // diff --git a/packages/db/tests/query/load-subset-subquery.test.ts b/packages/db/tests/query/load-subset-subquery.test.ts new file mode 100644 index 000000000..27ce7207c --- /dev/null +++ b/packages/db/tests/query/load-subset-subquery.test.ts @@ -0,0 +1,268 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { createCollection } from '../../src/collection/index.js' +import { + and, + createLiveQueryCollection, + eq, + gte, +} from '../../src/query/index.js' +import { PropRef, Value } from '../../src/query/ir.js' +import type { Collection } from '../../src/collection/index.js' +import type { + LoadSubsetOptions, + NonSingleResult, + UtilsRecord, +} from '../../src/types.js' +import type { OrderBy } from '../../src/query/ir.js' + +// Sample types for testing +type Order = { + id: number + scheduled_at: string + status: string + address_id: number +} + +type Charge = { + id: number + address_id: number + amount: number +} + +// Sample data +const sampleOrders: Array = [ + { + id: 1, + scheduled_at: `2024-01-15`, + status: `queued`, + address_id: 1, + }, + { + id: 2, + scheduled_at: `2024-01-10`, + status: `queued`, + address_id: 2, + }, + { + id: 3, + scheduled_at: `2024-01-20`, + status: `completed`, + address_id: 1, + }, +] + +const sampleCharges: Array = [ + { id: 1, address_id: 1, amount: 100 }, + { id: 2, address_id: 2, amount: 200 }, +] + +type ChargersCollection = Collection< + Charge, + string | number, + UtilsRecord, + never, + Charge +> & + NonSingleResult + +type OrdersCollection = Collection< + Order, + string | number, + UtilsRecord, + never, + Order +> & + NonSingleResult + +describe(`loadSubset with subqueries`, () => { + let chargesCollection: ChargersCollection + + beforeEach(() => { + // Create charges collection + chargesCollection = createCollection({ + id: `charges`, + getKey: (charge) => charge.id, + sync: { + sync: ({ begin, write, commit, markReady }) => { + begin() + for (const charge of sampleCharges) { + write({ type: `insert`, value: charge }) + } + commit() + markReady() + }, + }, + }) + }) + + function createOrdersCollectionWithTracking(): { + collection: OrdersCollection + loadSubsetCalls: Array + } { + const loadSubsetCalls: Array = [] + + const collection = createCollection({ + id: `orders`, + getKey: (order) => order.id, + syncMode: `on-demand`, + sync: { + sync: ({ begin, write, commit, markReady }) => { + begin() + for (const order of sampleOrders) { + write({ type: `insert`, value: order }) + } + commit() + markReady() + return { + loadSubset: vi.fn((options: LoadSubsetOptions) => { + loadSubsetCalls.push(options) + return Promise.resolve() + }), + } + }, + }, + }) + + return { collection, loadSubsetCalls } + } + + it(`should call loadSubset with where clause for direct query`, async () => { + const today = `2024-01-12` + const { collection: ordersCollection, loadSubsetCalls } = + createOrdersCollectionWithTracking() + + const directQuery = createLiveQueryCollection((q) => + q + .from({ order: ordersCollection }) + .where(({ order }) => gte(order.scheduled_at, today)) + .where(({ order }) => eq(order.status, `queued`)), + ) + + await directQuery.preload() + + // Verify loadSubset was called + expect(loadSubsetCalls.length).toBeGreaterThan(0) + + // Verify the last call (or any call) has the where clause + const lastCall = loadSubsetCalls[loadSubsetCalls.length - 1] + expect(lastCall).toBeDefined() + expect(lastCall!.where).toBeDefined() + + const expectedWhereClause = and( + gte(new PropRef([`scheduled_at`]), new Value(today)), + eq(new PropRef([`status`]), new Value(`queued`)), + ) + + expect(lastCall!.where).toEqual(expectedWhereClause) + }) + + it(`should call loadSubset with where clause for subquery`, async () => { + const today = `2024-01-12` + const { collection: ordersCollection, loadSubsetCalls } = + createOrdersCollectionWithTracking() + + const subqueryQuery = createLiveQueryCollection((q) => { + // Build subquery with filters + const prepaidOrderQ = q + .from({ prepaidOrder: ordersCollection }) + .where(({ prepaidOrder }) => gte(prepaidOrder.scheduled_at, today)) + .where(({ prepaidOrder }) => eq(prepaidOrder.status, `queued`)) + + // Use subquery in main query + return q + .from({ charge: chargesCollection }) + .fullJoin({ prepaidOrder: prepaidOrderQ }, ({ charge, prepaidOrder }) => + eq(charge.address_id, prepaidOrder.address_id), + ) + }) + + await subqueryQuery.preload() + + // Verify loadSubset was called for the orders collection + expect(loadSubsetCalls.length).toBeGreaterThan(0) + + // Verify the last call (or any call) has the where clause + const lastCall = loadSubsetCalls[loadSubsetCalls.length - 1] + expect(lastCall).toBeDefined() + expect(lastCall!.where).toBeDefined() + + const expectedWhereClause = and( + gte(new PropRef([`scheduled_at`]), new Value(today)), + eq(new PropRef([`status`]), new Value(`queued`)), + ) + + expect(lastCall!.where).toEqual(expectedWhereClause) + }) + + it(`should call loadSubset with orderBy clause for direct query`, async () => { + const { collection: ordersCollection, loadSubsetCalls } = + createOrdersCollectionWithTracking() + + const directQuery = createLiveQueryCollection((q) => + q + .from({ order: ordersCollection }) + .orderBy(({ order }) => order.scheduled_at, `desc`) + .limit(2), + ) + + await directQuery.preload() + + // Verify loadSubset was called + expect(loadSubsetCalls.length).toBeGreaterThan(0) + + // Verify the last call has the orderBy clause and limit + const lastCall = loadSubsetCalls[loadSubsetCalls.length - 1] + expect(lastCall).toBeDefined() + expect(lastCall!.orderBy).toBeDefined() + expect(lastCall!.limit).toBe(2) + + const expectedOrderBy: OrderBy = [ + { + expression: new PropRef([`scheduled_at`]), + compareOptions: { direction: `desc`, nulls: `first` }, + }, + ] + + expect(lastCall!.orderBy).toEqual(expectedOrderBy) + }) + + it(`should call loadSubset with orderBy clause for subquery`, async () => { + const { collection: ordersCollection, loadSubsetCalls } = + createOrdersCollectionWithTracking() + + const subqueryQuery = createLiveQueryCollection((q) => { + // Build subquery with orderBy and limit + const prepaidOrderQ = q + .from({ prepaidOrder: ordersCollection }) + .orderBy(({ prepaidOrder }) => prepaidOrder.scheduled_at, `desc`) + .limit(2) + + // Use subquery in main query + return q + .from({ charge: chargesCollection }) + .fullJoin({ prepaidOrder: prepaidOrderQ }, ({ charge, prepaidOrder }) => + eq(charge.address_id, prepaidOrder.address_id), + ) + }) + + await subqueryQuery.preload() + + // Verify loadSubset was called for the orders collection + expect(loadSubsetCalls.length).toBeGreaterThan(0) + + // Verify the last call has the orderBy clause and limit + const lastCall = loadSubsetCalls[loadSubsetCalls.length - 1] + expect(lastCall).toBeDefined() + expect(lastCall!.orderBy).toBeDefined() + expect(lastCall!.limit).toBe(2) + + const expectedOrderBy: OrderBy = [ + { + expression: new PropRef([`scheduled_at`]), + compareOptions: { direction: `desc`, nulls: `first` }, + }, + ] + + expect(lastCall!.orderBy).toEqual(expectedOrderBy) + }) +})