Consolidate query definition checks

Author: markerikson

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

@@ -1,7 +1,7 @@
 import type { ThunkDispatch, UnknownAction } from '@reduxjs/toolkit'
 import type { BaseQueryFn, BaseQueryMeta } from '../../baseQueryTypes'
 import type { BaseEndpointDefinition } from '../../endpointDefinitions'
-import { DefinitionType } from '../../endpointDefinitions'
+import { DefinitionType, isAnyQueryDefinition } from '../../endpointDefinitions'
 import type { QueryCacheKey, RootState } from '../apiState'
 import type {
   MutationResultSelectorResult,
@@ -328,10 +328,7 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({
     cacheDataLoaded.catch(() => {})
     lifecycleMap[queryCacheKey] = lifecycle
     const selector = (api.endpoints[endpointName] as any).select(
-      endpointDefinition.type === DefinitionType.query ||
-        endpointDefinition.type === DefinitionType.infinitequery
-        ? originalArgs
-        : queryCacheKey,
+      isAnyQueryDefinition(endpointDefinition) ? originalArgs : queryCacheKey,
     )
 
     const extra = mwApi.dispatch((_, __, extra) => extra)
@@ -340,8 +337,7 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({
       getCacheEntry: () => selector(mwApi.getState()),
       requestId,
       extra,
-      updateCachedData: (endpointDefinition.type === DefinitionType.query ||
-      endpointDefinition.type === DefinitionType.infinitequery
+      updateCachedData: (isAnyQueryDefinition(endpointDefinition)
         ? (updateRecipe: Recipe<any>) =>
             mwApi.dispatch(
               api.util.updateQueryData(

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

@@ -3,7 +3,7 @@ import type {
   BaseQueryFn,
   BaseQueryMeta,
 } from '../../baseQueryTypes'
-import { DefinitionType } from '../../endpointDefinitions'
+import { DefinitionType, isAnyQueryDefinition } from '../../endpointDefinitions'
 import type { Recipe } from '../buildThunks'
 import { isFulfilled, isPending, isRejected } from '../rtkImports'
 import type {
@@ -459,10 +459,7 @@ export const buildQueryLifecycleHandler: InternalHandlerBuilder = ({
         queryFulfilled.catch(() => {})
         lifecycleMap[requestId] = lifecycle
         const selector = (api.endpoints[endpointName] as any).select(
-          endpointDefinition.type === DefinitionType.query ||
-            endpointDefinition.type === DefinitionType.infinitequery
-            ? originalArgs
-            : requestId,
+          isAnyQueryDefinition(endpointDefinition) ? originalArgs : requestId,
         )
 
         const extra = mwApi.dispatch((_, __, extra) => extra)
@@ -471,8 +468,7 @@ export const buildQueryLifecycleHandler: InternalHandlerBuilder = ({
           getCacheEntry: () => selector(mwApi.getState()),
           requestId,
           extra,
-          updateCachedData: (endpointDefinition.type === DefinitionType.query ||
-          endpointDefinition.type === DefinitionType.infinitequery
+          updateCachedData: (isAnyQueryDefinition(endpointDefinition)
             ? (updateRecipe: Recipe<any>) =>
                 mwApi.dispatch(
                   api.util.updateQueryData(

packages/toolkit/src/query/endpointDefinitions.ts

@@ -912,6 +912,14 @@ export function isInfiniteQueryDefinition(
   return e.type === DefinitionType.infinitequery
 }
 
+export function isAnyQueryDefinition(
+  e: EndpointDefinition<any, any, any, any>,
+): e is
+  | QueryDefinition<any, any, any, any>
+  | InfiniteQueryDefinition<any, any, any, any, any> {
+  return isQueryDefinition(e) || isInfiniteQueryDefinition(e)
+}
+
 export type EndpointBuilder<
   BaseQuery extends BaseQueryFn,
   TagTypes extends string,