@tanstack/query-db-collection 0.3.0 → 1.0.1

package/src/query.ts

@@ -1,4 +1,4 @@
-import { QueryObserver } from "@tanstack/query-core"
+import { QueryObserver, hashKey } from "@tanstack/query-core"
 import {
   GetKeyRequiredError,
   QueryClientRequiredError,
@@ -6,23 +6,25 @@ import {
   QueryKeyRequiredError,
 } from "./errors"
 import { createWriteUtils } from "./manual-sync"
-import type {
-  QueryClient,
-  QueryFunctionContext,
-  QueryKey,
-  QueryObserverOptions,
-  QueryObserverResult,
-} from "@tanstack/query-core"
 import type {
   BaseCollectionConfig,
   ChangeMessage,
   CollectionConfig,
   DeleteMutationFnParams,
   InsertMutationFnParams,
+  LoadSubsetOptions,
   SyncConfig,
   UpdateMutationFnParams,
   UtilsRecord,
 } from "@tanstack/db"
+import type {
+  FetchStatus,
+  QueryClient,
+  QueryFunctionContext,
+  QueryKey,
+  QueryObserverOptions,
+  QueryObserverResult,
+} from "@tanstack/query-core"
 import type { StandardSchemaV1 } from "@standard-schema/spec"
 
 // Re-export for external use
@@ -42,6 +44,8 @@ type InferSchemaInput<T> = T extends StandardSchemaV1
     : Record<string, unknown>
   : Record<string, unknown>
 
+type TQueryKeyBuilder<TQueryKey> = (opts: LoadSubsetOptions) => TQueryKey
+
 /**
  * Configuration options for creating a Query Collection
  * @template T - The explicit type of items stored in the collection
@@ -63,7 +67,7 @@ export interface QueryCollectionConfig<
   TQueryData = Awaited<ReturnType<TQueryFn>>,
 > extends BaseCollectionConfig<T, TKey, TSchema> {
   /** The query key used by TanStack Query to identify this query */
-  queryKey: TQueryKey
+  queryKey: TQueryKey | TQueryKeyBuilder<TQueryKey>
   /** Function that fetches data from the server. Must return the complete collection state */
   queryFn: TQueryFn extends (
     context: QueryFunctionContext<TQueryKey>
@@ -201,9 +205,10 @@ interface QueryCollectionState {
   lastError: any
   errorCount: number
   lastErrorUpdatedAt: number
-  queryObserver:
-    | QueryObserver<Array<any>, any, Array<any>, Array<any>, any>
-    | undefined
+  observers: Map<
+    string,
+    QueryObserver<Array<any>, any, Array<any>, Array<any>, any>
+  >
 }
 
 /**
@@ -261,23 +266,40 @@ class QueryCollectionUtilsImpl {
 
   // Getters for QueryObserver state
   public get isFetching() {
-    return this.state.queryObserver?.getCurrentResult().isFetching ?? false
+    // check if any observer is fetching
+    return Array.from(this.state.observers.values()).some(
+      (observer) => observer.getCurrentResult().isFetching
+    )
   }
 
   public get isRefetching() {
-    return this.state.queryObserver?.getCurrentResult().isRefetching ?? false
+    // check if any observer is refetching
+    return Array.from(this.state.observers.values()).some(
+      (observer) => observer.getCurrentResult().isRefetching
+    )
   }
 
   public get isLoading() {
-    return this.state.queryObserver?.getCurrentResult().isLoading ?? false
+    // check if any observer is loading
+    return Array.from(this.state.observers.values()).some(
+      (observer) => observer.getCurrentResult().isLoading
+    )
   }
 
   public get dataUpdatedAt() {
-    return this.state.queryObserver?.getCurrentResult().dataUpdatedAt ?? 0
+    // compute the max dataUpdatedAt of all observers
+    return Math.max(
+      0,
+      ...Array.from(this.state.observers.values()).map(
+        (observer) => observer.getCurrentResult().dataUpdatedAt
+      )
+    )
   }
 
-  public get fetchStatus() {
-    return this.state.queryObserver?.getCurrentResult().fetchStatus ?? `idle`
+  public get fetchStatus(): Array<FetchStatus> {
+    return Array.from(this.state.observers.values()).map(
+      (observer) => observer.getCurrentResult().fetchStatus
+    )
   }
 }
 
@@ -522,6 +544,9 @@ export function queryCollectionOptions(
     ...baseCollectionConfig
   } = config
 
+  // Default to eager sync mode if not provided
+  const syncMode = baseCollectionConfig.syncMode ?? `eager`
+
   // Validate required parameters
 
   // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
@@ -548,181 +573,394 @@ export function queryCollectionOptions(
     lastError: undefined as any,
     errorCount: 0,
     lastErrorUpdatedAt: 0,
-    queryObserver: undefined,
+    observers: new Map<
+      string,
+      QueryObserver<Array<any>, any, Array<any>, Array<any>, any>
+    >(),
+  }
+
+  // hashedQueryKey → queryKey
+  const hashToQueryKey = new Map<string, QueryKey>()
+
+  // queryKey → Set<RowKey>
+  const queryToRows = new Map<string, Set<string | number>>()
+
+  // RowKey → Set<queryKey>
+  const rowToQueries = new Map<string | number, Set<string>>()
+
+  // queryKey → QueryObserver's unsubscribe function
+  const unsubscribes = new Map<string, () => void>()
+
+  // Helper function to add a row to the internal state
+  const addRow = (rowKey: string | number, hashedQueryKey: string) => {
+    const rowToQueriesSet = rowToQueries.get(rowKey) || new Set()
+    rowToQueriesSet.add(hashedQueryKey)
+    rowToQueries.set(rowKey, rowToQueriesSet)
+
+    const queryToRowsSet = queryToRows.get(hashedQueryKey) || new Set()
+    queryToRowsSet.add(rowKey)
+    queryToRows.set(hashedQueryKey, queryToRowsSet)
+  }
+
+  // Helper function to remove a row from the internal state
+  const removeRow = (rowKey: string | number, hashedQuerKey: string) => {
+    const rowToQueriesSet = rowToQueries.get(rowKey) || new Set()
+    rowToQueriesSet.delete(hashedQuerKey)
+    rowToQueries.set(rowKey, rowToQueriesSet)
+
+    const queryToRowsSet = queryToRows.get(hashedQuerKey) || new Set()
+    queryToRowsSet.delete(rowKey)
+    queryToRows.set(hashedQuerKey, queryToRowsSet)
+
+    return rowToQueriesSet.size === 0
   }
 
   const internalSync: SyncConfig<any>[`sync`] = (params) => {
     const { begin, write, commit, markReady, collection } = params
 
-    const observerOptions: QueryObserverOptions<
-      Array<any>,
-      any,
-      Array<any>,
-      Array<any>,
-      any
-    > = {
-      queryKey: queryKey,
-      queryFn: queryFn,
-      structuralSharing: true,
-      notifyOnChangeProps: `all`,
-      // Only include options that are explicitly defined to allow QueryClient defaultOptions to be used
-      ...(meta !== undefined && { meta }),
-      ...(enabled !== undefined && { enabled }),
-      ...(refetchInterval !== undefined && { refetchInterval }),
-      ...(retry !== undefined && { retry }),
-      ...(retryDelay !== undefined && { retryDelay }),
-      ...(staleTime !== undefined && { staleTime }),
-    }
-
-    const localObserver = new QueryObserver<
-      Array<any>,
-      any,
-      Array<any>,
-      Array<any>,
-      any
-    >(queryClient, observerOptions)
-
-    // Store reference for imperative refetch
-    state.queryObserver = localObserver
-
-    let isSubscribed = false
-    let actualUnsubscribeFn: (() => void) | null = null
-
-    type UpdateHandler = Parameters<typeof localObserver.subscribe>[0]
-    const handleQueryResult: UpdateHandler = (result) => {
-      if (result.isSuccess) {
-        // Clear error state
-        state.lastError = undefined
-        state.errorCount = 0
-
-        const rawData = result.data
-        const newItemsArray = select ? select(rawData) : rawData
-
-        if (
-          !Array.isArray(newItemsArray) ||
-          newItemsArray.some((item) => typeof item !== `object`)
-        ) {
-          const errorMessage = select
-            ? `@tanstack/query-db-collection: select() must return an array of objects. Got: ${typeof newItemsArray} for queryKey ${JSON.stringify(queryKey)}`
-            : `@tanstack/query-db-collection: queryFn must return an array of objects. Got: ${typeof newItemsArray} for queryKey ${JSON.stringify(queryKey)}`
-
-          console.error(errorMessage)
-          return
+    // Track whether sync has been started
+    let syncStarted = false
+
+    const createQueryFromOpts = (
+      opts: LoadSubsetOptions = {},
+      queryFunction: typeof queryFn = queryFn
+    ): true | Promise<void> => {
+      // Push the predicates down to the queryKey and queryFn
+      const key = typeof queryKey === `function` ? queryKey(opts) : queryKey
+      const hashedQueryKey = hashKey(key)
+      const extendedMeta = { ...meta, loadSubsetOptions: opts }
+
+      if (state.observers.has(hashedQueryKey)) {
+        // We already have a query for this queryKey
+        // Get the current result and return based on its state
+        const observer = state.observers.get(hashedQueryKey)!
+        const currentResult = observer.getCurrentResult()
+
+        if (currentResult.isSuccess) {
+          // Data is already available, return true synchronously
+          return true
+        } else if (currentResult.isError) {
+          // Error already occurred, reject immediately
+          return Promise.reject(currentResult.error)
+        } else {
+          // Query is still loading, wait for the first result
+          return new Promise<void>((resolve, reject) => {
+            const unsubscribe = observer.subscribe((result) => {
+              if (result.isSuccess) {
+                unsubscribe()
+                resolve()
+              } else if (result.isError) {
+                unsubscribe()
+                reject(result.error)
+              }
+            })
+          })
         }
+      }
+
+      const observerOptions: QueryObserverOptions<
+        Array<any>,
+        any,
+        Array<any>,
+        Array<any>,
+        any
+      > = {
+        queryKey: key,
+        queryFn: queryFunction,
+        meta: extendedMeta,
+        structuralSharing: true,
+        notifyOnChangeProps: `all`,
+
+        // Only include options that are explicitly defined to allow QueryClient defaultOptions to be used
+        ...(enabled !== undefined && { enabled }),
+        ...(refetchInterval !== undefined && { refetchInterval }),
+        ...(retry !== undefined && { retry }),
+        ...(retryDelay !== undefined && { retryDelay }),
+        ...(staleTime !== undefined && { staleTime }),
+      }
 
-        const currentSyncedItems: Map<string | number, any> = new Map(
-          collection._state.syncedData.entries()
-        )
-        const newItemsMap = new Map<string | number, any>()
-        newItemsArray.forEach((item) => {
-          const key = getKey(item)
-          newItemsMap.set(key, item)
+      const localObserver = new QueryObserver<
+        Array<any>,
+        any,
+        Array<any>,
+        Array<any>,
+        any
+      >(queryClient, observerOptions)
+
+      hashToQueryKey.set(hashedQueryKey, key)
+      state.observers.set(hashedQueryKey, localObserver)
+
+      // Create a promise that resolves when the query result is first available
+      const readyPromise = new Promise<void>((resolve, reject) => {
+        const unsubscribe = localObserver.subscribe((result) => {
+          if (result.isSuccess) {
+            unsubscribe()
+            resolve()
+          } else if (result.isError) {
+            unsubscribe()
+            reject(result.error)
+          }
         })
+      })
 
-        begin()
-
-        // Helper function for shallow equality check of objects
-        const shallowEqual = (
-          obj1: Record<string, any>,
-          obj2: Record<string, any>
-        ): boolean => {
-          // Get all keys from both objects
-          const keys1 = Object.keys(obj1)
-          const keys2 = Object.keys(obj2)
-
-          // If number of keys is different, objects are not equal
-          if (keys1.length !== keys2.length) return false
-
-          // Check if all keys in obj1 have the same values in obj2
-          return keys1.every((key) => {
-            // Skip comparing functions and complex objects deeply
-            if (typeof obj1[key] === `function`) return true
-            return obj1[key] === obj2[key]
-          })
-        }
+      // If sync has started or there are subscribers to the collection, subscribe to the query straight away
+      // This creates the main subscription that handles data updates
+      if (syncStarted || collection.subscriberCount > 0) {
+        subscribeToQuery(localObserver, hashedQueryKey)
+      }
+
+      // Tell tanstack query to GC the query when the subscription is unsubscribed
+      // The subscription is unsubscribed when the live query is GCed.
+      const subscription = opts.subscription
+      subscription?.once(`unsubscribed`, () => {
+        queryClient.removeQueries({ queryKey: key, exact: true })
+      })
+
+      return readyPromise
+    }
 
-        currentSyncedItems.forEach((oldItem, key) => {
-          const newItem = newItemsMap.get(key)
-          if (!newItem) {
-            write({ type: `delete`, value: oldItem })
-          } else if (
-            !shallowEqual(
-              oldItem as Record<string, any>,
-              newItem as Record<string, any>
-            )
+    type UpdateHandler = Parameters<QueryObserver[`subscribe`]>[0]
+
+    const makeQueryResultHandler = (queryKey: QueryKey) => {
+      const hashedQueryKey = hashKey(queryKey)
+      const handleQueryResult: UpdateHandler = (result) => {
+        if (result.isSuccess) {
+          // Clear error state
+          state.lastError = undefined
+          state.errorCount = 0
+
+          const rawData = result.data
+          const newItemsArray = select ? select(rawData) : rawData
+
+          if (
+            !Array.isArray(newItemsArray) ||
+            newItemsArray.some((item) => typeof item !== `object`)
           ) {
-            // Only update if there are actual differences in the properties
-            write({ type: `update`, value: newItem })
+            const errorMessage = select
+              ? `@tanstack/query-db-collection: select() must return an array of objects. Got: ${typeof newItemsArray} for queryKey ${JSON.stringify(queryKey)}`
+              : `@tanstack/query-db-collection: queryFn must return an array of objects. Got: ${typeof newItemsArray} for queryKey ${JSON.stringify(queryKey)}`
+
+            console.error(errorMessage)
+            return
           }
-        })
 
-        newItemsMap.forEach((newItem, key) => {
-          if (!currentSyncedItems.has(key)) {
-            write({ type: `insert`, value: newItem })
+          const currentSyncedItems: Map<string | number, any> = new Map(
+            collection._state.syncedData.entries()
+          )
+          const newItemsMap = new Map<string | number, any>()
+          newItemsArray.forEach((item) => {
+            const key = getKey(item)
+            newItemsMap.set(key, item)
+          })
+
+          begin()
+
+          // Helper function for shallow equality check of objects
+          const shallowEqual = (
+            obj1: Record<string, any>,
+            obj2: Record<string, any>
+          ): boolean => {
+            // Get all keys from both objects
+            const keys1 = Object.keys(obj1)
+            const keys2 = Object.keys(obj2)
+
+            // If number of keys is different, objects are not equal
+            if (keys1.length !== keys2.length) return false
+
+            // Check if all keys in obj1 have the same values in obj2
+            return keys1.every((key) => {
+              // Skip comparing functions and complex objects deeply
+              if (typeof obj1[key] === `function`) return true
+              return obj1[key] === obj2[key]
+            })
           }
-        })
 
-        commit()
+          currentSyncedItems.forEach((oldItem, key) => {
+            const newItem = newItemsMap.get(key)
+            if (!newItem) {
+              const needToRemove = removeRow(key, hashedQueryKey) // returns true if the row is no longer referenced by any queries
+              if (needToRemove) {
+                write({ type: `delete`, value: oldItem })
+              }
+            } else if (
+              !shallowEqual(
+                oldItem as Record<string, any>,
+                newItem as Record<string, any>
+              )
+            ) {
+              // Only update if there are actual differences in the properties
+              write({ type: `update`, value: newItem })
+            }
+          })
 
-        // Mark collection as ready after first successful query result
-        markReady()
-      } else if (result.isError) {
-        if (result.errorUpdatedAt !== state.lastErrorUpdatedAt) {
-          state.lastError = result.error
-          state.errorCount++
-          state.lastErrorUpdatedAt = result.errorUpdatedAt
-        }
+          newItemsMap.forEach((newItem, key) => {
+            addRow(key, hashedQueryKey)
+            if (!currentSyncedItems.has(key)) {
+              write({ type: `insert`, value: newItem })
+            }
+          })
 
-        console.error(
-          `[QueryCollection] Error observing query ${String(queryKey)}:`,
-          result.error
-        )
+          commit()
 
-        // Mark collection as ready even on error to avoid blocking apps
-        markReady()
+          // Mark collection as ready after first successful query result
+          markReady()
+        } else if (result.isError) {
+          if (result.errorUpdatedAt !== state.lastErrorUpdatedAt) {
+            state.lastError = result.error
+            state.errorCount++
+            state.lastErrorUpdatedAt = result.errorUpdatedAt
+          }
+
+          console.error(
+            `[QueryCollection] Error observing query ${String(queryKey)}:`,
+            result.error
+          )
+
+          // Mark collection as ready even on error to avoid blocking apps
+          markReady()
+        }
       }
+      return handleQueryResult
     }
 
-    const subscribeToQuery = () => {
-      if (!isSubscribed) {
-        actualUnsubscribeFn = localObserver.subscribe(handleQueryResult)
-        isSubscribed = true
-      }
+    const isSubscribed = (hashedQueryKey: string) => {
+      return unsubscribes.has(hashedQueryKey)
     }
 
-    const unsubscribeFromQuery = () => {
-      if (isSubscribed && actualUnsubscribeFn) {
-        actualUnsubscribeFn()
-        actualUnsubscribeFn = null
-        isSubscribed = false
+    const subscribeToQuery = (
+      observer: QueryObserver<Array<any>, any, Array<any>, Array<any>, any>,
+      hashedQueryKey: string
+    ) => {
+      if (!isSubscribed(hashedQueryKey)) {
+        const queryKey = hashToQueryKey.get(hashedQueryKey)!
+        const handleQueryResult = makeQueryResultHandler(queryKey)
+        const unsubscribeFn = observer.subscribe(handleQueryResult)
+        unsubscribes.set(hashedQueryKey, unsubscribeFn)
       }
     }
 
-    // Always subscribe when sync starts (this could be from preload(), startSync config, or first subscriber)
-    // We'll dynamically unsubscribe/resubscribe based on subscriber count to maintain staleTime behavior
-    subscribeToQuery()
+    const subscribeToQueries = () => {
+      state.observers.forEach(subscribeToQuery)
+    }
+
+    const unsubscribeFromQueries = () => {
+      unsubscribes.forEach((unsubscribeFn) => {
+        unsubscribeFn()
+      })
+      unsubscribes.clear()
+    }
+
+    // Mark that sync has started
+    syncStarted = true
 
     // Set up event listener for subscriber changes
     const unsubscribeFromCollectionEvents = collection.on(
       `subscribers:change`,
       ({ subscriberCount }) => {
         if (subscriberCount > 0) {
-          subscribeToQuery()
+          subscribeToQueries()
         } else if (subscriberCount === 0) {
-          unsubscribeFromQuery()
+          unsubscribeFromQueries()
         }
       }
     )
 
-    // Ensure we process any existing query data (QueryObserver doesn't invoke its callback automatically with initial
-    // state)
-    handleQueryResult(localObserver.getCurrentResult())
+    // If syncMode is eager, create the initial query without any predicates
+    if (syncMode === `eager`) {
+      // Catch any errors to prevent unhandled rejections
+      const initialResult = createQueryFromOpts({})
+      if (initialResult instanceof Promise) {
+        initialResult.catch(() => {
+          // Errors are already handled by the query result handler
+        })
+      }
+    } else {
+      // In on-demand mode, mark ready immediately since there's no initial query
+      markReady()
+    }
+
+    // Always subscribe when sync starts (this could be from preload(), startSync config, or first subscriber)
+    // We'll dynamically unsubscribe/resubscribe based on subscriber count to maintain staleTime behavior
+    subscribeToQueries()
+
+    // Ensure we process any existing query data (QueryObserver doesn't invoke its callback automatically with initial state)
+    state.observers.forEach((observer, hashedQueryKey) => {
+      const queryKey = hashToQueryKey.get(hashedQueryKey)!
+      const handleQueryResult = makeQueryResultHandler(queryKey)
+      handleQueryResult(observer.getCurrentResult())
+    })
+
+    // Subscribe to the query client's cache to handle queries that are GCed by tanstack query
+    const unsubscribeQueryCache = queryClient
+      .getQueryCache()
+      .subscribe((event) => {
+        const hashedKey = event.query.queryHash
+        if (event.type === `removed`) {
+          cleanupQuery(hashedKey)
+        }
+      })
+
+    function cleanupQuery(hashedQueryKey: string) {
+      // Unsubscribe from the query's observer
+      unsubscribes.get(hashedQueryKey)?.()
+
+      // Get all the rows that are in the result of this query
+      const rowKeys = queryToRows.get(hashedQueryKey) ?? new Set()
+
+      // Remove the query from these rows
+      rowKeys.forEach((rowKey) => {
+        const queries = rowToQueries.get(rowKey) // set of queries that reference this row
+        if (queries && queries.size > 0) {
+          queries.delete(hashedQueryKey)
+          if (queries.size === 0) {
+            // Reference count dropped to 0, we can GC the row
+            rowToQueries.delete(rowKey)
+
+            if (collection.has(rowKey)) {
+              begin()
+              write({ type: `delete`, value: collection.get(rowKey) })
+              commit()
+            }
+          }
+        }
+      })
 
-    return async () => {
+      // Remove the query from the internal state
+      unsubscribes.delete(hashedQueryKey)
+      state.observers.delete(hashedQueryKey)
+      queryToRows.delete(hashedQueryKey)
+      hashToQueryKey.delete(hashedQueryKey)
+    }
+
+    const cleanup = async () => {
       unsubscribeFromCollectionEvents()
-      unsubscribeFromQuery()
-      await queryClient.cancelQueries({ queryKey })
-      queryClient.removeQueries({ queryKey })
+      unsubscribeFromQueries()
+
+      const queryKeys = [...hashToQueryKey.values()]
+
+      hashToQueryKey.clear()
+      queryToRows.clear()
+      rowToQueries.clear()
+      state.observers.clear()
+      unsubscribeQueryCache()
+
+      await Promise.all(
+        queryKeys.map(async (queryKey) => {
+          await queryClient.cancelQueries({ queryKey })
+          queryClient.removeQueries({ queryKey })
+        })
+      )
+    }
+
+    // Create deduplicated loadSubset wrapper for non-eager modes
+    // This prevents redundant snapshot requests when multiple concurrent
+    // live queries request overlapping or subset predicates
+    const loadSubsetDedupe =
+      syncMode === `eager` ? undefined : createQueryFromOpts
+
+    return {
+      loadSubset: loadSubsetDedupe,
+      cleanup,
     }
   }
 
@@ -745,15 +983,15 @@ export function queryCollectionOptions(
    * @returns Promise that resolves when the refetch is complete, with QueryObserverResult
    */
   const refetch: RefetchFn = async (opts) => {
-    // Observer is created when sync starts. If never synced, nothing to refetch.
-
-    if (!state.queryObserver) {
-      return
-    }
-    // Return the QueryObserverResult for users to inspect
-    return state.queryObserver.refetch({
-      throwOnError: opts?.throwOnError,
+    const queryKeys = [...hashToQueryKey.values()]
+    const refetchPromises = queryKeys.map((queryKey) => {
+      const queryObserver = state.observers.get(hashKey(queryKey))!
+      return queryObserver.refetch({
+        throwOnError: opts?.throwOnError,
+      })
     })
+
+    await Promise.all(refetchPromises)
   }
 
   // Create write context for manual write operations
@@ -840,6 +1078,7 @@ export function queryCollectionOptions(
   return {
     ...baseCollectionConfig,
     getKey,
+    syncMode,
     sync: { sync: enhancedInternalSync },
     onInsert: wrappedOnInsert,
     onUpdate: wrappedOnUpdate,