Fixes to subscription

Author: kevin-dp

packages/db/src/collection/subscription.ts

@@ -1,11 +1,11 @@
 import { ensureIndexForExpression } from "../indexes/auto-index.js"
-import { and } from "../query/builder/functions.js"
+import { and, gt, lt } from "../query/index.js"
 import {
   createFilterFunctionFromExpression,
   createFilteredCallback,
 } from "./change-events.js"
 import type { IndexInterface } from "../indexes/base-index.js"
-import type { BasicExpression, OrderBy } from "../query/ir.js"
+import { Value, type BasicExpression, type OrderBy } from "../query/ir.js"
 import type { ChangeMessage } from "../types.js"
 import type { CollectionImpl } from "./index.js"
 
@@ -153,6 +153,7 @@ export class CollectionSubscription {
     limit,
     minValue,
   }: RequestLimitedSnapshotOptions) {
+    console.log("in requestLimitedSnapshot")
     if (!limit) throw new Error(`limit is required`)
 
     if (!this.orderByIndex) {
@@ -203,10 +204,19 @@ export class CollectionSubscription {
 
     this.callback(changes)
 
+    let whereWithValueFilter = where
+    if (typeof minValue !== `undefined`) {
+      // Only request data that we haven't seen yet (i.e. is bigger than the minValue)
+      const { expression, compareOptions } = orderBy[0]!
+      const operator = compareOptions.direction === `asc` ? gt : lt
+      const valueFilter = operator(expression, new Value(minValue))
+      whereWithValueFilter = where ? and(where, valueFilter) : valueFilter
+    }
+
     // Request the sync layer to load more data
     // don't await it, we will load the data into the collection when it comes in
     this.collection.syncMore({
-      where,
+      where: whereWithValueFilter,
       limit,
       orderBy,
     })

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

@@ -1,5 +1,5 @@
 import { Func, PropRef, Value } from "../ir.js"
-import type { BasicExpression } from "../ir.js"
+import type { BasicExpression, OrderBy } from "../ir.js"
 
 /**
  * Functions supported by the collection index system.
@@ -90,3 +90,26 @@ export function convertToBasicExpression(
     return new Func(whereClause.name, args)
   }
 }
+
+export function convertOrderByToBasicExpression(
+  orderBy: OrderBy,
+  collectionAlias: string
+): OrderBy {
+  const normalizedOrderBy = orderBy.map((clause) => {
+    const basicExp = convertToBasicExpression(
+      clause.expression,
+      collectionAlias
+    )
+
+    if (!basicExp) {
+      throw new Error(`Failed to convert orderBy expression to a basic expression: ${clause.expression}`)
+    }
+
+    return {
+      ...clause,
+      expression: basicExp,
+    }
+  })
+
+  return normalizedOrderBy
+}

packages/db/src/query/live/collection-subscriber.ts

@@ -1,5 +1,5 @@
 import { MultiSet } from "@tanstack/db-ivm"
-import { convertToBasicExpression } from "../compiler/expressions.js"
+import { convertOrderByToBasicExpression, convertToBasicExpression } from "../compiler/expressions.js"
 import type { FullSyncState } from "./types.js"
 import type { MultiSetArray, RootStreamBuilder } from "@tanstack/db-ivm"
 import type { Collection } from "../../collection/index.js"
@@ -16,26 +16,29 @@ export class CollectionSubscriber<
   // Keep track of the biggest value we've sent so far (needed for orderBy optimization)
   private biggest: any = undefined
 
+  private collectionAlias: string
+
   constructor(
     private collectionId: string,
     private collection: Collection,
     private config: Parameters<SyncConfig<TResult>[`sync`]>[0],
     private syncState: FullSyncState,
     private collectionConfigBuilder: CollectionConfigBuilder<TContext, TResult>
-  ) {}
-
-  subscribe(): CollectionSubscription {
-    const collectionAlias = findCollectionAlias(
+  ) {
+    this.collectionAlias = findCollectionAlias(
       this.collectionId,
       this.collectionConfigBuilder.query
-    )
-    const whereClause = this.getWhereClauseFromAlias(collectionAlias)
+    )!
+  }
+
+  subscribe(): CollectionSubscription {
+    const whereClause = this.getWhereClauseFromAlias(this.collectionAlias)
 
     if (whereClause) {
       // Convert WHERE clause to BasicExpression format for collection subscription
       const whereExpression = convertToBasicExpression(
         whereClause,
-        collectionAlias!
+        this.collectionAlias
       )
 
       if (whereExpression) {
@@ -164,11 +167,14 @@ export class CollectionSubscriber<
 
     subscription.setOrderByIndex(index)
 
+    // Normalize the orderBy clauses such that the references are relative to the collection
+    const normalizedOrderBy = convertOrderByToBasicExpression(orderBy, this.collectionAlias)
+
     // Load the first `offset + limit` values from the index
     // i.e. the K items from the collection that fall into the requested range: [offset, offset + limit[
     subscription.requestLimitedSnapshot({
       limit: offset + limit,
-      orderBy,
+      orderBy: normalizedOrderBy,
     })
 
     return subscription
@@ -234,9 +240,16 @@ export class CollectionSubscriber<
     const biggestSentValue = biggestSentRow
       ? valueExtractorForRawRow(biggestSentRow)
       : biggestSentRow
+    
+    // Normalize the orderBy clauses such that the references are relative to the collection
+    const normalizedOrderBy = convertOrderByToBasicExpression(
+      orderBy,
+      this.collectionAlias
+    )
+
     // Take the `n` items after the biggest sent value
     subscription.requestLimitedSnapshot({
-      orderBy,
+      orderBy: normalizedOrderBy,
       limit: n,
       minValue: biggestSentValue,
     })