Deduplicate endpoint definition lookups

Author: markerikson

packages/toolkit/src/query/apiTypes.ts

@@ -1,6 +1,6 @@
 import type { UnknownAction } from '@reduxjs/toolkit'
 import type { BaseQueryFn } from './baseQueryTypes'
-import type { CombinedState, CoreModule } from './core'
+import type { CombinedState, CoreModule, QueryKeys } from './core'
 import type { ApiModules } from './core/module'
 import type { CreateApiOptions } from './createApi'
 import type {
@@ -56,6 +56,14 @@ export interface ApiContext<Definitions extends EndpointDefinitions> {
   hasRehydrationInfo: (action: UnknownAction) => boolean
 }
 
+export const getEndpointDefinition = <
+  Definitions extends EndpointDefinitions,
+  EndpointName extends keyof Definitions,
+>(
+  context: ApiContext<Definitions>,
+  endpointName: EndpointName,
+) => context.endpointDefinitions[endpointName]
+
 export type Api<
   BaseQuery extends BaseQueryFn,
   Definitions extends EndpointDefinitions,

packages/toolkit/src/query/core/buildInitiate.ts

@@ -7,7 +7,7 @@ import type {
 } from '@reduxjs/toolkit'
 import type { Dispatch } from 'redux'
 import { asSafePromise } from '../../tsHelpers'
-import type { Api, ApiContext } from '../apiTypes'
+import { getEndpointDefinition, type Api, type ApiContext } from '../apiTypes'
 import type { BaseQueryError, QueryReturnValue } from '../baseQueryTypes'
 import type { InternalSerializeQueryArgs } from '../defaultSerializeQueryArgs'
 import {
@@ -304,7 +304,7 @@ export function buildInitiate({
 
   function getRunningQueryThunk(endpointName: string, queryArgs: any) {
     return (dispatch: Dispatch) => {
-      const endpointDefinition = context.endpointDefinitions[endpointName]
+      const endpointDefinition = getEndpointDefinition(context, endpointName)
       const queryCacheKey = serializeQueryArgs({
         queryArgs,
         endpointDefinition,

packages/toolkit/src/query/core/buildMiddleware/cacheCollection.ts

@@ -1,3 +1,4 @@
+import { getEndpointDefinition } from '@internal/query/apiTypes'
 import type { QueryDefinition } from '../../endpointDefinitions'
 import type { ConfigState, QueryCacheKey, QuerySubState } from '../apiState'
 import { isAnyOf } from '../rtkImports'
@@ -155,9 +156,10 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({
     api: SubMiddlewareApi,
     config: ConfigState<string>,
   ) {
-    const endpointDefinition = context.endpointDefinitions[
-      endpointName
-    ] as QueryDefinition<any, any, any, any>
+    const endpointDefinition = getEndpointDefinition(
+      context,
+      endpointName,
+    ) as QueryDefinition<any, any, any, any>
     const keepUnusedDataFor =
       endpointDefinition?.keepUnusedDataFor ?? config.keepUnusedDataFor
 

packages/toolkit/src/query/core/buildMiddleware/cacheLifecycle.ts

@@ -23,6 +23,7 @@ import type {
   PromiseWithKnownReason,
   SubMiddlewareApi,
 } from './types'
+import { getEndpointDefinition } from '@internal/query/apiTypes'
 
 export type ReferenceCacheLifecycle = never
 
@@ -317,7 +318,7 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({
     mwApi: SubMiddlewareApi,
     requestId: string,
   ) {
-    const endpointDefinition = context.endpointDefinitions[endpointName]
+    const endpointDefinition = getEndpointDefinition(context, endpointName)
     const onCacheEntryAdded = endpointDefinition?.onCacheEntryAdded
     if (!onCacheEntryAdded) return
 

packages/toolkit/src/query/core/buildMiddleware/queryLifecycle.ts

@@ -1,9 +1,10 @@
+import { getEndpointDefinition } from '@internal/query/apiTypes'
 import type {
   BaseQueryError,
   BaseQueryFn,
   BaseQueryMeta,
 } from '../../baseQueryTypes'
-import { DefinitionType, isAnyQueryDefinition } from '../../endpointDefinitions'
+import { isAnyQueryDefinition } from '../../endpointDefinitions'
 import type { Recipe } from '../buildThunks'
 import { isFulfilled, isPending, isRejected } from '../rtkImports'
 import type {
@@ -442,7 +443,7 @@ export const buildQueryLifecycleHandler: InternalHandlerBuilder = ({
         requestId,
         arg: { endpointName, originalArgs },
       } = action.meta
-      const endpointDefinition = context.endpointDefinitions[endpointName]
+      const endpointDefinition = getEndpointDefinition(context, endpointName)
       const onQueryStarted = endpointDefinition?.onQueryStarted
       if (onQueryStarted) {
         const lifecycle = {} as CacheLifecycle

packages/toolkit/src/query/createApi.ts

@@ -1,4 +1,10 @@
-import type { Api, ApiContext, Module, ModuleName } from './apiTypes'
+import {
+  getEndpointDefinition,
+  type Api,
+  type ApiContext,
+  type Module,
+  type ModuleName,
+} from './apiTypes'
 import type { CombinedState } from './core/apiState'
 import type { BaseQueryArg, BaseQueryFn } from './baseQueryTypes'
 import type { SerializeQueryArgs } from './defaultSerializeQueryArgs'
@@ -421,10 +427,10 @@ export function buildCreateApi<Modules extends [Module<any>, ...Module<any>[]]>(
             endpoints,
           )) {
             if (typeof partialDefinition === 'function') {
-              partialDefinition(context.endpointDefinitions[endpointName])
+              partialDefinition(getEndpointDefinition(context, endpointName))
             } else {
               Object.assign(
-                context.endpointDefinitions[endpointName] || {},
+                getEndpointDefinition(context, endpointName) || {},
                 partialDefinition,
               )
             }