split enhanceEndpoints into addTagTypes and enhanceEndpoint

packages/toolkit/src/query/apiTypes.ts

@@ -3,16 +3,21 @@ import type {
   EndpointBuilder,
   EndpointDefinition,
   ReplaceTagTypes,
+  ResultTypeFrom,
+  QueryArgFrom,
+  QueryDefinition,
+  MutationDefinition,
 } from './endpointDefinitions'
 import type {
   UnionToIntersection,
   NoInfer,
   WithRequiredProp,
+  OverrideAtKey,
 } from './tsHelpers'
 import type { CoreModule } from './core/module'
 import type { CreateApiOptions } from './createApi'
 import type { BaseQueryFn } from './baseQueryTypes'
-import type { CombinedState } from './core/apiState'
+import type { CombinedState, MutationKeys, QueryKeys } from './core/apiState'
 import type { AnyAction } from '@reduxjs/toolkit'
 
 export interface ApiModules<
@@ -92,6 +97,8 @@ export type Api<
   >
   /**
    *A function to enhance a generated API with additional information. Useful with code-generation.
+
+   @deprecated this will be replaced by `addTagTypes` and `enhanceEndpoint`
    */
   enhanceEndpoints<NewTagTypes extends string = never>(_: {
     addTagTypes?: readonly NewTagTypes[]
@@ -112,4 +119,65 @@ export type Api<
     TagTypes | NewTagTypes,
     Enhancers
   >
+
+  /**
+   *A function to enhance a generated API with additional information. Useful with code-generation.
+   */
+  addTagTypes<NewTagTypes extends string = never>(
+    ...addTagTypes: readonly NewTagTypes[]
+  ): Api<
+    BaseQuery,
+    ReplaceTagTypes<Definitions, TagTypes | NewTagTypes>,
+    ReducerPath,
+    TagTypes | NewTagTypes,
+    Enhancers
+  >
+
+  /**
+   *A function to enhance a generated API with additional information. Useful with code-generation.
+   */
+  enhanceEndpoint<
+    QueryName extends QueryKeys<Definitions>,
+    ResultType = ResultTypeFrom<Definitions[QueryName]>,
+    QueryArg = QueryArgFrom<Definitions[QueryName]>
+  >(
+    queryName: QueryName,
+    queryDefinition: Partial<
+      QueryDefinition<QueryArg, BaseQuery, TagTypes, ResultType, ReducerPath>
+    >
+  ): Api<
+    BaseQuery,
+    OverrideAtKey<
+      Definitions,
+      QueryName,
+      QueryDefinition<QueryArg, BaseQuery, TagTypes, ResultType, ReducerPath>
+    >,
+    ReducerPath,
+    TagTypes,
+    Enhancers
+  >
+
+  /**
+   *A function to enhance a generated API with additional information. Useful with code-generation.
+   */
+  enhanceEndpoint<
+    MutationName extends MutationKeys<Definitions>,
+    ResultType = ResultTypeFrom<Definitions[MutationName]>,
+    QueryArg = QueryArgFrom<Definitions[MutationName]>
+  >(
+    mutationName: MutationName,
+    mutationDefinition: Partial<
+      MutationDefinition<QueryArg, BaseQuery, TagTypes, ResultType, ReducerPath>
+    >
+  ): Api<
+    BaseQuery,
+    OverrideAtKey<
+      Definitions,
+      MutationName,
+      MutationDefinition<QueryArg, BaseQuery, TagTypes, ResultType, ReducerPath>
+    >,
+    ReducerPath,
+    TagTypes,
+    Enhancers
+  >
 }

packages/toolkit/src/query/createApi.ts

@@ -275,13 +275,24 @@ export function buildCreateApi<Modules extends [Module<any>, ...Module<any>[]]>(
 
     const api = {
       injectEndpoints,
+      addTagTypes(...addTagTypes) {
+        for (const eT of addTagTypes) {
+          if (!optionsWithDefaults.tagTypes!.includes(eT as any)) {
+            ;(optionsWithDefaults.tagTypes as any[]).push(eT)
+          }
+        }
+        return api
+      },
+      enhanceEndpoint(endpointName, partialDefinition) {
+        Object.assign(
+          context.endpointDefinitions[endpointName] || {},
+          partialDefinition
+        )
+        return api
+      },
       enhanceEndpoints({ addTagTypes, endpoints }) {
         if (addTagTypes) {
-          for (const eT of addTagTypes) {
-            if (!optionsWithDefaults.tagTypes!.includes(eT as any)) {
-              ;(optionsWithDefaults.tagTypes as any[]).push(eT)
-            }
-          }
+          api.addTagTypes(...addTagTypes)
         }
         if (endpoints) {
           for (const [endpointName, partialDefinition] of Object.entries(
@@ -290,8 +301,9 @@ export function buildCreateApi<Modules extends [Module<any>, ...Module<any>[]]>(
             if (typeof partialDefinition === 'function') {
               partialDefinition(context.endpointDefinitions[endpointName])
             } else {
-              Object.assign(
-                context.endpointDefinitions[endpointName] || {},
+              // @ts-expect-error Type 'string' does not satisfy the constraint 'never'.
+              api.enhanceEndpoint<string, unknown, unknown>(
+                endpointName,
                 partialDefinition
               )
             }

packages/toolkit/src/query/tests/codeSplitting.test.ts

@@ -0,0 +1,218 @@
+import { createApi, fetchBaseQuery } from '@reduxjs/toolkit/query'
+import { expectType } from './helpers'
+
+function buildInitialApi() {
+  return createApi({
+    baseQuery: fetchBaseQuery(),
+    endpoints(build) {
+      return {
+        someQuery: build.query<'ReturnedFromQuery', 'QueryArgument'>({
+          query() {
+            return '/'
+          },
+        }),
+        someMutation: build.query<'ReturnedFromMutation', 'MutationArgument'>({
+          query() {
+            return '/'
+          },
+        }),
+      }
+    },
+  })
+}
+let baseApi = buildInitialApi()
+beforeEach(() => {
+  baseApi = buildInitialApi()
+})
+const emptyApiState = {
+  [baseApi.reducerPath]: baseApi.reducer(undefined, { type: 'foo' }),
+}
+
+test('injectTagTypes', () => {
+  const injectedApi = baseApi.addTagTypes('Foo', 'Bar')
+
+  injectedApi.util.selectInvalidatedBy(emptyApiState, ['Foo'])
+  // @ts-expect-error
+  injectedApi.util.selectInvalidatedBy(emptyApiState, ['Baz'])
+})
+
+function getEndpointDetails<
+  E extends {
+    initiate(arg: any): any
+    select: (arg: any) => (state: any) => any
+  }
+>(
+  e: E
+): {
+  arg: Parameters<E['initiate']>[0]
+  returned: NonNullable<ReturnType<ReturnType<E['select']>>['data']>
+} {
+  return {} as any
+}
+
+describe('enhanceEndpoint', () => {
+  describe('query', () => {
+    test('no changes', () => {
+      const injectedApi = baseApi.enhanceEndpoint('someQuery', {})
+      const { arg, returned } = getEndpointDetails(
+        injectedApi.endpoints.someQuery
+      )
+
+      expectType<'QueryArgument'>(arg)
+      expectType<'ReturnedFromQuery'>(returned)
+    })
+
+    test('change return value through `tranformResponse`', () => {
+      const injectedApi = baseApi.enhanceEndpoint('someQuery', {
+        transformResponse(): 'Changed' {
+          return 'Changed'
+        },
+      })
+      const { arg, returned } = getEndpointDetails(
+        injectedApi.endpoints.someQuery
+      )
+
+      expectType<'QueryArgument'>(arg)
+      expectType<'Changed'>(returned)
+      // @ts-expect-error
+      expectType<'ReturnedFromQuery'>(returned)
+    })
+
+    test('change argument value through new `query` function', () => {
+      const injectedApi = baseApi.enhanceEndpoint('someQuery', {
+        query(arg: 'Changed') {
+          return '/'
+        },
+      })
+      const { arg, returned } = getEndpointDetails(
+        injectedApi.endpoints.someQuery
+      )
+
+      expectType<'Changed'>(arg)
+      // @ts-expect-error
+      expectType<'QueryArgument'>(arg)
+      expectType<'ReturnedFromQuery'>(returned)
+    })
+
+    test('change argument and return value through new `query` and `transformResponse` functions', () => {
+      const injectedApi = baseApi.enhanceEndpoint('someQuery', {
+        query(arg: 'Changed') {
+          return ''
+        },
+        transformResponse(): 'AlsoChanged' {
+          return 'AlsoChanged'
+        },
+      })
+      const { arg, returned } = getEndpointDetails(
+        injectedApi.endpoints.someQuery
+      )
+
+      expectType<'Changed'>(arg)
+      // @ts-expect-error
+      expectType<'QueryArgument'>(arg)
+      expectType<'AlsoChanged'>(returned)
+      // @ts-expect-error
+      expectType<'ReturnedFromQuery'>(returned)
+    })
+
+    test('change argument and return value through new `queryFn` function', () => {
+      const injectedApi = baseApi.enhanceEndpoint('someQuery', {
+        queryFn(arg: 'Changed') {
+          return { data: 'AlsoChanged' as const }
+        },
+      })
+      const { arg, returned } = getEndpointDetails(
+        injectedApi.endpoints.someQuery
+      )
+
+      expectType<'Changed'>(arg)
+      // @ts-expect-error
+      expectType<'QueryArgument'>(arg)
+      expectType<'AlsoChanged'>(returned)
+      // @ts-expect-error
+      expectType<'ReturnedFromQuery'>(returned)
+    })
+  })
+  describe('mutation', () => {
+    test('no changes', () => {
+      const injectedApi = baseApi.enhanceEndpoint('someMutation', {})
+      const { arg, returned } = getEndpointDetails(
+        injectedApi.endpoints.someMutation
+      )
+
+      expectType<'MutationArgument'>(arg)
+      expectType<'ReturnedFromMutation'>(returned)
+    })
+
+    test('change return value through `tranformResponse`', () => {
+      const injectedApi = baseApi.enhanceEndpoint('someMutation', {
+        transformResponse(): 'Changed' {
+          return 'Changed'
+        },
+      })
+      const { arg, returned } = getEndpointDetails(
+        injectedApi.endpoints.someMutation
+      )
+
+      expectType<'MutationArgument'>(arg)
+      expectType<'Changed'>(returned)
+      // @ts-expect-error
+      expectType<'ReturnedFromMutation'>(returned)
+    })
+
+    test('change argument value through new `query` function', () => {
+      const injectedApi = baseApi.enhanceEndpoint('someMutation', {
+        query(arg: 'Changed') {
+          return '/'
+        },
+      })
+      const { arg, returned } = getEndpointDetails(
+        injectedApi.endpoints.someMutation
+      )
+
+      expectType<'Changed'>(arg)
+      // @ts-expect-error
+      expectType<'MutationArgument'>(arg)
+      expectType<'ReturnedFromMutation'>(returned)
+    })
+
+    test('change argument and return value through new `queryFn` function', () => {
+      const injectedApi = baseApi.enhanceEndpoint('someMutation', {
+        query(arg: 'Changed') {
+          return ''
+        },
+        transformResponse(): 'AlsoChanged' {
+          return 'AlsoChanged'
+        },
+      })
+      const { arg, returned } = getEndpointDetails(
+        injectedApi.endpoints.someMutation
+      )
+
+      expectType<'Changed'>(arg)
+      // @ts-expect-error
+      expectType<'MutationArgument'>(arg)
+      expectType<'AlsoChanged'>(returned)
+      // @ts-expect-error
+      expectType<'ReturnedFromMutation'>(returned)
+    })
+
+    test('change argument and return value through new `queryFn` function', () => {
+      const injectedApi = baseApi.enhanceEndpoint('someMutation', {
+        queryFn(arg: 'Changed') {
+          return { data: 'AlsoChanged' as const }
+        },
+      })
+      const { arg, returned } = getEndpointDetails(
+        injectedApi.endpoints.someMutation
+      )
+
+      expectType<'Changed'>(arg)
+      // @ts-expect-error
+      expectType<'MutationArgument'>(arg)
+      expectType<'AlsoChanged'>(returned)
+      // @ts-expect-error
+      expectType<'ReturnedFromMutation'>(returned)
+    })
+  })
+})

packages/toolkit/src/query/tsHelpers.ts

@@ -47,3 +47,7 @@ export type IsAny<T, True, False = never> = true | false extends (
   : False
 
 export type CastAny<T, CastTo> = IsAny<T, CastTo, T>
+
+export type OverrideAtKey<Type, Key, NewKeyType> = {
+  [K in keyof Type]: Key extends K ? NewKeyType : Type[K]
+}