feat: Add select option to extract items while preserving metadata (#551)

* feat: Add select option to extract items while preserving metadata

* test: add tests

* test: add more tests

* throw type mismatch error schema level

* chore: remove a trailing comment

* chore: add implementation examples

* doc: add select to options

* add changeset

* Improve select validation with better error messages and type tests

- Enhanced error messages to include package name, error type, received value type, and queryKey for easier debugging
- Added negative type test to verify TypeScript catches wrapped response data without select function
- Updated existing runtime tests to handle improved error messaging
- Fixed eslint warnings for unnecessary conditions

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

---------

Co-authored-by: Kyle Mathews <[email protected]>
Co-authored-by: Claude <[email protected]>

Author: 3 people

.changeset/yellow-pears-rush.md

@@ -0,0 +1,5 @@
+---
+"@tanstack/query-db-collection": patch
+---
+
+query-collection now supports a `select` function to transform raw query results into an array of items. This is useful for APIs that return data with metadata or nested structures, ensuring metadata remains cached while collections work with the unwrapped array.

docs/collections/query-collection.md

@@ -55,6 +55,7 @@ The `queryCollectionOptions` function accepts the following options:
 
 ### Query Options
 
+- `select`: Function that lets extract array items when they’re wrapped with metadata
 - `enabled`: Whether the query should automatically run (default: `true`)
 - `refetchInterval`: Refetch interval in milliseconds
 - `retry`: Retry configuration for failed queries

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

@@ -52,25 +52,25 @@ type InferSchemaInput<T> = T extends StandardSchemaV1
  */
 export interface QueryCollectionConfig<
   T extends object = object,
-  TQueryFn extends (
+  TQueryFn extends (context: QueryFunctionContext<any>) => Promise<any> = (
     context: QueryFunctionContext<any>
-  ) => Promise<Array<any>> = (
-    context: QueryFunctionContext<any>
-  ) => Promise<Array<any>>,
+  ) => Promise<any>,
   TError = unknown,
   TQueryKey extends QueryKey = QueryKey,
   TKey extends string | number = string | number,
   TSchema extends StandardSchemaV1 = never,
+  TQueryData = Awaited<ReturnType<TQueryFn>>,
 > extends BaseCollectionConfig<T, TKey, TSchema> {
   /** The query key used by TanStack Query to identify this query */
   queryKey: TQueryKey
   /** Function that fetches data from the server. Must return the complete collection state */
   queryFn: TQueryFn extends (
     context: QueryFunctionContext<TQueryKey>
   ) => Promise<Array<any>>
-    ? TQueryFn
-    : (context: QueryFunctionContext<TQueryKey>) => Promise<Array<T>>
-
+    ? (context: QueryFunctionContext<TQueryKey>) => Promise<Array<T>>
+    : TQueryFn
+  /* Function that extracts array items from wrapped API responses (e.g metadata, pagination)  */
+  select?: (data: TQueryData) => Array<T>
   /** The TanStack Query client instance */
   queryClient: QueryClient
 
@@ -248,7 +248,77 @@ export interface QueryCollectionUtils<
  *     }
  *   })
  * )
+ *
+ * @example
+ * // The select option extracts the items array from a response with metadata
+ * const todosCollection = createCollection(
+ *   queryCollectionOptions({
+ *     queryKey: ['todos'],
+ *     queryFn: async () => fetch('/api/todos').then(r => r.json()),
+ *     select: (data) => data.items, // Extract the array of items
+ *     queryClient,
+ *     schema: todoSchema,
+ *     getKey: (item) => item.id,
+ *   })
+ * )
  */
+// Overload for when schema is provided and select present
+export function queryCollectionOptions<
+  T extends StandardSchemaV1,
+  TQueryFn extends (context: QueryFunctionContext<any>) => Promise<any>,
+  TError = unknown,
+  TQueryKey extends QueryKey = QueryKey,
+  TKey extends string | number = string | number,
+  TQueryData = Awaited<ReturnType<TQueryFn>>,
+>(
+  config: QueryCollectionConfig<
+    InferSchemaOutput<T>,
+    TQueryFn,
+    TError,
+    TQueryKey,
+    TKey,
+    T
+  > & {
+    schema: T
+    select: (data: TQueryData) => Array<InferSchemaInput<T>>
+  }
+): CollectionConfig<InferSchemaOutput<T>, TKey, T> & {
+  schema: T
+  utils: QueryCollectionUtils<
+    InferSchemaOutput<T>,
+    TKey,
+    InferSchemaInput<T>,
+    TError
+  >
+}
+
+// Overload for when no schema is provided and select present
+export function queryCollectionOptions<
+  T extends object,
+  TQueryFn extends (context: QueryFunctionContext<any>) => Promise<any> = (
+    context: QueryFunctionContext<any>
+  ) => Promise<any>,
+  TError = unknown,
+  TQueryKey extends QueryKey = QueryKey,
+  TKey extends string | number = string | number,
+  TQueryData = Awaited<ReturnType<TQueryFn>>,
+>(
+  config: QueryCollectionConfig<
+    T,
+    TQueryFn,
+    TError,
+    TQueryKey,
+    TKey,
+    never,
+    TQueryData
+  > & {
+    schema?: never // prohibit schema
+    select: (data: TQueryData) => Array<T>
+  }
+): CollectionConfig<T, TKey> & {
+  schema?: never // no schema in the result
+  utils: QueryCollectionUtils<T, TKey, T, TError>
+}
 
 // Overload for when schema is provided
 export function queryCollectionOptions<
@@ -308,6 +378,7 @@ export function queryCollectionOptions(
   const {
     queryKey,
     queryFn,
+    select,
     queryClient,
     enabled,
     refetchInterval,
@@ -390,16 +461,18 @@ export function queryCollectionOptions(
         lastError = undefined
         errorCount = 0
 
-        const newItemsArray = result.data
+        const rawData = result.data
+        const newItemsArray = select ? select(rawData) : rawData
 
         if (
           !Array.isArray(newItemsArray) ||
           newItemsArray.some((item) => typeof item !== `object`)
         ) {
-          console.error(
-            `[QueryCollection] queryFn did not return an array of objects. Skipping update.`,
-            newItemsArray
-          )
+          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
         }
 

packages/query-db-collection/tests/query.test-d.ts

@@ -263,10 +263,10 @@ describe(`Query collection type resolution tests`, () => {
       const options = queryCollectionOptions({
         queryClient,
         queryKey: [`schema-priority`],
-        // @ts-expect-error – queryFn doesn't match the schema type
         queryFn: async () => {
           return [] as Array<UserType>
         },
+        // @ts-expect-error – queryFn doesn't match the schema type
         schema: userSchema,
         getKey: (item) => item.id,
       })
@@ -301,4 +301,106 @@ describe(`Query collection type resolution tests`, () => {
       expectTypeOf(collection.toArray).toEqualTypeOf<Array<TodoType>>()
     })
   })
+
+  describe(`select type inference`, () => {
+    it(`queryFn type inference`, () => {
+      const dataSchema = z.object({
+        id: z.string(),
+        name: z.string(),
+        email: z.string(),
+      })
+
+      const options = queryCollectionOptions({
+        queryClient,
+        queryKey: [`x-queryFn-infer`],
+        queryFn: async (): Promise<Array<z.infer<typeof dataSchema>>> => {
+          return [] as Array<z.infer<typeof dataSchema>>
+        },
+        schema: dataSchema,
+        getKey: (item) => item.id,
+      })
+
+      type ExpectedType = z.infer<typeof dataSchema>
+      expectTypeOf(options.getKey).parameters.toEqualTypeOf<[ExpectedType]>()
+    })
+
+    it(`should error when queryFn returns wrapped data without select`, () => {
+      const userData = z.object({
+        id: z.string(),
+        name: z.string(),
+        email: z.string(),
+      })
+
+      type UserDataType = z.infer<typeof userData>
+
+      type WrappedResponse = {
+        metadata: string
+        data: Array<UserDataType>
+      }
+
+      queryCollectionOptions({
+        queryClient,
+        queryKey: [`wrapped-no-select`],
+        // @ts-expect-error - queryFn returns wrapped data but no select provided
+        queryFn: (): Promise<WrappedResponse> => {
+          return Promise.resolve({
+            metadata: `example`,
+            data: [],
+          })
+        },
+        // @ts-expect-error - schema type conflicts with queryFn return type
+        schema: userData,
+        // @ts-expect-error - item type is inferred as object due to type mismatch
+        getKey: (item) => item.id,
+      })
+    })
+
+    it(`select properly extracts array from wrapped response`, () => {
+      const userData = z.object({
+        id: z.string(),
+        name: z.string(),
+        email: z.string(),
+      })
+
+      type UserDataType = z.infer<typeof userData>
+
+      type MetaDataType<T> = {
+        metaDataOne: string
+        metaDataTwo: string
+        data: T
+      }
+
+      const metaDataObject: ResponseType = {
+        metaDataOne: `example meta data`,
+        metaDataTwo: `example meta data`,
+        data: [
+          {
+            id: `1`,
+            name: `carter`,
+            email: `[email protected]`,
+          },
+        ],
+      }
+
+      type ResponseType = MetaDataType<Array<UserDataType>>
+
+      const selectUserData = (data: ResponseType) => {
+        return data.data
+      }
+
+      queryCollectionOptions({
+        queryClient,
+        queryKey: [`x-queryFn-infer`],
+        queryFn: async (): Promise<ResponseType> => {
+          return metaDataObject
+        },
+        select: selectUserData,
+        schema: userData,
+        getKey: (item) => item.id,
+      })
+
+      // Should infer ResponseType as select parameter type
+      expectTypeOf(selectUserData).parameters.toEqualTypeOf<[ResponseType]>()
+    })
+  })
 })