Tag type discrimination POC

src/createQueryClient.spec-d.ts

@@ -25,9 +25,9 @@ describe('query', () => {
 
     test('tags', async () => {
       const action = () => 'response'
-      const { query } = createQueryClient()
-      const numberTag = tag<number>()
-      const stringTag = tag<string>()
+      const { query, setQueryData } = createQueryClient()
+      const numberTag = tag<number, 'count'>('count')
+      const stringTag = tag<string, 'name'>('name')
       const untypedTag = tag()
 
       // @ts-expect-error - number tag not assignable to string action
@@ -41,6 +41,11 @@ describe('query', () => {
 
       query(action, [], { tags: [untypedTag] })
       query(action, [], { tags: () => [untypedTag] })
+
+      setQueryData([numberTag, stringTag], {
+        count: (data) => data + 1,
+        name: (data) => data + 'bar',
+      })
     })
   })
 })
@@ -101,8 +106,8 @@ describe('defineQuery', () => {
 describe('setQueryData', () => {
   test('tags', async () => {
     const { setQueryData } = createQueryClient()
-    const numberTag = tag<number>()
-    const stringTag = tag<string>()
+    const numberTag = tag<number, 'count'>('count')
+    const stringTag = tag<string, 'name'>('name')
     const untypedTag = tag()
 
     setQueryData(untypedTag, (data) => {
@@ -130,28 +135,82 @@ describe('setQueryData', () => {
       return 2
     })
 
-    // this is kinda interesting, no matter the data the return type is the union :thinking:
-    // so there's not really a type safe way to update multiple queries at once
-    setQueryData([numberTag, stringTag], (data) => {
-      expectTypeOf(data).toEqualTypeOf<number | string>()
-      return 'foo'
+    // multiple tags with distinct kinds: simple callback collapses to never, must use object form
+    setQueryData([numberTag, stringTag], {
+      count: (data) => {
+        expectTypeOf(data).toEqualTypeOf<number>()
+        return data + 1
+      },
+      name: (data) => {
+        expectTypeOf(data).toEqualTypeOf<string>()
+        return data + 'bar'
+      },
     })
 
-    setQueryData([untypedTag, stringTag, numberTag], (data) => {
-      expectTypeOf(data).toEqualTypeOf<unknown>()
-      return 'foo'
+    // untyped tag has the default kind, so its handler is keyed by 'default'
+    setQueryData([untypedTag, stringTag, numberTag], {
+      default: (data) => {
+        expectTypeOf(data).toEqualTypeOf<unknown>()
+        return 'foo'
+      },
+      count: (data) => {
+        expectTypeOf(data).toEqualTypeOf<number>()
+        return data
+      },
+      name: (data) => {
+        expectTypeOf(data).toEqualTypeOf<string>()
+        return data
+      },
     })
 
-    // @ts-expect-error - number tag not assignable to string action
+    // @ts-expect-error - number tag with string return
     setQueryData(numberTag, (data) => {
       expectTypeOf(data).toEqualTypeOf<number>()
       return 'string'
     })
 
-    // @ts-expect-error - number tag not assignable to string action
-    setQueryData([numberTag, stringTag], (data) => {
-      expectTypeOf(data).toEqualTypeOf<number | string>()
-      return []
+    // @ts-expect-error - object handler returning wrong type for kind
+    setQueryData([numberTag, stringTag], {
+      count: () => 'wrong',
+      name: (data) => data,
+    })
+
+    // @ts-expect-error - missing handler for one of the kinds
+    setQueryData([numberTag, stringTag], {
+      count: (data) => data,
+    })
+  })
+
+  test('tags with shared kind', () => {
+    const { setQueryData } = createQueryClient()
+    const userTagA = tag<{ id: number }, 'user'>('user')
+    const userTagB = tag<{ id: number }, 'user'>('user')
+
+    // same kind, same type — simple callback works
+    setQueryData([userTagA, userTagB], (data) => {
+      expectTypeOf(data).toEqualTypeOf<{ id: number }>()
+      return data
+    })
+
+    // same kind, same type — object form also works
+    setQueryData([userTagA, userTagB], {
+      user: (data) => {
+        expectTypeOf(data).toEqualTypeOf<{ id: number }>()
+        return data
+      },
+    })
+  })
+
+  test('tags with same kind but different types collapse data to never', () => {
+    const { setQueryData } = createQueryClient()
+    const aTag = tag<number, 'shared'>('shared')
+    const bTag = tag<string, 'shared'>('shared')
+
+    setQueryData([aTag, bTag], {
+      shared: (data) => {
+        expectTypeOf(data).toEqualTypeOf<never>()
+        return data
+      },
     })
   })
 

src/createQueryClient.spec.ts

@@ -599,6 +599,28 @@ describe('setQueryData', () => {
     expect(numberQuery.data).toBe(2)
   })
 
+  test('tags with object setter dispatches by kind', async () => {
+    const { setQueryData, query } = createQueryClient()
+    const stringTag = tag<string, 'name'>('name')
+    const numberTag = tag<number, 'count'>('count')
+
+    const stringAction = () => 'foo'
+    const numberAction = () => 1
+
+    const stringQuery = query(stringAction, [], { tags: [stringTag] })
+    const numberQuery = query(numberAction, [], { tags: [numberTag] })
+
+    await vi.runOnlyPendingTimersAsync()
+
+    setQueryData([stringTag, numberTag], {
+      name: (data) => data + '-bar',
+      count: (data) => data + 10,
+    })
+
+    expect(stringQuery.data).toBe('foo-bar')
+    expect(numberQuery.data).toBe(11)
+  })
+
   test('action', async () => {
     const { setQueryData, query } = createQueryClient()
 

src/createQueryClient.ts

@@ -60,7 +60,7 @@ export function createQueryClient(options?: ClientOptions): QueryClient {
 
   const setQueryData: SetQueryData = (
     param1: QueryTag | QueryTag[] | QueryAction,
-    param2: Parameters<QueryAction> | QueryDataSetter,
+    param2: Parameters<QueryAction> | QueryDataSetter | Record<string, QueryDataSetter>,
     param3?: QueryDataSetter,
   ): void => {
     const setDataForGroups = (groups: QueryGroup[], setter: QueryDataSetter): void => {
@@ -73,11 +73,19 @@ export function createQueryClient(options?: ClientOptions): QueryClient {
     }
 
     if (isQueryTag(param1) || isQueryTags(param1)) {
-      const tags = param1
-      const setter = param2 as QueryDataSetter
-      const groups = getQueryGroups(tags)
-
-      setDataForGroups(groups, setter)
+      const tags = isArray(param1) ? param1 : [param1]
+      const setter = param2 as QueryDataSetter | Record<string, QueryDataSetter>
+
+      if (typeof setter === 'function') {
+        const groups = getQueryGroups(tags)
+        setDataForGroups(groups, setter)
+      } else {
+        for (const tag of tags) {
+          const handler = setter[tag.kind]
+          const groups = getQueryGroups(tag)
+          setDataForGroups(groups, handler)
+        }
+      }
 
       return
     }

src/tag.spec-d.ts

@@ -16,23 +16,35 @@ test('tag function returns a tag factory when a callback is provided', () => {
 
   const value = factory('foo')
 
-  expectTypeOf(value).toEqualTypeOf<QueryTag<Unset>>()
+  expectTypeOf(value).toEqualTypeOf<QueryTag<Unset, 'default'>>()
 })
 
 test('tag function returns a typed tag when data generic is provided', () => {
   const value = tag<string>()
 
-  expectTypeOf(value).toEqualTypeOf<QueryTag<string>>()
+  expectTypeOf(value).toEqualTypeOf<QueryTag<string, 'default'>>()
 })
 
 test('tag factory returns a typed tag when data generic is provided', () => {
   const factory = tag<string, string>((value: string) => value)
 
-  expectTypeOf(factory).toEqualTypeOf<QueryTagFactory<string, string>>()
+  expectTypeOf(factory).toEqualTypeOf<QueryTagFactory<string, string, 'default'>>()
 
   const value = factory('foo')
 
-  expectTypeOf(value).toEqualTypeOf<QueryTag<string>>()
+  expectTypeOf(value).toEqualTypeOf<QueryTag<string, 'default'>>()
+})
+
+test('tag function preserves kind literal', () => {
+  const value = tag('count')
+
+  expectTypeOf(value).toEqualTypeOf<QueryTag<Unset, 'count'>>()
+})
+
+test('tag function preserves kind literal with explicit data and kind generics', () => {
+  const value = tag<number, 'count'>('count')
+
+  expectTypeOf(value).toEqualTypeOf<QueryTag<number, 'count'>>()
 })
 
 test('query from query function with tags callback is called with the query data', () => {

src/tag.ts

@@ -1,24 +1,30 @@
 import { createSequence } from './createSequence'
 import { getTagKey } from './getTagKey'
-import { QueryTagFactory, QueryTagCallback, QueryTag, Unset, unset } from './types/tags'
+import { QueryTagFactory, QueryTagCallback, QueryTag, Unset, unset, DEFAULT_TAG_KIND, DefaultTagKind } from './types/tags'
 
 const createTagId = createSequence()
 
-function createQueryTag(id: number, value: unknown): QueryTag {
+function createQueryTag(id: number, kind: string, value: unknown): QueryTag {
   return {
     data: unset,
+    kind,
     key: getTagKey(id, value),
-  }
+  } as QueryTag
 }
 
-export function tag<const TData = Unset>(): QueryTag<TData>
-export function tag<const TData = Unset, TInput = unknown>(callback: QueryTagCallback<TInput>): QueryTagFactory<TData, TInput>
-export function tag(callback?: QueryTagCallback): QueryTag | QueryTagFactory {
+export function tag<const TData = Unset, const TKind extends string = DefaultTagKind>(kind?: TKind): QueryTag<TData, TKind>
+export function tag<const TData = Unset, TInput = unknown, const TKind extends string = DefaultTagKind>(callback: QueryTagCallback<TInput>, kind?: TKind): QueryTagFactory<TData, TInput, TKind>
+export function tag(callbackOrKind?: QueryTagCallback | string, maybeKind?: string): QueryTag | QueryTagFactory {
   const id = createTagId()
 
-  if (callback) {
-    return (value) => createQueryTag(id, callback(value))
+  if (typeof callbackOrKind === 'function') {
+    const callback = callbackOrKind
+    const kind = maybeKind ?? DEFAULT_TAG_KIND
+
+    return (value) => createQueryTag(id, kind, callback(value))
   }
 
-  return createQueryTag(id, undefined)
+  const kind = callbackOrKind ?? DEFAULT_TAG_KIND
+
+  return createQueryTag(id, kind, undefined)
 }

src/types/client.ts

@@ -1,7 +1,7 @@
 import { MutationFunction, MutationComposition, DefineMutation } from './mutation'
 import { Query, QueryOptions, QueryAction, QueryActionArgs, QueryData } from './query'
-import { QueryTag, QueryTagType } from './tags'
-import { DefaultValue } from './utilities'
+import { QueryTag, QueryTagType, QueryTagKind } from './tags'
+import { DefaultValue, UnionToIntersection } from './utilities'
 
 export type QueryClient = {
   query: QueryFunction,
@@ -62,8 +62,18 @@ export type DefinedQuery<
 
 export type QueryDataSetter<T = unknown> = (data: T) => T
 
+type SetQueryDataSimpleData<TQueryTag extends QueryTag> =
+  UnionToIntersection<TQueryTag extends any ? QueryTagType<TQueryTag> : never>
+
+type SetQueryDataKindData<TQueryTag extends QueryTag, TKind extends string> =
+  UnionToIntersection<TQueryTag extends QueryTag<any, TKind> ? QueryTagType<TQueryTag> : never>
+
+export type SetQueryDataValue<TQueryTag extends QueryTag> =
+  | QueryDataSetter<SetQueryDataSimpleData<TQueryTag>>
+  | { [K in QueryTagKind<TQueryTag>]: QueryDataSetter<SetQueryDataKindData<TQueryTag, K>> }
+
 export type SetQueryData = {
-  <TQueryTag extends QueryTag>(tag: TQueryTag | TQueryTag[], setter: QueryDataSetter<QueryTagType<TQueryTag>>): void,
+  <TQueryTag extends QueryTag>(tag: TQueryTag | TQueryTag[], setter: SetQueryDataValue<TQueryTag>): void,
   <TAction extends QueryAction>(action: TAction, setter: QueryDataSetter<QueryData<TAction>>): void,
   <TAction extends QueryAction>(action: TAction, parameters: Parameters<TAction>, setter: QueryDataSetter<QueryData<TAction>>): void,
 }

src/types/tags.ts

@@ -3,26 +3,35 @@ import { TagKey } from '@/getTagKey'
 export const unset = Symbol('unset')
 export type Unset = typeof unset
 
+export const DEFAULT_TAG_KIND = 'default'
+export type DefaultTagKind = typeof DEFAULT_TAG_KIND
+
 export type QueryTag<
-  TData = unknown
+  TData = unknown,
+  TKind extends string = string
 > = {
   /**
    * @private
    * @internal
    * This property is unused, but necessary to preserve the type for TData because unused generics are ignored by typescript.
    */
   data: TData,
+  kind: TKind,
   key: TagKey,
 }
 
-export type QueryTagType<TQueryTag extends QueryTag> = TQueryTag extends QueryTag<infer TData>
+export type QueryTagType<TQueryTag extends QueryTag> = TQueryTag extends QueryTag<infer TData, any>
   ? TData extends Unset
     ? unknown
     : TData
   : never
 
+export type QueryTagKind<TQueryTag extends QueryTag> = TQueryTag extends QueryTag<any, infer TKind>
+  ? TKind
+  : never
+
 export function isQueryTag(tag: unknown): tag is QueryTag {
-  return typeof tag === 'object' && tag !== null && 'data' in tag && 'key' in tag
+  return typeof tag === 'object' && tag !== null && 'data' in tag && 'kind' in tag && 'key' in tag
 }
 
 export function isQueryTags(tags: unknown): tags is QueryTag[] {
@@ -35,5 +44,6 @@ export type QueryTagCallback<
 
 export type QueryTagFactory<
   TData = unknown,
-  TInput = unknown
-> = (value: TInput) => QueryTag<TData>
+  TInput = unknown,
+  TKind extends string = string
+> = (value: TInput) => QueryTag<TData, TKind>

src/types/utilities.ts

@@ -1 +1,4 @@
 export type DefaultValue<TValue, TDefault> = unknown extends TValue ? TDefault : TValue
+
+export type UnionToIntersection<U> =
+  (U extends any ? (k: U) => void : never) extends (k: infer I) => void ? I : never