fix: support params in use document projection

packages/core/src/projection/getProjectionState.ts

@@ -22,6 +22,7 @@ export interface ProjectionOptions<
   TProjectId extends string = string,
 > extends DocumentHandle<TDocumentType, TDataset, TProjectId> {
   projection: TProjection
+  params?: Record<string, unknown>
 }
 
 /**
@@ -82,11 +83,12 @@ export const _getProjectionState = bindActionByDataset(
       return state.values[documentId]?.[projectionHash] ?? STABLE_EMPTY_PROJECTION
     },
     onSubscribe: ({state}, options: ProjectionOptions<ValidProjection, string, string, string>) => {
-      const {projection, ...docHandle} = options
+      const {projection, params, ...docHandle} = options
       const subscriptionId = insecureRandomId()
       const documentId = getPublishedId(docHandle.documentId)
       const validProjection = validateProjection(projection)
-      const projectionHash = hashString(validProjection)
+      const paramsHash = params ? `:${hashString(JSON.stringify(params))}` : ''
+      const projectionHash = hashString(`${validProjection}${paramsHash}`)
 
       state.set('addSubscription', (prev) => ({
         documentProjections: {
@@ -96,6 +98,13 @@ export const _getProjectionState = bindActionByDataset(
             [projectionHash]: validProjection,
           },
         },
+        projectionParams: {
+          ...prev.projectionParams,
+          [documentId]: {
+            ...prev.projectionParams[documentId],
+            [projectionHash]: params,
+          },
+        },
         subscriptions: {
           ...prev.subscriptions,
           [documentId]: {
@@ -119,12 +128,14 @@ export const _getProjectionState = bindActionByDataset(
 
             const nextSubscriptions = {...prev.subscriptions}
             const nextDocumentProjections = {...prev.documentProjections}
+            const nextProjectionParams = {...prev.projectionParams}
             const nextValues = {...prev.values}
 
             // clean up the subscription and documentProjection if there are no subscribers
             if (!hasSubscribersForProjection) {
               delete nextSubscriptions[documentId]![projectionHash]
               delete nextDocumentProjections[documentId]![projectionHash]
+              delete nextProjectionParams[documentId]![projectionHash]
 
               const currentProjectionValue = prev.values[documentId]?.[projectionHash]
               if (currentProjectionValue && nextValues[documentId]) {
@@ -146,12 +157,14 @@ export const _getProjectionState = bindActionByDataset(
             if (!hasAnySubscribersForDocument) {
               delete nextSubscriptions[documentId]
               delete nextDocumentProjections[documentId]
+              delete nextProjectionParams[documentId]
               // Keep nextValues[documentId] as cache
             }
 
             return {
               subscriptions: nextSubscriptions,
               documentProjections: nextDocumentProjections,
+              projectionParams: nextProjectionParams,
               values: nextValues,
             }
           })

packages/core/src/projection/projectionQuery.ts

@@ -18,11 +18,17 @@ interface CreateProjectionQueryResult {
   params: Record<string, unknown>
 }
 
-type ProjectionMap = Record<string, {projection: ValidProjection; documentIds: Set<string>}>
+type ProjectionMap = Record<
+  string,
+  {projection: ValidProjection; params?: Record<string, unknown>; documentIds: Set<string>}
+>
 
 export function createProjectionQuery(
   documentIds: Set<string>,
   documentProjections: {[TDocumentId in string]?: DocumentProjections},
+  projectionParams?: {
+    [TDocumentId in string]?: {[hash: string]: Record<string, unknown> | undefined}
+  },
 ): CreateProjectionQueryResult {
   const projections = Array.from(documentIds)
     .flatMap((id) => {
@@ -32,11 +38,12 @@ export function createProjectionQuery(
       return Object.entries(projectionsForDoc).map(([projectionHash, projection]) => ({
         documentId: id,
         projection,
+        params: projectionParams?.[id]?.[projectionHash],
         projectionHash,
       }))
     })
-    .reduce<ProjectionMap>((acc, {documentId, projection, projectionHash}) => {
-      const obj = acc[projectionHash] ?? {documentIds: new Set(), projection}
+    .reduce<ProjectionMap>((acc, {documentId, projection, params, projectionHash}) => {
+      const obj = acc[projectionHash] ?? {documentIds: new Set(), projection, params}
       obj.documentIds.add(documentId)
 
       acc[projectionHash] = obj
@@ -45,11 +52,20 @@ export function createProjectionQuery(
 
   const query = `[${Object.entries(projections)
     .map(([projectionHash, {projection}]) => {
-      return `...*[_id in $__ids_${projectionHash}]{_id,_type,_updatedAt,"__projectionHash":"${projectionHash}","result":{...${projection}}}`
+      // Rename parameter references inside the projection to avoid collisions between batches
+      const rewrittenProjection = (projection as ValidProjection).replace(
+        /\$[A-Za-z_][A-Za-z0-9_]*/g,
+        (match) => {
+          const paramName = match.slice(1)
+          return `$__p_${projectionHash}_${paramName}`
+        },
+      )
+
+      return `...*[_id in $__ids_${projectionHash}]{_id,_type,_updatedAt,"__projectionHash":"${projectionHash}","result":{...${rewrittenProjection}}}`
     })
     .join(',')}]`
 
-  const params = Object.fromEntries(
+  const idsParams = Object.fromEntries(
     Object.entries(projections).map(([projectionHash, value]) => {
       const idsInProjection = Array.from(value.documentIds).flatMap((id) => [
         getPublishedId(id),
@@ -60,6 +76,15 @@ export function createProjectionQuery(
     }),
   )
 
+  // Merge all custom params across unique projection hashes. Since hashes include params hash, keys won't collide between different param sets.
+  const customParams = Object.fromEntries(
+    Object.entries(projections).flatMap(([projectionHash, {params}]) =>
+      Object.entries(params ?? {}).map(([k, v]) => [`__p_${projectionHash}_${k}`, v]),
+    ),
+  )
+
+  const params = {...idsParams, ...customParams}
+
   return {query, params}
 }
 

packages/core/src/projection/projectionStore.ts

@@ -8,6 +8,7 @@ export const projectionStore = defineStore<ProjectionStoreState>({
     return {
       values: {},
       documentProjections: {},
+      projectionParams: {},
       subscriptions: {},
     }
   },

packages/core/src/projection/subscribeToStateAndFetchBatches.ts

@@ -87,7 +87,11 @@
     .pipe(
       switchMap(([ids, documentProjections]) => {
         if (!ids.size) return EMPTY
-        const {query, params} = createProjectionQuery(ids, documentProjections)
+        const {query, params} = createProjectionQuery(
+          ids,
+          documentProjections,
+          state.get().projectionParams,
+        )
         const controller = new AbortController()
 
         return new Observable<ProjectionQueryResult[]>((observer) => {
@@ -150,7 +154,7 @@
       error: (err) => {
         // eslint-disable-next-line no-console
         console.error('Error fetching projection batches:', err)
         // TODO: Potentially update state to reflect error state for affected projections?
       },
     })
 

packages/core/src/projection/types.ts

@@ -20,6 +20,10 @@ export interface DocumentProjections {
   [projectionHash: string]: ValidProjection
 }
 
+export interface DocumentProjectionParams {
+  [projectionHash: string]: Record<string, unknown> | undefined
+}
+
 interface DocumentProjectionSubscriptions {
   [projectionHash: string]: {
     [subscriptionId: string]: true
@@ -41,6 +45,13 @@ export interface ProjectionStoreState<TValue extends object = object> {
     [documentId: string]: DocumentProjections
   }
 
+  /**
+   * A map of document IDs to their projection params, organized by projection hash
+   */
+  projectionParams: {
+    [documentId: string]: DocumentProjectionParams
+  }
+
   /**
    * A map of document IDs to their subscriptions, organized by projection hash
    */

packages/react/src/hooks/projection/useDocumentProjection.test.tsx

@@ -280,4 +280,91 @@ describe('useDocumentProjection', () => {
     expect(subscribe).toHaveBeenCalled()
     expect(screen.getByText('Title')).toBeInTheDocument()
   })
+
+  test('it forwards params to core when provided', async () => {
+    const params = {q: 'Star'}
+
+    // Trigger suspense first
+    getCurrent.mockReturnValueOnce({
+      data: null,
+      isPending: true,
+    })
+
+    const resolvedData = {
+      data: {title: 'Resolved Title', description: 'Resolved Description'},
+      isPending: false,
+    }
+
+    ;(resolveProjection as Mock).mockReturnValueOnce(Promise.resolve(resolvedData))
+    getCurrent.mockReturnValue(resolvedData)
+    subscribe.mockReturnValue(() => {})
+
+    function WithParamsComponent({
+      projection,
+      ...docHandle
+    }: DocumentHandle & {projection: ValidProjection}) {
+      const ref = useRef(null)
+      const {data} = useDocumentProjection<ProjectionResult>({
+        ...docHandle,
+        projection,
+        params,
+        ref,
+      })
+      return (
+        <div ref={ref}>
+          <h1>{data.title}</h1>
+        </div>
+      )
+    }
+
+    render(
+      <Suspense fallback={<div>Loading...</div>}>
+        <WithParamsComponent {...mockDocument} projection="{title}" />
+      </Suspense>,
+    )
+
+    await act(async () => {
+      intersectionObserverCallback([{isIntersecting: true} as IntersectionObserverEntry])
+      await Promise.resolve()
+    })
+
+    const getProjectionCalls = (getProjectionState as unknown as Mock).mock.calls
+    const getProjectionCall = getProjectionCalls[getProjectionCalls.length - 1]
+    const resolveProjectionCalls = (resolveProjection as Mock).mock.calls
+    const resolveProjectionCall = resolveProjectionCalls[resolveProjectionCalls.length - 1]
+
+    expect(getProjectionCall[1].params).toEqual(params)
+    expect(resolveProjectionCall[1].params).toEqual(params)
+    expect(screen.getByText('Resolved Title')).toBeInTheDocument()
+  })
+
+  test('it works without params and does not set params field', async () => {
+    getCurrent.mockReturnValue({
+      data: {title: 'No Params', description: 'Desc'},
+      isPending: false,
+    })
+    subscribe.mockReturnValue(() => {})
+
+    function NoParamsComponent({
+      projection,
+      ...docHandle
+    }: DocumentHandle & {projection: ValidProjection}) {
+      const {data} = useDocumentProjection<ProjectionResult>({
+        ...docHandle,
+        projection,
+      })
+      return <h1>{data.title}</h1>
+    }
+
+    render(
+      <Suspense fallback={<div>Loading...</div>}>
+        <NoParamsComponent {...mockDocument} projection="{title}" />
+      </Suspense>,
+    )
+
+    const getProjectionCalls = (getProjectionState as unknown as Mock).mock.calls
+    const getProjectionCall = getProjectionCalls[getProjectionCalls.length - 1]
+    expect('params' in getProjectionCall[1]).toBe(false)
+    expect(screen.getByText('No Params')).toBeInTheDocument()
+  })
 })

packages/react/src/hooks/projection/useDocumentProjection.ts

@@ -179,13 +179,22 @@ export function useDocumentProjection<TData extends object>(
 export function useDocumentProjection<TData extends object>({
   ref,
   projection,
+  params,
   ...docHandle
 }: useDocumentProjectionOptions): useDocumentProjectionResults<TData> {
   const instance = useSanityInstance(docHandle)
-  const stateSource = getProjectionState<TData>(instance, {...docHandle, projection})
+  const stateSource = getProjectionState<TData>(instance, {
+    ...docHandle,
+    projection,
+    ...(params ? {params} : {}),
+  })
 
   if (stateSource.getCurrent()?.data === null) {
-    throw resolveProjection(instance, {...docHandle, projection})
+    throw resolveProjection(instance, {
+      ...docHandle,
+      projection,
+      ...(params ? {params} : {}),
+    })
   }
 
   // Create subscribe function for useSyncExternalStore