Wrapping RTKQ Hooks (and getting the typing right)

[StefanBRas]

TLDR; I want to make a generic wrapper around RTKQ hooks, but when i try to pass the options argument to the wrapped hook, it breaks the typing. See the codebox and notice the optionsData has type any, while the others correctly infers the type.

---

Hello!

I'm using RTK Query. Most of my hooks use some value from my redux store, and some use several. To avoid having to use useSelectors everytime i use these useQuery hooks, I can either

1. Use a custom BaseQuery where i use getState() and then modify the args being passed or,
2. Wrap the useQueries in a custom hook.

Now, if I do 1 i have to manually invalidate the cache whenever the value from the store changes, so i want to go with 2.
This is easy enough, for example (omitting redux stuff just for clarity):

export interface QueryArgs {
  arg1: string;
  arg2: string;
}

export interface QueryResult {
  value: string;
}

const wrapperManual = (args: QueryArgs) => () => {
  return useGetStuffQuery(args);
};

export const useGetStuffQueryWrappedManual = wrapperManual({
  arg1: "a",
  arg2: "b"
});

I'm ignoring the options parameter for now.
This works but requires me to do it with every hook. So i make a generic wrapper, that takes the name of the endpoint and the name of arguments which is given in the wrapper (and thus removed from the inner hooks' arguments).
I use some type helper too, you can see them here:

Typing helpers

type EndpointKey = keyof typeof api.endpoints;
type Endpoint<EndpointName extends EndpointKey> = typeof api.endpoints[EndpointName];
type InferArgs<E> = E extends ApiEndpointQuery<QueryDefinition<infer A, any, any, any>,any> ? A : never;
type InferArgsKeys<E> = keyof InferArgs<E>;
type InferUseQuery<E> = E extends ApiEndpointQuery<infer QD, any> ? UseQuery<QD> : never;
type InferQD<E> = E extends ApiEndpointQuery<infer QD, any> ? QD : never;
type InferOptions<E extends EndpointKey> = Parameters<Endpoint<E>["useQuery"]>[1];
type InferQdParms<E> = E extends ApiEndpointQuery<QueryDefinition<any, infer B, infer C, infer D>,any>? [B, C, D]: never;
type ReturnQuery<E, A> = E extends ApiEndpointQuery<QueryDefinition<any, infer B, infer C, infer D>,any>
  ? UseQuery<QueryDefinition<A, B, C, D>>
  : never;
type PartialArgs<E extends EndpointKey> = Partial<InferArgs<Endpoint<E>>>;
type PickArgs<
  E extends EndpointKey,
  K extends InferArgsKeys<Endpoint<E>>
> = Pick<InferArgs<Endpoint<E>>, K>;
type OmitArgs<
E extends EndpointKey, 
K extends InferArgsKeys<E>> = Omit<InferArgs<Endpoint<E>>,K>;

Generic version:

function wrapperBasic<
  E extends EndpointKey,
  K extends keyof InferArgs<Endpoint<E>>
>(hook: InferUseQuery<Endpoint<E>>, args: Pick<InferArgs<Endpoint<E>>, K>) {
  function useWrapper(
    oargs: Omit<InferArgs<Endpoint<E>>, K> | typeof skipToken
  ) {
    if (oargs === skipToken) {
      return hook(skipToken);
    } else {
      return hook({ ...oargs, ...args } as InferArgs<Endpoint<E>>);
    }
  }
  return useWrapper;
}

export const useGetStuffQueryWrappedBasic = wrapperBasic<"getStuff", "arg1">(
  useGetStuffQuery,
  { arg1: "a" }
);

This works great and if i use the useGetStuffQueryWrappedBasic hook in a component, the args are correctly inferred to be only {args2: string} now. Nice! So I want the wrapped hook to pass the options to the inner hook so i do:

unction wrapperOptions<
  E extends EndpointKey,
  K extends keyof InferArgs<Endpoint<E>>
>(hook: InferUseQuery<Endpoint<E>>, args: Pick<InferArgs<Endpoint<E>>, K>) {
  function useWrapper(
    oargs: Omit<InferArgs<Endpoint<E>>, K> | typeof skipToken,
    options?: InferOptions<E>
  ) {
    if (oargs === skipToken) {
      return hook(skipToken, options);
    } else {
      return hook({ ...oargs, ...args } as InferArgs<Endpoint<E>>, options);
    }
  }
  useWrapper;
}

but whenever I try to use this wrapper, the data will come out as any even though the only difference is adding the options.

I have solved this for now by making an assertion so i instead return

...
type NewArgs = Omit<InferArgs<Endpoint<E>>, K>;
return useWrapper as ReturnQuery<Endpoint<E>, NewArgs>;

But of course I'd rather avoid this assertion if possible.

I'd love to hear your thoughts on this and you can try this codesandbox to play around with it.

[GriffinSauce]

@StefanBRas did you get anywhere with this in the end?

[StefanBRas]

Nope, I haven't had the chance to work with RTKQ since this really. I still think it would be nice to have though, maybe even as part of RTKQ.

[GriffinSauce]

Thanks for responding so quickly 🙏

Yeah I agree. I haven't been able to make a wrapper factory work but I've gotten single hooks to work. (I assume this means it's possible and entirely down to my own incompetence with types ;) )

[StefanBRas]

No problem!

Phryneas warned me about this before hand too - on discord he wrote:

But fair warning: making a generalized version of it and typing it correctly will be difficult.

But depending on the scope of the project, I think just doing the manual assertion in the end is good enough. The types are only complicated because you have to dynamically infer a lot - the return type itself is just the same as the original hook, but with some arguments omitted.

[GriffinSauce]

Yep, so I ended up just wrapping each hook individually and using some helpers (inspired by your code 🙏 ) to tie it together without having to infer everything. It's a little boilerplatey but it works well and is way less complex, which is also a positive. Cheers!

[lostm1nd]

I was trying to do something similar by wrapping the query hooks so I can add some common behaviour. The problem for me remains the same - cannot get the typing right so that I dont lose the BE payload type. My code is something along those lines:

import React from 'react';

import { NotificationManager } from '@styleguide/composite/notifications';
import { UseMutation } from '@reduxjs/toolkit/dist/query/react/buildHooks';
import { MutationDefinition } from '@reduxjs/toolkit/dist/query';

const useGenericNotification = <T extends UseMutation<MutationDefinition<any, any, string, any, string>>>(
  mutation: T,
  { notification = 'Success' } = {}
) => {
  const [mutate, status] = mutation();

  const mutateWithNotification = React.useCallback(
    async (data) => {
      try {
        const response = await mutate(data).unwrap();

        NotificationManager.success(notification);

        return response;
      } catch (error) {
        NotificationManager.error(error.toString());
        throw error;
      }
    },
    [mutate, notification]
  );

  return [mutateWithNotification, status] as const;
};

export { useGenericNotification };

Ideas are welcome ; )

[phryneas]

Can you give the CodeSandbox build of #2276 a try and see if that is what you need?

[TreeMage]

Hi,

not sure if anyone is still interested in this, but the following works perfectly for me:

import { BaseQueryFn, MutationDefinition } from '@reduxjs/toolkit/query';
import { UseMutation } from '@reduxjs/toolkit/dist/query/react/buildHooks';

const useWrappedMutation = <U, V extends BaseQueryFn, W>(
  mutation: UseMutation<MutationDefinition<U, V, string, W>>
): ReturnType<UseMutation<MutationDefinition<U, V, string, W>>> => {
  const [mutate, status] = mutation();
  // Do your think with the mutation/status ...
  return [mutate, status];
};

[Roman86]

here is the helper, in case you need to have a small/nice code to skip when some of arguments is undefined:

// utils to put somewhere separately:
const undefinedFilter = ((v: any) => typeof v !== 'undefined') as any as <T>(
  x: T | undefined,
) => x is NonNullable<T>;

function allFieldsDefined<O extends Record<string, any>>(o: O): o is PropsNonNullable<O> {
  return Object.values(o).every(undefinedFilter);
}

// main utility
import { skipToken } from '@reduxjs/toolkit/query/react';

function paramsOrSkip<P>(p: Partial<P>): P | typeof skipToken {
  return allFieldsDefined(p) ? p : skipToken;
}

// your cool hook with required param inside (projectId in this example), fed with another hook that can return undefined;
// our RTK query will be skipped while any of arguments is undefined
export function useFetchSomethingQuery() {
  return api.useFetchSomethingQuery(
    paramsOrSkip({ // params infer from the query which is cool
      projectId: useProjectId(), // no conditional call here, so you may invoke hooks directly without defining constants
      somethingElse: useSomethingElse(),
    }),
  );
}