change to subscription per source alias, rather than per colleciton

Author: samwillis

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

@@ -40,6 +40,8 @@ export interface CompilationResult {
   pipeline: ResultStream
   /** Map of collection aliases to their WHERE clauses for index optimization */
   collectionWhereClauses: Map<string, BasicExpression<boolean>>
+  /** Map of alias to underlying collection id used during compilation */
+  aliasToCollectionId: Record<string, string>
 }
 
 /**
@@ -78,6 +80,10 @@ export function compileQuery(
   // Create a copy of the inputs map to avoid modifying the original
   const allInputs = { ...inputs }
 
+  // Track alias to collection id relationships discovered during compilation so
+  // the live layer can subscribe to every alias the optimiser introduces.
+  const aliasToCollectionId: Record<string, string> = {}
+
   // Create a map of table aliases to inputs
   const tables: Record<string, KeyedStream> = {}
 
@@ -95,7 +101,8 @@ export function compileQuery(
     lazyCollections,
     optimizableOrderByCollections,
     cache,
-    queryMapping
+    queryMapping,
+    aliasToCollectionId
   )
   tables[mainTableAlias] = mainInput
 
@@ -128,7 +135,8 @@ export function compileQuery(
       lazyCollections,
       optimizableOrderByCollections,
       rawQuery,
-      compileQuery
+      compileQuery,
+      aliasToCollectionId
     )
   }
 
@@ -287,6 +295,7 @@ export function compileQuery(
       collectionId: mainCollectionId,
       pipeline: result,
       collectionWhereClauses,
+      aliasToCollectionId,
     }
     cache.set(rawQuery, compilationResult)
 
@@ -315,6 +324,7 @@ export function compileQuery(
     collectionId: mainCollectionId,
     pipeline: result,
     collectionWhereClauses,
+    aliasToCollectionId,
   }
   cache.set(rawQuery, compilationResult)
 
@@ -333,14 +343,16 @@ function processFrom(
   lazyCollections: Set<string>,
   optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
   cache: QueryCache,
-  queryMapping: QueryMapping
+  queryMapping: QueryMapping,
+  aliasToCollectionId: Record<string, string>
 ): { alias: string; input: KeyedStream; collectionId: string } {
   switch (from.type) {
     case `collectionRef`: {
       const input = allInputs[from.alias] ?? allInputs[from.collection.id]
       if (!input) {
         throw new CollectionInputNotFoundError(from.alias, from.collection.id)
       }
+      aliasToCollectionId[from.alias] = from.collection.id
       return { alias: from.alias, input, collectionId: from.collection.id }
     }
     case `queryRef`: {
@@ -360,6 +372,8 @@ function processFrom(
         queryMapping
       )
 
+      Object.assign(aliasToCollectionId, subQueryResult.aliasToCollectionId)
+
       // Extract the pipeline from the compilation result
       const subQueryInput = subQueryResult.pipeline
 

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

@@ -63,7 +63,8 @@ export function processJoins(
   lazyCollections: Set<string>,
   optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
   rawQuery: QueryIR,
-  onCompileSubquery: CompileQueryFn
+  onCompileSubquery: CompileQueryFn,
+  aliasToCollectionId: Record<string, string>
 ): NamespacedAndKeyedStream {
   let resultPipeline = pipeline
 
@@ -83,7 +84,8 @@ export function processJoins(
       lazyCollections,
       optimizableOrderByCollections,
       rawQuery,
-      onCompileSubquery
+      onCompileSubquery,
+      aliasToCollectionId
     )
   }
 
@@ -108,8 +110,11 @@ function processJoin(
   lazyCollections: Set<string>,
   optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
   rawQuery: QueryIR,
-  onCompileSubquery: CompileQueryFn
+  onCompileSubquery: CompileQueryFn,
+  aliasToCollectionId: Record<string, string>
 ): NamespacedAndKeyedStream {
+  const isCollectionRef = joinClause.from.type === `collectionRef`
+
   // Get the joined table alias and input stream
   const {
     alias: joinedTableAlias,
@@ -125,11 +130,17 @@ function processJoin(
     optimizableOrderByCollections,
     cache,
     queryMapping,
-    onCompileSubquery
+    onCompileSubquery,
+    aliasToCollectionId
   )
 
   // Add the joined table to the tables map
   tables[joinedTableAlias] = joinedInput
+  if (isCollectionRef) {
+    // Only direct collection references form new alias bindings. Subquery
+    // aliases reuse the mapping returned from the recursive compilation above.
+    aliasToCollectionId[joinedTableAlias] = joinedCollectionId
+  }
 
   const mainCollection = collections[mainTableId]
   const joinedCollection = collections[joinedCollectionId]
@@ -260,9 +271,9 @@ function processJoin(
             subscriptions[`__collection:${lazyCollection.id}`]
 
           if (!lazyCollectionSubscription) {
-            // The alias was not subscribed (e.g. belongs to a nested subquery),
-            // so we skip the lazy loading optimization for this join.
-            return
+            throw new Error(
+              `Internal error: subscription for collection is missing in join pipeline. Make sure the live query collection sets the subscription before running the pipeline.`
+            )
           }
 
           if (lazyCollectionSubscription.hasLoadedInitialState()) {
@@ -401,14 +412,16 @@ function processJoinSource(
   optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
   cache: QueryCache,
   queryMapping: QueryMapping,
-  onCompileSubquery: CompileQueryFn
+  onCompileSubquery: CompileQueryFn,
+  aliasToCollectionId: Record<string, string>
 ): { alias: string; input: KeyedStream; collectionId: string } {
   switch (from.type) {
     case `collectionRef`: {
       const input = allInputs[from.alias] ?? allInputs[from.collection.id]
       if (!input) {
         throw new CollectionInputNotFoundError(from.alias, from.collection.id)
       }
+      aliasToCollectionId[from.alias] = from.collection.id
       return { alias: from.alias, input, collectionId: from.collection.id }
     }
     case `queryRef`: {
@@ -428,6 +441,10 @@ function processJoinSource(
         queryMapping
       )
 
+      // Pull the nested alias map up so the caller can subscribe to those aliases
+      // and keep the current alias pointing at the subquery's collection.
+      Object.assign(aliasToCollectionId, subQueryResult.aliasToCollectionId)
+
       // Extract the pipeline from the compilation result
       const subQueryInput = subQueryResult.pipeline
 

packages/db/src/query/live/collection-config-builder.ts

@@ -32,8 +32,9 @@ export class CollectionConfigBuilder<
   private readonly id: string
   readonly query: QueryIR
   private readonly collections: Record<string, Collection<any, any, any>>
-  private readonly aliasToCollectionId: Record<string, string>
   private readonly collectionByAlias: Record<string, Collection<any, any, any>>
+  // Populated during compilation to include optimizer-generated aliases
+  private compiledAliasToCollectionId: Record<string, string> = {}
 
   // WeakMap to store the keys of the results
   // so that we can retrieve them in the getKey function
@@ -72,13 +73,11 @@ export class CollectionConfigBuilder<
     this.collections = extractCollectionsFromQuery(this.query)
     const collectionAliasesById = extractCollectionAliases(this.query)
 
-    this.aliasToCollectionId = {}
     this.collectionByAlias = {}
     for (const [collectionId, aliases] of collectionAliasesById.entries()) {
       const collection = this.collections[collectionId]
       if (!collection) continue
       for (const alias of aliases) {
-        this.aliasToCollectionId[alias] = collectionId
         this.collectionByAlias[alias] = collection
       }
     }
@@ -111,11 +110,15 @@ export class CollectionConfigBuilder<
   }
 
   getCollectionIdForAlias(alias: string): string {
-    const collectionId = this.aliasToCollectionId[alias]
-    if (!collectionId) {
-      throw new Error(`Unknown collection alias "${alias}"`)
+    const compiled = this.compiledAliasToCollectionId[alias]
+    if (compiled) {
+      return compiled
     }
-    return collectionId
+    const collection = this.collectionByAlias[alias]
+    if (collection) {
+      return collection.id
+    }
+    throw new Error(`Unknown collection alias "${alias}"`)
   }
 
   // The callback function is called after the graph has run.
@@ -224,11 +227,8 @@ export class CollectionConfigBuilder<
       ])
     )
 
-    // Compile the query and get both pipeline and collection WHERE clauses
-    const {
-      pipeline: pipelineCache,
-      collectionWhereClauses: collectionWhereClausesCache,
-    } = compileQuery(
+    // Compile the query and capture alias metadata produced during optimisation
+    let compilation = compileQuery(
       this.query,
       this.inputsCache as Record<string, KeyedStream>,
       this.collections,
@@ -238,8 +238,37 @@ export class CollectionConfigBuilder<
       this.optimizableOrderByCollections
     )
 
-    this.pipelineCache = pipelineCache
-    this.collectionWhereClausesCache = collectionWhereClausesCache
+    this.pipelineCache = compilation.pipeline
+    this.collectionWhereClausesCache = compilation.collectionWhereClauses
+    this.compiledAliasToCollectionId = compilation.aliasToCollectionId
+    // Optimized queries can introduce aliases beyond those declared on the
+    // builder. If that happens, provision inputs for the missing aliases and
+    // recompile so the pipeline is fully wired before execution.
+    const missingAliases = Object.keys(this.compiledAliasToCollectionId).filter(
+      (alias) => !Object.hasOwn(this.inputsCache!, alias)
+    )
+
+    if (missingAliases.length > 0) {
+      for (const alias of missingAliases) {
+        this.inputsCache[alias] = this.graphCache.newInput<any>()
+      }
+
+      compilation = compileQuery(
+        this.query,
+        this.inputsCache as Record<string, KeyedStream>,
+        this.collections,
+        this.subscriptions,
+        this.lazyCollectionsCallbacks,
+        this.lazyCollections,
+        this.optimizableOrderByCollections,
+        new WeakMap(),
+        new WeakMap()
+      )
+
+      this.pipelineCache = compilation.pipeline
+      this.collectionWhereClausesCache = compilation.collectionWhereClauses
+      this.compiledAliasToCollectionId = compilation.aliasToCollectionId
+    }
   }
 
   private maybeCompileBasePipeline() {
@@ -355,31 +384,41 @@ export class CollectionConfigBuilder<
     config: Parameters<SyncConfig<TResult>[`sync`]>[0],
     syncState: FullSyncState
   ) {
-    const loaders = Object.entries(this.collectionByAlias).map(
-      ([alias, collection]) => {
-        const collectionId = this.aliasToCollectionId[alias]!
-        const collectionSubscriber = new CollectionSubscriber(
-          alias,
-          collectionId,
-          collection,
-          config,
-          syncState,
-          this
-        )
-
-        const subscription = collectionSubscriber.subscribe()
-        this.subscriptions[alias] = subscription
-        const collectionKey = `__collection:${collectionId}`
-        this.subscriptions[collectionKey] = subscription
-
-        const loadMore = collectionSubscriber.loadMoreIfNeeded.bind(
-          collectionSubscriber,
-          subscription
-        )
-
-        return loadMore
-      }
-    )
+    const compiledAliases = Object.entries(this.compiledAliasToCollectionId)
+    if (compiledAliases.length === 0) {
+      throw new Error(
+        `Compiler returned no alias metadata for query '${this.id}'. This should not happen; please report.`
+      )
+    }
+
+    // Subscribe to each alias the compiler reported.
+    const aliasEntries = compiledAliases
+
+    const loaders = aliasEntries.map(([alias, collectionId]) => {
+      const collection =
+        this.collectionByAlias[alias] ?? this.collections[collectionId]!
+
+      const collectionSubscriber = new CollectionSubscriber(
+        alias,
+        collectionId,
+        collection,
+        config,
+        syncState,
+        this
+      )
+
+      const subscription = collectionSubscriber.subscribe()
+      this.subscriptions[alias] = subscription
+      const collectionKey = `__collection:${collectionId}`
+      this.subscriptions[collectionKey] = subscription
+
+      const loadMore = collectionSubscriber.loadMoreIfNeeded.bind(
+        collectionSubscriber,
+        subscription
+      )
+
+      return loadMore
+    })
 
     const loadMoreDataCallback = () => {
       loaders.map((loader) => loader())