feat(mutation): add an optional key option to mutations

playground/src/pages/ecom/item/[id].vue

@@ -25,6 +25,7 @@ const itemAvailability = ref()
 watch(() => item.value?.availability, (value) => itemAvailability.value = value)
 
 const { mutate: bookProduct } = useMutation({
+  key: (product) => ['book-item', product.id],
   keys: (product) => [['items'], ['items', product.id]],
   mutation: async (product: ProductListItem) => {
     await delay(Math.random() * 1000 + 200)

src/entry-options.ts

@@ -0,0 +1,7 @@
+import type { EntryNodeKey } from './tree-map'
+import type { _ObjectFlat } from './utils'
+
+/**
+ * Key used to identify a query or a mutation. Always an array.
+ */
+export type EntryKey = Array<EntryNodeKey | _ObjectFlat>

src/index.ts

@@ -17,7 +17,10 @@ export { defineQuery } from './define-query'
 // export { type UseQueryKeyList } from './query-keys'
 
 export {
-  type UseQueryKey,
+  type EntryKey,
+} from './entry-options'
+
+export {
   type UseQueryOptions,
   type UseQueryOptionsWithDefaults,
 } from './query-options'

src/plugins/invalidate-queries.ts

@@ -1,16 +1,16 @@
 /**
  * @module @pinia/colada/plugins/invalidate-queries
  */
-import type { UseQueryKey } from '@pinia/colada'
+import type { EntryKey } from '@pinia/colada'
 
 export interface UseMutationOptionsInvalidateQueries<TResult, TVars> {
   /**
    * Keys to invalidate if the mutation succeeds so that `useQuery()` refetch if used. Only available if
    * `@pinia/colada/plugins/invalidate-queries` is installed.
    */
   invalidateKeys?:
-    | UseQueryKey[]
-    | ((data: TResult, vars: TVars) => UseQueryKey[])
+    | EntryKey[]
+    | ((data: TResult, vars: TVars) => EntryKey[])
 }
 
 declare module '@pinia/colada' {

src/query-options.ts

@@ -1,6 +1,5 @@
 import { type InjectionKey, type MaybeRefOrGetter, inject } from 'vue'
-import type { EntryNodeKey } from './tree-map'
-import type { _ObjectFlat } from './utils'
+import type { EntryKey } from './entry-options'
 import type { QueryPluginOptions } from './query-plugin'
 import type { ErrorDefault } from './types-extension'
 
@@ -9,11 +8,6 @@ import type { ErrorDefault } from './types-extension'
  */
 export type _RefetchOnControl = boolean | 'always'
 
-/**
- * Key used to identify a query. Always an array.
- */
-export type UseQueryKey = Array<EntryNodeKey | _ObjectFlat>
-
 /**
  * Context object passed to the `query` function of `useQuery()`.
  * @see {@link UseQueryOptions}
@@ -59,7 +53,7 @@ export interface UseQueryOptions<TResult = unknown, TError = ErrorDefault> {
    * })
    * ```
    */
-  key: MaybeRefOrGetter<UseQueryKey>
+  key: MaybeRefOrGetter<EntryKey>
 
   /**
    * The function that will be called to fetch the data. It **must** be async.

src/query-store.ts

@@ -12,7 +12,8 @@
 } from 'vue'
 import { stringifyFlatObject } from './utils'
 import { type EntryNodeKey, TreeMapNode } from './tree-map'
-import type { UseQueryKey, UseQueryOptionsWithDefaults } from './query-options'
+import type { EntryKey } from './entry-options'
+import type { UseQueryOptionsWithDefaults } from './query-options'
 import type { ErrorDefault } from './types-extension'
 import type { defineQuery } from './define-query'
 
@@ -152,7 +153,7 @@
   }
 
   function ensureEntry<TResult = unknown, TError = ErrorDefault>(
-    keyRaw: UseQueryKey,
+    keyRaw: EntryKey,
     options: UseQueryOptionsWithDefaults<TResult, TError>,
   ): UseQueryEntry<TResult, TError> {
     if (process.env.NODE_ENV !== 'production' && keyRaw.length === 0) {
@@ -187,7 +188,7 @@
    * @param options.refetch - if true, it will refetch the data
    */
   function invalidateEntry(
-    key: UseQueryKey,
+    key: EntryKey,
     {
       refetch: shouldRefetch = true,
       exact = false,
@@ -310,7 +311,7 @@
 
   // TODO: tests, remove function version
   function setQueryData<TResult = unknown>(
-    key: UseQueryKey,
+    key: EntryKey,
     data: TResult | ((data: Ref<TResult | undefined>) => void),
   ) {
     const entry = caches.get(key.map(stringifyFlatObject)) as
@@ -331,7 +332,7 @@
   }
 
   function getQueryData<TResult = unknown>(
-    key: UseQueryKey,
+    key: EntryKey,
   ): TResult | undefined {
     const entry = caches.get(key.map(stringifyFlatObject)) as
       | UseQueryEntry<TResult>

src/use-mutation.test-d.ts

@@ -2,6 +2,16 @@ import { expectTypeOf, it } from 'vitest'
 import { useMutation } from './use-mutation'
 
 it('types the parameters for the key', () => {
+  useMutation({
+    mutation: (_one: string) => Promise.resolve({ name: 'foo' }),
+    key(one) {
+      expectTypeOf(one).toBeString()
+      return ['foo']
+    },
+  })
+})
+
+it('types the parameters for the keys', () => {
   const { mutate, mutateAsync } = useMutation({
     mutation: (_one: string) => Promise.resolve({ name: 'foo' }),
     keys(result, one) {

src/use-mutation.ts

@@ -1,18 +1,22 @@
 import { computed, shallowRef } from 'vue'
 import type { ComputedRef, ShallowRef } from 'vue'
 import { useQueryCache } from './query-store'
-import type { UseQueryKey } from './query-options'
+import type { EntryKey } from './entry-options'
 import type { ErrorDefault } from './types-extension'
 import { type _Awaitable, noop } from './utils'
 
+type _MutationKey<TVars> =
+  | EntryKey
+  | ((vars: TVars) => EntryKey)
+
 // TODO: move to a plugin
 /**
  * The keys to invalidate when a mutation succeeds.
  * @internal
  */
 type _MutationKeys<TVars, TResult> =
-  | UseQueryKey[]
-  | ((data: TResult, vars: TVars) => UseQueryKey[])
+  | EntryKey[]
+  | ((data: TResult, vars: TVars) => EntryKey[])
 
 /**
  * The status of the mutation.
@@ -69,6 +73,8 @@ export interface UseMutationOptions<
    */
   mutation: (vars: TVars, context: NoInfer<TContext>) => Promise<TResult>
 
+  key?: _MutationKey<TVars>
+
   // TODO: move this to a plugin that calls invalidateEntry()
   /**
    * Keys to invalidate if the mutation succeeds so that `useQuery()` refetch if used.

src/use-query.ts

@@ -13,8 +13,8 @@
 import { IS_CLIENT, computedRef, useEventListener } from './utils'
 import { type _UseQueryEntry_State, useQueryCache } from './query-store'
 import { useQueryOptions } from './query-options'
+import type { EntryKey } from './entry-options'
 import type {
-  UseQueryKey,
   UseQueryOptions,
   UseQueryOptionsWithDefaults,
 } from './query-options'
@@ -176,9 +176,9 @@
  * @returns - the computed key
  */
 function _computedKeyWithWarnings(
-  key: MaybeRefOrGetter<UseQueryKey>,
+  key: MaybeRefOrGetter<EntryKey>,
   warnChecksMap: WeakMap<object, boolean>,
-): () => UseQueryKey {
+): () => EntryKey {
   const componentInstance = getCurrentInstance()
   // probably correct scope, no need to warn
   if (!componentInstance) return () => toValue(key)