Tag type discrimination POC

src/createQueryClient.spec-d.ts

@@ -25,22 +25,26 @@ describe('query', () => {
 
     test('tags', async () => {
       const action = () => 'response'
-      const { query } = createQueryClient()
-      const numberTag = tag<number>()
+      const { query, setQueryData } = createQueryClient()
       const stringTag = tag<string>()
       const untypedTag = tag()
 
-      // @ts-expect-error - number tag not assignable to string action
-      query(action, [], { tags: [numberTag, stringTag] })
-
-      // @ts-expect-error - number tag not assignable to string action
-      query(action, [], { tags: () => [numberTag, stringTag] })
-
       query(action, [], { tags: [stringTag, untypedTag] })
       query(action, [], { tags: () => [stringTag, untypedTag] })
 
       query(action, [], { tags: [untypedTag] })
       query(action, [], { tags: () => [untypedTag] })
+
+      setQueryData(stringTag, (data) => {
+        expectTypeOf(data).toEqualTypeOf<string>()
+        return data + 'bar'
+      })
+
+      // @ts-expect-error - sharedTag has data: never, can't return anything useful
+      setQueryData(untypedTag, (data) => {
+        expectTypeOf(data).toEqualTypeOf<never>()
+        return 'could be corrupting'
+      })
     })
   })
 })
@@ -99,59 +103,78 @@ describe('defineQuery', () => {
 })
 
 describe('setQueryData', () => {
-  test('tags', async () => {
+  test('naked tag is not setQueryData-able (data: never)', () => {
+    const { setQueryData } = createQueryClient()
+    const myTag = tag()
+
+    // @ts-expect-error - data: never, return must be never (effectively impossible)
+    setQueryData(myTag, () => 'anything')
+  })
+
+  test('typed tag passes data through with its declared type', () => {
     const { setQueryData } = createQueryClient()
     const numberTag = tag<number>()
-    const stringTag = tag<string>()
-    const untypedTag = tag()
 
-    setQueryData(untypedTag, (data) => {
-      expectTypeOf(data).toEqualTypeOf<unknown>()
-      return 'foo'
+    setQueryData(numberTag, (data) => {
+      expectTypeOf(data).toEqualTypeOf<number>()
+      return data + 1
     })
 
+    // @ts-expect-error - returning wrong type
     setQueryData(numberTag, (data) => {
       expectTypeOf(data).toEqualTypeOf<number>()
-      return 2
+      return 'wrong type'
     })
+  })
 
-    setQueryData(stringTag, (data) => {
-      expectTypeOf(data).toEqualTypeOf<string>()
-      return 'new string'
-    })
+  test('can type with single kind as well', () => {
+    const { setQueryData } = createQueryClient()
+    const usersTag = tag<{ user: { id: number } }>(['user'])
 
-    setQueryData([untypedTag], (data) => {
-      expectTypeOf(data).toEqualTypeOf<unknown>()
-      return 'foo'
+    setQueryData(usersTag.user, (data) => {
+      expectTypeOf(data).toEqualTypeOf<{ id: number }>()
+      return data
     })
+  })
 
-    setQueryData([numberTag], (data) => {
-      expectTypeOf(data).toEqualTypeOf<number>()
-      return 2
-    })
+  test('object handler on parent typed per kind', () => {
+    const { setQueryData } = createQueryClient()
 
-    // 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'
-    })
+    type User = { id: string, name: string, email: string }
+    type Potato = { genus: string, species: string }
+
+    const genericTag = tag<{ user: User, potato: Potato }>(['user', 'potato'])
 
-    setQueryData([untypedTag, stringTag, numberTag], (data) => {
-      expectTypeOf(data).toEqualTypeOf<unknown>()
-      return 'foo'
+    setQueryData(genericTag, {
+      user: (data) => {
+        expectTypeOf(data).toEqualTypeOf<User>()
+        return data
+      },
+      potato: (data) => {
+        expectTypeOf(data).toEqualTypeOf<Potato>()
+        return data
+      },
     })
 
-    // @ts-expect-error - number tag not assignable to string action
-    setQueryData(numberTag, (data) => {
-      expectTypeOf(data).toEqualTypeOf<number>()
-      return 'string'
+    // function form on parent is locked when kinds are disjoint:
+    // setter takes UnionToIntersection<User | Potato> = never
+    // @ts-expect-error - return type cannot satisfy never
+    setQueryData(genericTag, () => ({ id: '123', name: 'John', email: 'a@b' }))
+
+    // function form on a specific kind works
+    setQueryData(genericTag.user, (data) => {
+      expectTypeOf(data).toEqualTypeOf<User>()
+      return data
     })
+  })
 
-    // @ts-expect-error - number tag not assignable to string action
-    setQueryData([numberTag, stringTag], (data) => {
-      expectTypeOf(data).toEqualTypeOf<number | string>()
-      return []
+  test('object handler must cover every declared kind', () => {
+    const { setQueryData } = createQueryClient()
+    const genericTag = tag<{ a: number, b: string }>(['a', 'b'])
+
+    // @ts-expect-error - missing 'b' kind in handler
+    setQueryData(genericTag, {
+      a: (data) => data,
     })
   })
 
@@ -252,12 +275,12 @@ describe('refreshQueryData', () => {
   test('tags', () => {
     const { refreshQueryData } = createQueryClient()
 
-    const numberTag = tag<number>()
-    const stringTag = tag<string>()
+    const sharedTag = tag<{ count: number, name: string }>(['count', 'name'])
     const action = (param: number) => param
 
-    refreshQueryData(numberTag)
-    refreshQueryData([numberTag, stringTag])
+    refreshQueryData(sharedTag)
+    refreshQueryData(sharedTag.count)
+    refreshQueryData(sharedTag.name)
     refreshQueryData(action)
     refreshQueryData(action, [2])
 

src/createQueryClient.spec.ts

@@ -562,41 +562,50 @@ describe('setQueryData', () => {
 
     await vi.runOnlyPendingTimersAsync()
 
-    setQueryData(stringTag, () => {
-      return 'bar'
-    })
-
-    setQueryData(numberTag, () => {
-      return 2
-    })
+    setQueryData(stringTag, () => 'bar')
+    setQueryData(numberTag, () => 2)
 
     expect(stringQuery.data).toBe('bar')
     expect(numberQuery.data).toBe(2)
   })
 
-  test('tags', async () => {
+  test('kind tag setter only matches that kind\'s queries', async () => {
     const { setQueryData, query } = createQueryClient()
-    const stringTag = tag<string>()
-    const numberTag = tag<number>()
+    const sharedTag = tag<{ name: string, count: number }>(['name', 'count'])
 
     const stringAction = () => 'foo'
     const numberAction = () => 1
 
-    const stringQuery = query(stringAction, [], { tags: [stringTag] })
-    const numberQuery = query(numberAction, [], { tags: [numberTag] })
+    const stringQuery = query(stringAction, [], { tags: [sharedTag.name] })
+    const numberQuery = query(numberAction, [], { tags: [sharedTag.count] })
 
     await vi.runOnlyPendingTimersAsync()
 
-    setQueryData([stringTag], () => {
-      return 'bar'
-    })
+    setQueryData(sharedTag.name, (data) => data + '-bar')
 
-    setQueryData([numberTag], () => {
-      return 2
+    expect(stringQuery.data).toBe('foo-bar')
+    expect(numberQuery.data).toBe(1)
+  })
+
+  test('object handler on parent dispatches to each kind', async () => {
+    const { setQueryData, query } = createQueryClient()
+    const sharedTag = tag<{ name: string, count: number }>(['name', 'count'])
+
+    const stringAction = () => 'foo'
+    const numberAction = () => 1
+
+    const stringQuery = query(stringAction, [], { tags: [sharedTag.name] })
+    const numberQuery = query(numberAction, [], { tags: [sharedTag.count] })
+
+    await vi.runOnlyPendingTimersAsync()
+
+    setQueryData(sharedTag, {
+      name: (data) => data + '-bar',
+      count: (data) => data + 10,
     })
 
-    expect(stringQuery.data).toBe('bar')
-    expect(numberQuery.data).toBe(2)
+    expect(stringQuery.data).toBe('foo-bar')
+    expect(numberQuery.data).toBe(11)
   })
 
   test('action', async () => {
@@ -656,8 +665,8 @@ describe('refreshQueryData', () => {
 
     const numberAction = vi.fn()
     const stringAction = vi.fn()
-    const numberTag = tag<number>()
-    const stringTag = tag<string>()
+    const numberTag = tag()
+    const stringTag = tag()
 
     query(numberAction, [], { tags: [numberTag] })
     query(stringAction, [], { tags: [stringTag] })
@@ -795,7 +804,7 @@ describe('mutate', () => {
     [undefined],
   ])('refreshes tagged queries: %s', async (refreshQueryData) => {
     const { mutate, query } = createQueryClient()
-    const numberTag = tag<number>()
+    const numberTag = tag()
     const queryAction = vi.fn()
     const mutationAction = vi.fn()
 
@@ -817,7 +826,7 @@ describe('mutate', () => {
 
   test('does not refresh tagged queries if refreshQueryData is false', async () => {
     const { mutate, query } = createQueryClient()
-    const numberTag = tag<number>()
+    const numberTag = tag()
     const queryAction = vi.fn()
     const mutationAction = vi.fn()
 
@@ -839,7 +848,7 @@ describe('mutate', () => {
 
   test('does not refresh tagged queries if the action throws an error', async () => {
     const { mutate, query } = createQueryClient()
-    const numberTag = tag<number>()
+    const numberTag = tag()
     const queryAction = vi.fn()
     const mutationAction = vi.fn(() => {
       throw new Error()
@@ -1047,7 +1056,7 @@ describe('useMutation', () => {
   test('setQueryDataBefore and setQueryDataAfter are called when the mutation is executed', async () => {
     const { useMutation, query } = createQueryClient()
     const { promise, resolve } = Promise.withResolvers<void>()
-    const numberTag = tag<number>()
+    const numberTag = tag<void>()
     const queryAction = vi.fn()
     const mutationAction = vi.fn(() => promise)
     const setQueryDataBefore = vi.fn()
@@ -1306,8 +1315,8 @@ describe('defineMutation', () => {
           const { promise, resolve } = Promise.withResolvers<void>()
           const mutationPayload = 1
           const mutationAction = (value: number) => promise.then(() => value)
-          const tagA = tag()
-          const tagB = tag()
+          const tagA = tag<number>()
+          const tagB = tag<number>()
           const queryAResponse = 1
           const queryBResponse = 1
           const queryAAction = () => queryAResponse
@@ -1627,8 +1636,8 @@ describe('defineMutation', () => {
           const { promise, resolve } = Promise.withResolvers<void>()
           const mutationPayload = 1
           const mutationAction = (value: number) => promise.then(() => value)
-          const tagA = tag()
-          const tagB = tag()
+          const tagA = tag<number>()
+          const tagB = tag<number>()
           const queryAResponse = 1
           const queryBResponse = 1
           const queryAAction = () => queryAResponse

src/createQueryClient.ts

@@ -13,7 +13,7 @@ import {
 } from './types/client'
 import { createQueryGroups } from './createQueryGroups'
 import { createUseQuery } from './createUseQuery'
-import { isQueryTag, isQueryTags, QueryTag } from './types/tags'
+import { isQueryTag, QueryTag } from './types/tags'
 import { isArray } from './utilities/arrays'
 import { assertNever } from './utilities/assert'
 import { QueryGroup } from './createQueryGroup'
@@ -59,8 +59,8 @@ export function createQueryClient(options?: ClientOptions): QueryClient {
   }
 
   const setQueryData: SetQueryData = (
-    param1: QueryTag | QueryTag[] | QueryAction,
-    param2: Parameters<QueryAction> | QueryDataSetter,
+    param1: QueryTag | QueryAction,
+    param2: Parameters<QueryAction> | QueryDataSetter | Record<string, QueryDataSetter>,
     param3?: QueryDataSetter,
   ): void => {
     const setDataForGroups = (groups: QueryGroup[], setter: QueryDataSetter): void => {
@@ -72,12 +72,25 @@ export function createQueryClient(options?: ClientOptions): QueryClient {
       })
     }
 
-    if (isQueryTag(param1) || isQueryTags(param1)) {
-      const tags = param1
-      const setter = param2 as QueryDataSetter
-      const groups = getQueryGroups(tags)
+    if (isQueryTag(param1)) {
+      const queryTag = param1
 
-      setDataForGroups(groups, setter)
+      if (typeof param2 === 'function') {
+        const setter = param2
+        const groups = getQueryGroups(queryTag)
+
+        setDataForGroups(groups, setter)
+      } else {
+        const handler = param2 as unknown as Record<string, QueryDataSetter>
+
+        for (const kind of queryTag.kinds) {
+          const childTag = (queryTag as unknown as Record<string, QueryTag>)[kind]
+          const setter = handler[kind]
+          const groups = getQueryGroups(childTag)
+
+          setDataForGroups(groups, setter)
+        }
+      }
 
       return
     }
@@ -107,12 +120,12 @@ export function createQueryClient(options?: ClientOptions): QueryClient {
   }
 
   const refreshQueryData: RefreshQueryData = (
-    param1: QueryTag | QueryTag[] | QueryAction,
+    param1: QueryTag | QueryAction,
     param2?: Parameters<QueryAction>,
   ): void => {
-    if (isQueryTag(param1) || isQueryTags(param1)) {
-      const tags = param1
-      const groups = getQueryGroups(tags)
+    if (isQueryTag(param1)) {
+      const queryTag = param1
+      const groups = getQueryGroups(queryTag)
 
       groups.forEach((group) => {
         group.execute()
@@ -133,7 +146,7 @@ export function createQueryClient(options?: ClientOptions): QueryClient {
       return
     }
 
-    assertNever(param1, 'Invalid arguments given to setQueryData')
+    assertNever(param1, 'Invalid arguments given to refreshQueryData')
   }
 
   const mutate: MutationFunction = (action, parameters, options) => {