add onStart onError onSuccess hooks to query (#130)

* add onStart onError onSuccess hooks to query

* code review
rename original type QueryApi to BaseQueryApi, move to baseQueryTypes
introduce new type QueryApi to endpointDefinitions
add `await` both at query&mutationThunk `return` property

* also run tests on pull requests

* Remove signal from QueryApi, add endpoint desc to queries documentation

* Update types, add basic lifecycle method tests

Co-authored-by: Lenz Weber <mail@lenzw.de>
Co-authored-by: Matt Sutkowski <msutkowski@gmail.com>

.github/workflows/main.yml

@@ -1,5 +1,5 @@
 name: CI
-on: [push]
+on: [push, pull_request]
 jobs:
   build:
     name: Build, lint, and test on Node ${{ matrix.node }} and ${{ matrix.os }}

docs/api/createApi.md

@@ -120,13 +120,18 @@ Endpoints are just a set of operations that you want to perform against your ser
   - Used by `mutations` for [cache invalidation](../concepts/mutations#advanced-mutations-with-revalidation) purposes.
   - Expects the same shapes as `provides`
 
-- `onStart`, `onError` and `onSuccess` _(optional)_
+- `onStart`, `onError` and `onSuccess` _(optional)_ - Available to both [queries](../concepts/queries) and [mutations](../concepts/mutations)
   - Can be used in `mutations` for [optimistic updates](../concepts/optimistic-updates).
-  - ```ts title="signatures"
+  - ```ts title="Mutation lifecycle signatures"
     function onStart(arg: QueryArg, mutationApi: MutationApi<ReducerPath, Context>): void;
     function onError(arg: QueryArg, mutationApi: MutationApi<ReducerPath, Context>, error: unknown): void;
     function onSuccess(arg: QueryArg, mutationApi: MutationApi<ReducerPath, Context>, result: ResultType): void;
     ```
+  - ```ts title="Query lifecycle signatures"
+    function onStart(arg: QueryArg, queryApi: QueryApi<ReducerPath, Context>): void;
+    function onError(arg: QueryArg, queryApi: QueryApi<ReducerPath, Context>, error: unknown): void;
+    function onSuccess(arg: QueryArg, queryApi: QueryApi<ReducerPath, Context>, result: ResultType): void;
+    ```
 
 #### How endpoints get used
 

docs/concepts/mutations.md

@@ -18,9 +18,13 @@ const api = createApi({
   endpoints: (build) => ({
     updatePost: build.mutation({
       query: ({ id, ...patch }) => ({ url: `post/${id}`, method: 'PATCH', body: patch }),
+      // Pick out data and prevent nested properties in a hook or selector
+      transformResponse: (response) => response.data,
       // onStart, onSuccess, onError are useful for optimistic updates
-      onStart({ id, ...patch }, mutationApi) {},
-      onSuccess({ id }, { dispatch, getState, extra, requestId, context }, result) {}, // result is the server response, the 2nd parameter is the destructured `mutationApi`
+      // The 2nd parameter is the destructured `mutationApi`
+      onStart({ id, ...patch }, { dispatch, getState, extra, requestId, context }) {},
+      // `result` is the server response
+      onSuccess({ id }, mutationApi, result) {},
       onError({ id }, { dispatch, getState, extra, requestId, context }) {},
       invalidates: ['Post'],
     }),

docs/concepts/queries.md

@@ -13,6 +13,25 @@ By default, RTK Query ships with [`fetchBaseQuery`](../api/fetchBaseQuery), whic
 
 > Depending on your environment, you may need to polyfill `fetch` with `node-fetch` or `cross-fetch` if you choose to use `fetchBaseQuery` or `fetch` on its own.
 
+```js title="Example of all query endpoint options"
+const api = createApi({
+  baseQuery,
+  endpoints: (build) => ({
+    getPost: build.query({
+      query: (id) => ({ url: `post/${id}` }),
+      // Pick out data and prevent nested properties in a hook or selector
+      transformResponse: (response) => response.data,
+      // The 2nd parameter is the destructured `queryApi`
+      onStart(id, { dispatch, getState, extra, requestId, context }) {},
+      // `result` is the server response
+      onSuccess(id, queryApi, result) {},
+      onError(id, queryApi) {},
+      provides: (_, id) => [{ type: 'Post', id }],
+    }),
+  }),
+});
+```
+
 ### Performing queries with React Hooks
 
 If you're using React Hooks, RTK Query does a few additional things for you. The primary benefit is that you get a render-optimized hook that allows you to have 'background fetching' as well as [derived booleans](#query-hook-return-types) for convenience.

src/baseQueryTypes.ts

@@ -1,9 +1,15 @@
-import { QueryApi } from './core/buildThunks';
+import { ThunkDispatch } from '@reduxjs/toolkit';
 import { MaybePromise, UnwrapPromise } from './tsHelpers';
 
+export interface BaseQueryApi {
+  signal?: AbortSignal;
+  dispatch: ThunkDispatch<any, any, any>;
+  getState: () => unknown;
+}
+
 export type BaseQueryFn<Args = any, Result = unknown, Error = unknown, DefinitionExtraOptions = {}> = (
   args: Args,
-  api: QueryApi,
+  api: BaseQueryApi,
   extraOptions: DefinitionExtraOptions
 ) => MaybePromise<{ error: Error; data?: undefined } | { error?: undefined; data?: Result }>;
 

src/core/buildThunks.ts

@@ -8,6 +8,7 @@ import {
   EndpointDefinitions,
   MutationApi,
   MutationDefinition,
+  QueryApi,
   QueryArgFrom,
   QueryDefinition,
   ResultTypeFrom,
@@ -95,12 +96,6 @@ export interface ThunkResult {
 export type QueryThunk = AsyncThunk<ThunkResult, QueryThunkArg<any>, {}>;
 export type MutationThunk = AsyncThunk<ThunkResult, MutationThunkArg<any>, {}>;
 
-export interface QueryApi {
-  signal?: AbortSignal;
-  dispatch: ThunkDispatch<any, any, any>;
-  getState: () => unknown;
-}
-
 function defaultTransformResponse(baseQueryReturnValue: unknown) {
   return baseQueryReturnValue;
 }
@@ -196,18 +191,37 @@ export function buildThunks<
     { state: RootState<any, string, ReducerPath> }
   >(
     `${reducerPath}/executeQuery`,
-    async (arg, { signal, rejectWithValue, dispatch, getState }) => {
-      const result = await baseQuery(
-        arg.internalQueryArgs,
-        { signal, dispatch, getState },
-        endpointDefinitions[arg.endpoint].extraOptions as any
-      );
-      if (result.error) return rejectWithValue(result.error);
+    async (arg, { signal, rejectWithValue, ...api }) => {
+      const endpoint = endpointDefinitions[arg.endpoint] as QueryDefinition<any, any, any, any>;
 
-      return {
-        fulfilledTimeStamp: Date.now(),
-        result: await (endpointDefinitions[arg.endpoint].transformResponse ?? defaultTransformResponse)(result.data),
+      const context: Record<string, any> = {};
+      const queryApi: QueryApi<ReducerPath, any> = {
+        ...api,
+        context,
       };
+
+      if (endpoint.onStart) endpoint.onStart(arg.originalArgs, queryApi);
+
+      try {
+        const result = await baseQuery(
+          arg.internalQueryArgs,
+          { signal, dispatch: api.dispatch, getState: api.getState },
+          endpoint.extraOptions as any
+        );
+        if (result.error) throw new HandledError(result.error);
+        if (endpoint.onSuccess) endpoint.onSuccess(arg.originalArgs, queryApi, result.data);
+        return {
+          fulfilledTimeStamp: Date.now(),
+          result: await (endpoint.transformResponse ?? defaultTransformResponse)(result.data),
+        };
+      } catch (error) {
+        if (endpoint.onError)
+          endpoint.onError(arg.originalArgs, queryApi, error instanceof HandledError ? error.value : error);
+        if (error instanceof HandledError) {
+          return rejectWithValue(error.value);
+        }
+        throw error;
+      }
     },
     {
       condition(arg, { getState }) {
@@ -245,23 +259,23 @@ export function buildThunks<
     const endpoint = endpointDefinitions[arg.endpoint] as MutationDefinition<any, any, any, any>;
 
     const context: Record<string, any> = {};
-    const mutationApi = {
+    const mutationApi: MutationApi<ReducerPath, any> = {
       ...api,
       context,
-    } as MutationApi<ReducerPath, any>;
+    };
 
     if (endpoint.onStart) endpoint.onStart(arg.originalArgs, mutationApi);
     try {
       const result = await baseQuery(
         arg.internalQueryArgs,
         { signal, dispatch: api.dispatch, getState: api.getState },
-        endpointDefinitions[arg.endpoint].extraOptions as any
+        endpoint.extraOptions as any
       );
       if (result.error) throw new HandledError(result.error);
       if (endpoint.onSuccess) endpoint.onSuccess(arg.originalArgs, mutationApi, result.data);
       return {
         fulfilledTimeStamp: Date.now(),
-        result: (endpointDefinitions[arg.endpoint].transformResponse ?? defaultTransformResponse)(result.data),
+        result: await (endpoint.transformResponse ?? defaultTransformResponse)(result.data),
       };
     } catch (error) {
       if (endpoint.onError)

src/endpointDefinitions.ts

@@ -35,16 +35,28 @@ type ResultDescription<EntityTypes extends string, ResultType, QueryArg> =
   | ReadonlyArray<EntityDescription<EntityTypes>>
   | GetResultDescriptionFn<EntityTypes, ResultType, QueryArg>;
 
+export interface QueryApi<ReducerPath extends string, Context extends {}> {
+  dispatch: ThunkDispatch<RootState<any, any, ReducerPath>, unknown, AnyAction>;
+  getState(): RootState<any, any, ReducerPath>;
+  extra: unknown;
+  requestId: string;
+  context: Context;
+}
+
 export type QueryDefinition<
   QueryArg,
   BaseQuery extends BaseQueryFn,
   EntityTypes extends string,
   ResultType,
-  _ReducerPath extends string = string
+  ReducerPath extends string = string,
+  Context = Record<string, any>
 > = BaseEndpointDefinition<QueryArg, BaseQuery, ResultType> & {
   type: DefinitionType.query;
   provides?: ResultDescription<EntityTypes, ResultType, QueryArg>;
   invalidates?: never;
+  onStart?(arg: QueryArg, queryApi: QueryApi<ReducerPath, Context>): void;
+  onError?(arg: QueryArg, queryApi: QueryApi<ReducerPath, Context>, error: unknown): void;
+  onSuccess?(arg: QueryArg, queryApi: QueryApi<ReducerPath, Context>, result: ResultType): void;
 };
 
 export interface MutationApi<ReducerPath extends string, Context extends {}> {

test/createApi.test.ts

@@ -1,7 +1,9 @@
-import { configureStore } from '@reduxjs/toolkit';
+import { configureStore, createAction, createReducer } from '@reduxjs/toolkit';
 import { Api, createApi, fetchBaseQuery } from '@rtk-incubator/rtk-query';
 import { QueryDefinition, MutationDefinition } from '@internal/endpointDefinitions';
 import { ANY, expectType, setupApiStore, waitMs } from './helpers';
+import { server } from './mocks/server';
+import { rest } from 'msw';
 
 test('sensible defaults', () => {
   const api = createApi({
@@ -316,3 +318,87 @@ describe('additional transformResponse behaviors', () => {
     expect(result.data).toEqual({ value: 'success', banana: 'bread' });
   });
 });
+
+describe('query endpoint lifecycles - onStart, onSuccess, onError', () => {
+  const initialState = {
+    count: null as null | number,
+  };
+  const setCount = createAction<number>('setCount');
+  const testReducer = createReducer(initialState, (builder) => {
+    builder.addCase(setCount, (state, action) => {
+      state.count = action.payload;
+    });
+  });
+
+  type SuccessResponse = { value: 'success' };
+  const api = createApi({
+    baseQuery: fetchBaseQuery({ baseUrl: 'http://example.com' }),
+    endpoints: (build) => ({
+      echo: build.mutation({
+        query: () => ({ method: 'PUT', url: '/echo' }),
+      }),
+      query: build.query<SuccessResponse, void>({
+        query: () => '/success',
+        onStart: (_, api) => {
+          api.dispatch(setCount(0));
+        },
+        onSuccess: (_, api) => {
+          api.dispatch(setCount(1));
+        },
+        onError: (_, api) => {
+          api.dispatch(setCount(-1));
+        },
+      }),
+      mutation: build.mutation<SuccessResponse, void>({
+        query: () => ({ url: '/success', method: 'POST' }),
+        onStart: (_, api) => {
+          api.dispatch(setCount(0));
+        },
+        onSuccess: (_, api) => {
+          api.dispatch(setCount(1));
+        },
+        onError: (_, api) => {
+          api.dispatch(setCount(-1));
+        },
+      }),
+    }),
+  });
+
+  const storeRef = setupApiStore(api, { testReducer });
+
+  test('query lifecycle events fire properly', async () => {
+    // We intentionally fail the first request so we can test all lifecycles
+    server.use(
+      rest.get('http://example.com/success', (_, res, ctx) => res.once(ctx.status(500), ctx.json({ value: 'failed' })))
+    );
+
+    expect(storeRef.store.getState().testReducer.count).toBe(null);
+    const failAttempt = storeRef.store.dispatch(api.endpoints.query.initiate());
+    expect(storeRef.store.getState().testReducer.count).toBe(0);
+    await failAttempt;
+    expect(storeRef.store.getState().testReducer.count).toBe(-1);
+
+    const successAttempt = storeRef.store.dispatch(api.endpoints.query.initiate());
+    expect(storeRef.store.getState().testReducer.count).toBe(0);
+    await successAttempt;
+    expect(storeRef.store.getState().testReducer.count).toBe(1);
+  });
+
+  test('mutation lifecycle events fire properly', async () => {
+    // We intentionally fail the first request so we can test all lifecycles
+    server.use(
+      rest.post('http://example.com/success', (_, res, ctx) => res.once(ctx.status(500), ctx.json({ value: 'failed' })))
+    );
+
+    expect(storeRef.store.getState().testReducer.count).toBe(null);
+    const failAttempt = storeRef.store.dispatch(api.endpoints.mutation.initiate());
+    expect(storeRef.store.getState().testReducer.count).toBe(0);
+    await failAttempt;
+    expect(storeRef.store.getState().testReducer.count).toBe(-1);
+
+    const successAttempt = storeRef.store.dispatch(api.endpoints.mutation.initiate());
+    expect(storeRef.store.getState().testReducer.count).toBe(0);
+    await successAttempt;
+    expect(storeRef.store.getState().testReducer.count).toBe(1);
+  });
+});