feat: add shouldCacheResult option

package/src/context.test.ts

@@ -1,5 +1,5 @@
 import type {QueryStore} from '@sanity/react-loader'
-import {CacheShort, WithCache} from '@shopify/hydrogen'
+import {CacheShort, type WithCache} from '@shopify/hydrogen'
 import groq from 'groq'
 import {beforeEach, describe, expect, it, vi} from 'vitest'
 
@@ -25,9 +25,11 @@ let withCache = vi.hoisted<WithCache | null>(() => null)
 
 vi.mock('@shopify/hydrogen', async (importOriginal) => {
   const module = await importOriginal<typeof import('@shopify/hydrogen')>()
-  withCache = vi
-    .fn()
-    .mockImplementation(module.createWithCache({cache, waitUntil: () => Promise.resolve()}))
+  withCache = module.createWithCache({
+    cache,
+    waitUntil: () => Promise.resolve(),
+    request: new Request('https://example.com'),
+  })
 
   return {
     ...module,
@@ -44,6 +46,8 @@ vi.mock('@sanity/react-loader', async (importOriginal) => {
   }
 })
 
+const runWithCache = vi.spyOn(withCache!, 'run')
+
 const client = createClient({projectId: 'my-project-id', dataset: 'my-dataset'})
 
 const query = groq`true`
@@ -73,17 +77,25 @@ describe('the Sanity request context', () => {
     })
 
     await contextWithDefaultStrategy.loadQuery<boolean>(query, params)
-    expect(withCache).toHaveBeenCalledWith(hashedQuery, defaultStrategy, expect.any(Function))
-    expect(cache.put).toHaveBeenCalled()
+    expect(runWithCache).toHaveBeenNthCalledWith(
+      1,
+      expect.objectContaining({cacheKey: hashedQuery, cacheStrategy: defaultStrategy}),
+      expect.any(Function),
+    )
+    expect(cache.put).toHaveBeenCalledOnce()
   })
 
   it('queries should use the cache strategy passed in `loadQuery`', async () => {
     const strategy = CacheShort()
     await sanityContext.loadQuery<boolean>(query, params, {
       hydrogen: {cache: strategy},
     })
-    expect(withCache).toHaveBeenCalledWith(hashedQuery, strategy, expect.any(Function))
-    expect(cache.put).toHaveBeenCalled()
+    expect(runWithCache).toHaveBeenNthCalledWith(
+      1,
+      expect.objectContaining({cacheKey: hashedQuery, cacheStrategy: strategy}),
+      expect.any(Function),
+    )
+    expect(cache.put).toHaveBeenCalledOnce()
   })
 })
 
@@ -122,6 +134,6 @@ describe('when configured for preview', () => {
   it(`shouldn't cache queries`, async () => {
     await previewContext.loadQuery<boolean>(query)
     expect(loadQuery).toHaveBeenCalledOnce()
-    expect(cache.put).not.toHaveBeenCalled()
+    expect(cache.put).not.toHaveBeenCalledOnce()
   })
 })

package/src/context.ts

@@ -62,10 +62,18 @@ type HydrogenResponseQueryOptions = Omit<ResponseQueryOptions, 'next' | 'cache'>
   hydrogen?: 'hydrogen' extends keyof RequestInit ? RequestInit['hydrogen'] : never
 }
 
-type LoadQueryOptions = Pick<
+type LoadQueryOptions<T> = Pick<
   HydrogenResponseQueryOptions,
   'perspective' | 'hydrogen' | 'useCdn' | 'stega' | 'headers' | 'tag'
->
+> & {
+  hydrogen?: {
+    /**
+     * Whether to cache the result of the query or not.
+     * @default () => true
+     */
+    shouldCacheResult?: (value: QueryResponseInitial<T>) => boolean
+  }
+}
 
 export type SanityContext = {
   /**
@@ -76,7 +84,7 @@ export type SanityContext = {
   loadQuery<T = any>(
     query: string,
     params?: QueryParams,
-    options?: LoadQueryOptions,
+    options?: LoadQueryOptions<T>,
   ): Promise<QueryResponseInitial<T>>
 
   client: SanityClient
@@ -125,7 +133,7 @@ export function createSanityContext(options: CreateSanityContextOptions): Sanity
     async loadQuery<T>(
       query: string,
       params: QueryParams | QueryWithoutParams,
-      loaderOptions?: LoadQueryOptions,
+      loaderOptions?: LoadQueryOptions<T>,
     ): Promise<QueryResponseInitial<T>> {
       // Don't store response if preview is enabled
       const cacheStrategy =
@@ -134,20 +142,28 @@ export function createSanityContext(options: CreateSanityContextOptions): Sanity
           : loaderOptions?.hydrogen?.cache || defaultStrategy || DEFAULT_CACHE_STRATEGY
 
       const queryHash = await hashQuery(query, params)
+      const shouldCacheResult = loaderOptions?.hydrogen?.shouldCacheResult ?? (() => true)
 
       return await (withCache
-        ? withCache(queryHash, cacheStrategy, async ({addDebugData}) => {
-            if (process.env.NODE_ENV === 'development') {
-              // Name displayed in the subrequest profiler
-              const displayName = loaderOptions?.hydrogen?.debug?.displayName || 'query Sanity'
-
-              addDebugData({
-                displayName,
-              })
-            }
-
-            return await loadQuery<T>(query, params, loaderOptions)
-          })
+        ? withCache.run(
+            {
+              cacheKey: queryHash,
+              cacheStrategy,
+              shouldCacheResult,
+            },
+            async ({addDebugData}) => {
+              if (process.env.NODE_ENV === 'development') {
+                // Name displayed in the subrequest profiler
+                const displayName = loaderOptions?.hydrogen?.debug?.displayName || 'query Sanity'
+
+                addDebugData({
+                  displayName,
+                })
+              }
+
+              return await loadQuery<T>(query, params, loaderOptions)
+            },
+          )
         : loadQuery<T>(query, params, loaderOptions))
     },
     client,

package/src/loader.test.ts

@@ -33,8 +33,8 @@ function waitUntil() {
   return Promise.resolve()
 }
 
-type WithCache = ReturnType<typeof createWithCache>
-const withCache: WithCache = vi.fn().mockImplementation(createWithCache({cache, waitUntil}))
+const withCache = createWithCache({cache, waitUntil, request: new Request('https://example.com')})
+const runWithCache = vi.spyOn(withCache, 'run')
 
 const client = createClient({projectId: 'my-project-id', dataset: 'my-dataset'})
 
@@ -63,15 +63,23 @@ describe('the loader', () => {
     })
 
     await loaderWithDefaultStrategy.loadQuery<boolean>(query, params)
-    expect(withCache).toHaveBeenCalledWith(hashedQuery, strategy, expect.any(Function))
-    expect(cache.put).toHaveBeenCalled()
+    expect(runWithCache).toHaveBeenNthCalledWith(
+      1,
+      expect.objectContaining({cacheKey: hashedQuery, cacheStrategy: strategy}),
+      expect.any(Function),
+    )
+    expect(cache.put).toHaveBeenCalledOnce()
   })
 
   it('queries should use the cache strategy passed in `loadQuery`', async () => {
     const strategy = CacheShort()
     await loader.loadQuery<boolean>(query, params, {hydrogen: {cache: strategy}})
-    expect(withCache).toHaveBeenCalledWith(hashedQuery, strategy, expect.any(Function))
-    expect(cache.put).toHaveBeenCalled()
+    expect(runWithCache).toHaveBeenNthCalledWith(
+      1,
+      expect.objectContaining({cacheKey: hashedQuery, cacheStrategy: strategy}),
+      expect.any(Function),
+    )
+    expect(cache.put).toHaveBeenCalledOnce()
   })
 })
 
@@ -108,6 +116,6 @@ describe('when configured for preview', () => {
   it(`shouldn't cache queries`, async () => {
     await previewLoader.loadQuery<boolean>(query)
     expect(loadQuery).toHaveBeenCalledOnce()
-    expect(cache.put).not.toHaveBeenCalled()
+    expect(cache.put).not.toHaveBeenCalledOnce()
   })
 })

package/src/loader.ts

@@ -61,10 +61,18 @@ type HydrogenResponseQueryOptions = Omit<ResponseQueryOptions, 'next' | 'cache'>
   hydrogen?: 'hydrogen' extends keyof RequestInit ? RequestInit['hydrogen'] : never
 }
 
-type LoadQueryOptions = Pick<
+type LoadQueryOptions<T> = Pick<
   HydrogenResponseQueryOptions,
   'perspective' | 'hydrogen' | 'useCdn' | 'stega' | 'headers' | 'tag'
->
+> & {
+  hydrogen?: {
+    /**
+     * Whether to cache the result of the query or not.
+     * @default () => true
+     */
+    shouldCacheResult?: (value: QueryResponseInitial<T>) => boolean
+  }
+}
 
 export type SanityLoader = {
   /**
@@ -75,7 +83,7 @@ export type SanityLoader = {
   loadQuery<T = any>(
     query: string,
     params?: QueryParams,
-    options?: LoadQueryOptions,
+    options?: LoadQueryOptions<T>,
   ): Promise<QueryResponseInitial<T>>
 
   client: SanityClient
@@ -123,7 +131,7 @@ export function createSanityLoader(options: CreateSanityLoaderOptions): SanityLo
     async loadQuery<T>(
       query: string,
       params: QueryParams | QueryWithoutParams,
-      loaderOptions?: LoadQueryOptions,
+      loaderOptions?: LoadQueryOptions<T>,
     ): Promise<QueryResponseInitial<T>> {
       // Don't store response if preview is enabled
       const cacheStrategy =
@@ -133,19 +141,28 @@ export function createSanityLoader(options: CreateSanityLoaderOptions): SanityLo
 
       const queryHash = await hashQuery(query, params)
 
-      return await withCache(queryHash, cacheStrategy, async ({addDebugData}) => {
-        // eslint-disable-next-line no-process-env
-        if (process.env.NODE_ENV === 'development') {
-          // Name displayed in the subrequest profiler
-          const displayName = loaderOptions?.hydrogen?.debug?.displayName || 'query Sanity'
-
-          addDebugData({
-            displayName,
-          })
-        }
-
-        return await loadQuery<T>(query, params, loaderOptions)
-      })
+      const shouldCacheResult = loaderOptions?.hydrogen?.shouldCacheResult ?? (() => true)
+
+      return await withCache.run(
+        {
+          cacheKey: queryHash,
+          cacheStrategy,
+          shouldCacheResult,
+        },
+        async ({addDebugData}) => {
+          // eslint-disable-next-line no-process-env
+          if (process.env.NODE_ENV === 'development') {
+            // Name displayed in the subrequest profiler
+            const displayName = loaderOptions?.hydrogen?.debug?.displayName || 'query Sanity'
+
+            addDebugData({
+              displayName,
+            })
+          }
+
+          return await loadQuery<T>(query, params, loaderOptions)
+        },
+      )
     },
     client,
     preview,