Return query collection from useLiveQuery (#74)

Co-authored-by: Sam Willis <sam.willis@gmail.com>

Author: KyleAMathews

packages/react-db/src/useLiveQuery.ts

@@ -2,6 +2,7 @@ import { useEffect, useMemo, useState } from "react"
 import { useStore } from "@tanstack/react-store"
 import { compileQuery, queryBuilder } from "@tanstack/db"
 import type {
+  Collection,
   Context,
   InitialQueryBuilder,
   QueryBuilder,
@@ -12,6 +13,7 @@ import type {
 export interface UseLiveQueryReturn<T extends object> {
   state: Map<string, T>
   data: Array<T>
+  collection: Collection<T>
 }
 
 export function useLiveQuery<
@@ -50,5 +52,6 @@ export function useLiveQuery<
   return {
     state,
     data,
+    collection: compiledQuery.results,
   }
 }

packages/react-db/tests/useLiveQuery.test.tsx

@@ -17,6 +17,7 @@ type Person = {
   age: number
   email: string
   isActive: boolean
+  team: string
 }
 
 type Issue = {
@@ -33,20 +34,23 @@ const initialPersons: Array<Person> = [
     age: 30,
     email: `john.doe@example.com`,
     isActive: true,
+    team: `team1`,
   },
   {
     id: `2`,
     name: `Jane Doe`,
     age: 25,
     email: `jane.doe@example.com`,
     isActive: true,
+    team: `team2`,
   },
   {
     id: `3`,
     name: `John Smith`,
     age: 35,
     email: `john.smith@example.com`,
     isActive: true,
+    team: `team1`,
   },
 ]
 
@@ -592,6 +596,116 @@ describe(`Query Collections`, () => {
     // Restore console.log
     console.log = originalConsoleLog
   })
+
+  it(`should be able to query a result collection`, async () => {
+    const emitter = mitt()
+
+    // Create collection with mutation capability
+    const collection = new Collection<Person>({
+      id: `optimistic-changes-test`,
+      sync: {
+        sync: ({ begin, write, commit }) => {
+          // Listen for sync events
+          emitter.on(`*`, (_, changes) => {
+            begin()
+            ;(changes as Array<PendingMutation>).forEach((change) => {
+              write({
+                key: change.key,
+                type: change.type,
+                value: change.changes as Person,
+              })
+            })
+            commit()
+          })
+        },
+      },
+    })
+
+    // Sync from initial state
+    act(() => {
+      emitter.emit(
+        `sync`,
+        initialPersons.map((person) => ({
+          key: person.id,
+          type: `insert`,
+          changes: person,
+        }))
+      )
+    })
+
+    // Initial query
+    const { result } = renderHook(() => {
+      return useLiveQuery((q) =>
+        q
+          .from({ collection })
+          .where(`@age`, `>`, 30)
+          .keyBy(`@id`)
+          .select(`@id`, `@name`, `@team`)
+          .orderBy({ "@id": `asc` })
+      )
+    })
+
+    // Grouped query derived from initial query
+    const { result: groupedResult } = renderHook(() => {
+      return useLiveQuery((q) =>
+        q
+          .from({ queryResult: result.current.collection })
+          .groupBy(`@team`)
+          .keyBy(`@team`)
+          .select(`@team`, { count: { COUNT: `@id` } })
+      )
+    })
+
+    // Verify initial grouped results
+    expect(groupedResult.current.state.size).toBe(1)
+    expect(groupedResult.current.state.get(`team1`)).toEqual({
+      team: `team1`,
+      count: 1,
+    })
+
+    // Insert two new users in different teams
+    act(() => {
+      emitter.emit(`sync`, [
+        {
+          key: `5`,
+          type: `insert`,
+          changes: {
+            id: `5`,
+            name: `Sarah Jones`,
+            age: 32,
+            email: `sarah.jones@example.com`,
+            isActive: true,
+            team: `team1`,
+          },
+        },
+        {
+          key: `6`,
+          type: `insert`,
+          changes: {
+            id: `6`,
+            name: `Mike Wilson`,
+            age: 38,
+            email: `mike.wilson@example.com`,
+            isActive: true,
+            team: `team2`,
+          },
+        },
+      ])
+    })
+
+    await waitForChanges()
+
+    // Verify the grouped results include the new team members
+    expect(groupedResult.current.state.size).toBe(2)
+    expect(groupedResult.current.state.get(`team1`)).toEqual({
+      team: `team1`,
+      count: 2,
+    })
+    expect(groupedResult.current.state.get(`team2`)).toEqual({
+      team: `team2`,
+      count: 1,
+    })
+  })
 })
 
 async function waitForChanges(ms = 0) {