better mapping of subquery aliases

Author: samwillis

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

@@ -42,6 +42,8 @@ export interface CompilationResult {
   sourceWhereClauses: Map<string, BasicExpression<boolean>>
   /** Map of alias to underlying collection id used during compilation */
   aliasToCollectionId: Record<string, string>
+  /** Map of outer alias to inner alias for subquery aliasing (e.g., 'activeUser' → 'user') */
+  aliasRemapping: Record<string, string>
 }
 
 /**
@@ -88,6 +90,11 @@ export function compileQuery(
   // the live layer can subscribe to every alias the optimizer introduces.
   const aliasToCollectionId: Record<string, string> = {}
 
+  // Track alias remapping for subqueries (outer alias → inner alias)
+  // e.g., when .join({ activeUser: subquery }) where subquery uses .from({ user: collection })
+  // we store: aliasRemapping['activeUser'] = 'user'
+  const aliasRemapping: Record<string, string> = {}
+
   // Create a map of source aliases to input streams.
   // Inputs MUST be keyed by alias (e.g., `{ employee: input1, manager: input2 }`),
   // not by collection ID. This enables per-alias subscriptions where different aliases
@@ -109,7 +116,8 @@ export function compileQuery(
     optimizableOrderByCollections,
     cache,
     queryMapping,
-    aliasToCollectionId
+    aliasToCollectionId,
+    aliasRemapping
   )
   sources[mainSource] = mainInput
 
@@ -143,7 +151,8 @@ export function compileQuery(
       optimizableOrderByCollections,
       rawQuery,
       compileQuery,
-      aliasToCollectionId
+      aliasToCollectionId,
+      aliasRemapping
     )
   }
 
@@ -303,6 +312,7 @@ export function compileQuery(
       pipeline: result,
       sourceWhereClauses,
       aliasToCollectionId,
+      aliasRemapping,
     }
     cache.set(rawQuery, compilationResult)
 
@@ -332,6 +342,7 @@ export function compileQuery(
     pipeline: result,
     sourceWhereClauses,
     aliasToCollectionId,
+    aliasRemapping,
   }
   cache.set(rawQuery, compilationResult)
 
@@ -351,7 +362,8 @@ function processFrom(
   optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
   cache: QueryCache,
   queryMapping: QueryMapping,
-  aliasToCollectionId: Record<string, string>
+  aliasToCollectionId: Record<string, string>,
+  aliasRemapping: Record<string, string>
 ): { alias: string; input: KeyedStream; collectionId: string } {
   switch (from.type) {
     case `collectionRef`: {
@@ -383,7 +395,20 @@ function processFrom(
         queryMapping
       )
 
+      // Pull up the inner alias mappings
       Object.assign(aliasToCollectionId, subQueryResult.aliasToCollectionId)
+      Object.assign(aliasRemapping, subQueryResult.aliasRemapping)
+
+      // For subqueries, the outer alias (from.alias) may differ from inner aliases.
+      // Find the inner alias that corresponds to the subquery's main collection and create a remapping.
+      const innerAlias = Object.keys(subQueryResult.aliasToCollectionId).find(
+        (alias) =>
+          subQueryResult.aliasToCollectionId[alias] ===
+          subQueryResult.collectionId
+      )
+      if (innerAlias && innerAlias !== from.alias) {
+        aliasRemapping[from.alias] = innerAlias
+      }
 
       // Extract the pipeline from the compilation result
       const subQueryInput = subQueryResult.pipeline

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

@@ -64,7 +64,8 @@ export function processJoins(
   optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
   rawQuery: QueryIR,
   onCompileSubquery: CompileQueryFn,
-  aliasToCollectionId: Record<string, string>
+  aliasToCollectionId: Record<string, string>,
+  aliasRemapping: Record<string, string>
 ): NamespacedAndKeyedStream {
   let resultPipeline = pipeline
 
@@ -85,7 +86,8 @@ export function processJoins(
       optimizableOrderByCollections,
       rawQuery,
       onCompileSubquery,
-      aliasToCollectionId
+      aliasToCollectionId,
+      aliasRemapping
     )
   }
 
@@ -111,7 +113,8 @@ function processJoin(
   optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
   rawQuery: QueryIR,
   onCompileSubquery: CompileQueryFn,
-  aliasToCollectionId: Record<string, string>
+  aliasToCollectionId: Record<string, string>,
+  aliasRemapping: Record<string, string>
 ): NamespacedAndKeyedStream {
   const isCollectionRef = joinClause.from.type === `collectionRef`
 
@@ -131,7 +134,8 @@ function processJoin(
     cache,
     queryMapping,
     onCompileSubquery,
-    aliasToCollectionId
+    aliasToCollectionId,
+    aliasRemapping
   )
 
   // Add the joined source to the sources map
@@ -270,23 +274,18 @@ function processJoin(
           const lazyAliasCandidate =
             activeSource === `main` ? joinedSource : mainSource
 
-          // The alias candidate might be a subquery alias without a direct subscription.
-          // In that case, find an alias from aliasToCollectionId that maps to the lazy collection.
-          let lazySourceSubscription = subscriptions[lazyAliasCandidate]
-          if (!lazySourceSubscription) {
-            // Search for any alias that maps to the lazy collection ID
-            const matchingAlias = Object.entries(aliasToCollectionId).find(
-              ([_alias, collId]) => collId === lazySource.id
-            )?.[0]
-
-            if (matchingAlias) {
-              lazySourceSubscription = subscriptions[matchingAlias]
-            }
-          }
+          // Find the subscription for lazy loading.
+          // For subqueries, the outer join alias (e.g., 'activeUser') may differ from the
+          // inner alias (e.g., 'user'). Use aliasRemapping to resolve outer → inner alias.
+          // Example: .join({ activeUser: subquery }) where subquery uses .from({ user: collection })
+          // → aliasRemapping['activeUser'] = 'user'
+          const resolvedAlias =
+            aliasRemapping[lazyAliasCandidate] || lazyAliasCandidate
+          const lazySourceSubscription = subscriptions[resolvedAlias]
 
           if (!lazySourceSubscription) {
             throw new Error(
-              `Internal error: subscription for alias '${lazyAliasCandidate}' (collection '${lazySource.id}') is missing in join pipeline. Available aliases: ${Object.keys(subscriptions).join(`, `)}. This indicates a bug in alias tracking.`
+              `Internal error: subscription for alias '${resolvedAlias}' (remapped from '${lazyAliasCandidate}', collection '${lazySource.id}') is missing in join pipeline. Available aliases: ${Object.keys(subscriptions).join(`, `)}. This indicates a bug in alias tracking.`
             )
           }
 
@@ -427,7 +426,8 @@ function processJoinSource(
   cache: QueryCache,
   queryMapping: QueryMapping,
   onCompileSubquery: CompileQueryFn,
-  aliasToCollectionId: Record<string, string>
+  aliasToCollectionId: Record<string, string>,
+  aliasRemapping: Record<string, string>
 ): { alias: string; input: KeyedStream; collectionId: string } {
   switch (from.type) {
     case `collectionRef`: {
@@ -459,9 +459,20 @@ 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.
+      // Pull up the inner alias mappings
       Object.assign(aliasToCollectionId, subQueryResult.aliasToCollectionId)
+      Object.assign(aliasRemapping, subQueryResult.aliasRemapping)
+
+      // For subqueries, the outer alias (from.alias) may differ from inner aliases.
+      // Find the inner alias that corresponds to the subquery's main collection and create a remapping.
+      const innerAlias = Object.keys(subQueryResult.aliasToCollectionId).find(
+        (alias) =>
+          subQueryResult.aliasToCollectionId[alias] ===
+          subQueryResult.collectionId
+      )
+      if (innerAlias && innerAlias !== from.alias) {
+        aliasRemapping[from.alias] = innerAlias
+      }
 
       // Extract the pipeline from the compilation result
       const subQueryInput = subQueryResult.pipeline

packages/db/tests/query/compiler/subqueries.test.ts

@@ -273,9 +273,11 @@ describe(`Query2 Subqueries`, () => {
 
       const usersSubscription = usersCollection.subscribeChanges(() => {})
       const issuesSubscription = issuesCollection.subscribeChanges(() => {})
+
+      // Create subscriptions keyed by alias (matching production behavior)
       const subscriptions: Record<string, CollectionSubscription> = {
-        [usersCollection.id]: usersSubscription,
-        [issuesCollection.id]: issuesSubscription,
+        issue: issuesSubscription,
+        user: usersSubscription,
       }
 
       // Compile and execute the query
@@ -297,17 +299,6 @@ describe(`Query2 Subqueries`, () => {
       )
       const { pipeline } = compilation
 
-      for (const [alias, collectionId] of Object.entries(
-        compilation.aliasToCollectionId
-      )) {
-        if (!subscriptions[alias]) {
-          subscriptions[alias] =
-            collectionId === usersCollection.id
-              ? usersSubscription
-              : issuesSubscription
-        }
-      }
-
       // Since we're doing a left join, the alias on the right (from the subquery) should be handled lazily
       // The subquery uses 'user' alias, but the join uses 'activeUser' - we expect the lazy alias
       // to be the one that's marked (which is 'activeUser' since it's the joinedTableAlias)
@@ -377,9 +368,11 @@ describe(`Query2 Subqueries`, () => {
 
       const usersSubscription = usersCollection.subscribeChanges(() => {})
       const issuesSubscription = issuesCollection.subscribeChanges(() => {})
+
+      // Create subscriptions keyed by alias (matching production behavior)
       const subscriptions: Record<string, CollectionSubscription> = {
-        [usersCollection.id]: usersSubscription,
-        [issuesCollection.id]: issuesSubscription,
+        issue: issuesSubscription,
+        user: usersSubscription,
       }
 
       const dummyCallbacks = {