Fixes to electric collection + unit test

Author: kevin-dp

packages/electric-db-collection/src/electric.ts

@@ -15,6 +15,7 @@ import {
 import { compileSQL } from "./sql-compiler"
 import type {
   BaseCollectionConfig,
+  Collection,
   CollectionConfig,
   DeleteMutationFnParams,
   Fn,
@@ -419,7 +420,10 @@ function createElectricSync<T extends Row<unknown>>(
       })
 
       return {
-        onLoadMore: (opts) => onLoadMore(stream, params, opts),
+        onLoadMore: (opts) => {
+          const snapshotParams = compileSQL<T>(opts)
+          return stream.requestSnapshot(snapshotParams)
+        },
         cleanup: () => {
           // Unsubscribe from the stream
           unsubscribeStream()
@@ -433,21 +437,3 @@ function createElectricSync<T extends Row<unknown>>(
   }
 }
 
-async function onLoadMore<T extends Row<unknown>>(
-  stream: ShapeStream<T>,
-  syncParams: Parameters<SyncConfig<T>[`sync`]>[0],
-  options: OnLoadMoreOptions
-) {
-  const { begin, write, commit } = syncParams
-  const snapshotParams = compileSQL<T>(options)
-  const snapshot = await stream.requestSnapshot(snapshotParams)
-
-  begin()
-  snapshot.data.forEach((row) => {
-    write({
-      type: `insert`,
-      value: row.value,
-    })
-  })
-  commit()
-}

packages/electric-db-collection/src/pg-serializer.ts

@@ -3,6 +3,10 @@ export function serialize(value: unknown): string {
     return `'${value}'`
   }
 
+  if (typeof value === `number`) {
+    return value.toString()
+  }
+
   if (value === null || value === undefined) {
     return `NULL`
   }
@@ -19,9 +23,5 @@ export function serialize(value: unknown): string {
     return `ARRAY[${value.map(serialize).join(`,`)}]`
   }
 
-  if (typeof value === `object`) {
-    throw new Error(`Cannot serialize object: ${JSON.stringify(value)}`)
-  }
-
-  return value.toString()
+  throw new Error(`Cannot serialize value: ${JSON.stringify(value)}`)
 }

packages/electric-db-collection/src/sql-compiler.ts

@@ -30,6 +30,7 @@ export function compileSQL<T>(
 
   // Serialize the values in the params array into PG formatted strings
   // and transform the array into a Record<string, string>
+  console.log("params", params)
   const paramsRecord = params.reduce(
     (acc, param, index) => {
       acc[`${index + 1}`] = serialize(param)
@@ -58,7 +59,8 @@ function compileBasicExpression(
     case `val`:
       params.push(exp.value)
       return `$${params.length}`
-    case `ref`:
+    case `ref`: 
+      // TODO: doesn't yet support JSON(B) values which could be accessed with nested props
       if (exp.path.length !== 1) {
         throw new Error(
           `Compiler can't handle nested properties: ${exp.path.join(`.`)}`

packages/electric-db-collection/tests/electric-live-query.test.ts

@@ -54,10 +54,30 @@ const sampleUsers: Array<User> = [
 
 // Mock the ShapeStream module
 const mockSubscribe = vi.fn()
+const mockRequestSnapshot = vi.fn()
 const mockStream = {
   subscribe: mockSubscribe,
+  requestSnapshot: (...args: any) => {
+    mockRequestSnapshot(...args)
+    const results = mockRequestSnapshot.mock.results
+    const lastResult = results[results.length - 1]!.value
+
+    const subscribers = mockSubscribe.mock.calls.map(args => args[0])
+    subscribers.forEach(subscriber => subscriber(lastResult.data.map((row: any) => ({
+      type: `insert`,
+      value: row.value,
+      key: row.key,
+    }))))
+  }
 }
 
+// Mock the requestSnapshot method
+// to return an empty array of data
+// since most tests don't use it
+mockRequestSnapshot.mockResolvedValue({
+  data: []
+})
+
 vi.mock(`@electric-sql/client`, async () => {
   const actual = await vi.importActual(`@electric-sql/client`)
   return {
@@ -437,4 +457,98 @@ describe.each([
     // Clean up
     subscription.unsubscribe()
   })
+  if (autoIndex === `eager`) {
+    it.only(`should load more data via requestSnapshot when creating live query with higher limit`, async () => {
+      // Reset mocks
+      vi.clearAllMocks()
+      mockRequestSnapshot.mockResolvedValue({
+        data: [
+          { key: 5, value: { id: 5, name: `Eve`, age: 30, email: `[email protected]`, active: true } },
+          { key: 6, value: { id: 6, name: `Frank`, age: 35, email: `[email protected]`, active: true } },
+        ],
+      })
+
+      // Initial sync with limited data
+      simulateInitialSync()
+      expect(electricCollection.status).toBe(`ready`)
+      expect(electricCollection.size).toBe(4)
+      expect(mockRequestSnapshot).toHaveBeenCalledTimes(0)
+
+      // Create first live query with limit of 2
+      const limitedLiveQuery = createLiveQueryCollection({
+        id: `limited-users-live-query`,
+        startSync: true,
+        query: (q) =>
+          q
+            .from({ user: electricCollection })
+            .where(({ user }) => eq(user.active, true))
+            .select(({ user }) => ({
+              id: user.id,
+              name: user.name,
+              active: user.active,
+              age: user.age,
+            }))
+            .orderBy(({ user }) => user.age, `asc`)
+            .limit(2),
+      })
+
+      expect(limitedLiveQuery.status).toBe(`ready`)
+      expect(limitedLiveQuery.size).toBe(2) // Only first 2 active users
+      expect(mockRequestSnapshot).toHaveBeenCalledTimes(1)
+
+      const callArgs = (index: number) => mockRequestSnapshot.mock.calls[index]?.[0]
+      expect(callArgs(0)).toMatchObject({
+        params: { "1": "true" },
+        where: "active = $1",
+        orderBy: "age NULLS FIRST",
+        limit: 2,
+      })
+
+      // Create second live query with higher limit of 5
+      const expandedLiveQuery = createLiveQueryCollection({
+        id: `expanded-users-live-query`,
+        startSync: true,
+        query: (q) =>
+          q
+            .from({ user: electricCollection })
+            .where(({ user }) => eq(user.active, true))
+            .select(({ user }) => ({
+              id: user.id,
+              name: user.name,
+              active: user.active,
+            }))
+            .orderBy(({ user }) => user.age, `asc`)
+            .limit(6),
+      })
+
+      // Wait for the live query to process
+      await new Promise((resolve) => setTimeout(resolve, 0))
+
+      // Verify that requestSnapshot was called with the correct parameters
+      expect(mockRequestSnapshot).toHaveBeenCalledTimes(3)
+
+      // Check that first it requested a limit of 6 users
+      expect(callArgs(1)).toMatchObject({
+        params: { "1": "true" },
+        where: "active = $1",
+        orderBy: "age NULLS FIRST",
+        limit: 6,
+      })
+
+      // After this initial snapshot for the new live query it receives all 3 users from the local collection
+      // so it still needs 3 more users to reach the limit of 6 so it requests 3 more to the sync layer
+      expect(callArgs(2)).toMatchObject({
+        params: { "1": "true", "2": "25" },
+        where: "active = $1 AND age > $2",
+        orderBy: "age NULLS FIRST",
+        limit: 3,
+      })
+
+      // The sync layer won't provide any more users so the DB is exhausted and it stops (i.e. doesn't request more) 
+
+      // The expanded live query should now have more data
+      expect(expandedLiveQuery.status).toBe(`ready`)
+      expect(expandedLiveQuery.size).toBe(5) // Alice, Bob, Dave from initial + Eve and Frank from additional data
+    })
+  }
 })