add onQuery lifecycle

src/query/core/buildMiddleware/index.ts

@@ -19,6 +19,7 @@ import { build as buildPolling } from './polling'
 import { BuildMiddlewareInput } from './types'
 import { build as buildWindowEventHandling } from './windowEventHandling'
 import { build as buildCacheLifecycle } from './cacheLifecycle'
+import { build as buildQueryLifecycle } from './queryLifecycle'
 
 export function buildMiddleware<
   Definitions extends EndpointDefinitions,
@@ -38,6 +39,7 @@ export function buildMiddleware<
     buildPolling,
     buildWindowEventHandling,
     buildCacheLifecycle,
+    buildQueryLifecycle,
   ].map((build) =>
     build({
       ...((input as any) as BuildMiddlewareInput<

src/query/core/buildMiddleware/queryLifecycle.ts

@@ -0,0 +1,52 @@
+import { isPending, isRejected, isFulfilled } from '@reduxjs/toolkit'
+import { toOptionalPromise } from '../../utils/toOptionalPromise'
+import { SubMiddlewareBuilder } from './types'
+
+export const build: SubMiddlewareBuilder = ({
+  context,
+  queryThunk,
+  mutationThunk,
+}) => {
+  type CacheLifecycle = {
+    resolve(value: unknown): unknown
+    reject(value: unknown): unknown
+  }
+  const lifecycleMap: Record<string, CacheLifecycle> = {}
+
+  const isPendingThunk = isPending(queryThunk, mutationThunk)
+  const isRejectedThunk = isRejected(queryThunk, mutationThunk)
+  const isFullfilledThunk = isFulfilled(queryThunk, mutationThunk)
+
+  return (mwApi) => (next) => (action): any => {
+    const result = next(action)
+
+    if (isPendingThunk(action)) {
+      const {
+        requestId,
+        arg: { endpointName, originalArgs },
+      } = action.meta
+      const { onQuery } = context.endpointDefinitions[endpointName]
+      if (onQuery) {
+        const lifecycle = {} as CacheLifecycle
+        const resultPromise = toOptionalPromise(
+          new Promise((resolve, reject) => {
+            lifecycle.resolve = resolve
+            lifecycle.reject = reject
+          })
+        )
+        lifecycleMap[requestId] = lifecycle
+        onQuery(originalArgs, mwApi, { resultPromise })
+      }
+    } else if (isFullfilledThunk(action)) {
+      const { requestId } = action.meta
+      lifecycleMap[requestId].resolve(action.payload)
+      delete lifecycleMap[requestId]
+    } else if (isRejectedThunk(action)) {
+      const { requestId } = action.meta
+      lifecycleMap[requestId].reject(action.payload ?? action.error)
+      delete lifecycleMap[requestId]
+    }
+
+    return result
+  }
+}

src/query/endpointDefinitions.ts

@@ -99,6 +99,20 @@ export type BaseEndpointDefinition<
       cleanup: Promise<void>
     }
   ): Promise<void> | void
+  onQuery?(
+    arg: QueryArg,
+    api: LifecycleApi,
+    promises: {
+      /**
+       * Promise that will resolve with the (transformed) query result.
+       
+       * If the query fails, this promise will reject with the error.
+       *
+       * This allows you to `await` for the query to finish.
+       */
+      resultPromise: OptionalPromise<ResultType>
+    }
+  ): Promise<void> | void
 } & HasRequiredProps<
     BaseQueryExtraOptions<BaseQuery>,
     { extraOptions: BaseQueryExtraOptions<BaseQuery> },