prefetchQuery, useSuspenseQuery with zsa

[davve5]

Is there any way to use zsa-react-query with prefetchQuery and useSuspenseQuery?
It would be nice to be able to prefetch data with validation, type safety and also to stay with react-query ecosystem

[cptlchad]

I didn't manage to preload data yet

[musjj]

Would love to see this feature too!

[strangebot]

Just picking up zsa for the first time. I'd also like to see support for prefetchQuery as well.

[davve5]

I was able to create a prefetchQuery for zsa. You need to configure the QueryKeyFactory in another file so that it can be used in both server and client code. Here's how I configured it:

get-query-client.ts

import {
  QueryClient,
  defaultShouldDehydrateQuery,
  isServer,
} from '@tanstack/react-query';

function makeQueryClient() {
  return new QueryClient({
    defaultOptions: {
      queries: {
        staleTime: 60 * 1000,
      },
      dehydrate: {
        // include pending queries in dehydration
        shouldDehydrateQuery: (query) =>
          defaultShouldDehydrateQuery(query) ||
          query.state.status === 'pending',
      },
    },
  });
}

let browserQueryClient: QueryClient | undefined = undefined;

export function getQueryClient() {
  if (isServer) {
    // Server: always make a new query client
    return makeQueryClient();
  } else {
    // Browser: make a new query client if we don't already have one
    // This is very important, so we don't re-make a new client if React
    // suspends during the initial render. This may not be needed if we
    // have a suspense boundary BELOW the creation of the query client
    if (!browserQueryClient) browserQueryClient = makeQueryClient();
    return browserQueryClient;
  }
}

server-actions-key-facotry.ts

type ServerActionsKeyFactory<TKey extends string[]> = {
  [key: string]: (...args: string[]) => TKey;
};

export const customCreateServerActionsKeyFactory = <
  const TKeys extends string[],
  const TFactory extends ServerActionsKeyFactory<TKeys>
>(
  factory: TFactory
) => {
  return factory;
};

keys.ts

import { customCreateServerActionsKeyFactory } from './server-actions-key-facotry';

export const QueryKeyFactory = customCreateServerActionsKeyFactory({
  someKey: (id: string) => ['some-key', id],
});

server-actions-hooks.ts

import { useInfiniteQuery, useMutation, useQuery } from '@tanstack/react-query';
import {
  setupServerActionHooks,
} from 'zsa-react-query';
import { QueryKeyFactory } from './keys';

const {
  useServerActionQuery,
  useServerActionMutation,
  useServerActionInfiniteQuery,
} = setupServerActionHooks({
  hooks: {
    useQuery: useQuery,
    useMutation: useMutation,
    useInfiniteQuery: useInfiniteQuery,
  },
  queryKeyFactory: QueryKeyFactory,
});

export {
  useServerActionInfiniteQuery,
  useServerActionMutation,
  useServerActionQuery,
};

server-actions-prefetch.ts

import { QueryClient, useQuery } from '@tanstack/react-query';
import {
  type inferServerActionError,
  type inferServerActionInput,
  type inferServerActionReturnData,
  type TAnyZodSafeFunctionHandler,
} from 'zsa';
import { getQueryClient } from './get-query-client';

export const prefetchQuery = async <
  THandler extends TAnyZodSafeFunctionHandler,
  TInitialData extends
    | undefined
    | inferServerActionReturnData<THandler>
    | (() => inferServerActionReturnData<THandler>)
>(
  action: THandler,
  options: Omit<
    Parameters<
      typeof useQuery<
        inferServerActionReturnData<THandler>,
        inferServerActionError<THandler>,
        inferServerActionReturnData<THandler>
      >
    >[0],
    'queryFn' | 'initialData'
  > & {
    input: inferServerActionInput<THandler>;
    initialData?: TInitialData;
  }
): Promise<InstanceType<typeof QueryClient>> => {
  const queryClient = getQueryClient();
  await queryClient.prefetchQuery<
    inferServerActionReturnData<THandler>,
    inferServerActionError<THandler>,
    inferServerActionReturnData<THandler>
  >({
    ...options,
    queryFn: async () => {
      const result = await action(options.input);

      if (!result) return;

      const [data, err] = result;

      if (err) {
        throw err;
      }

      return data;
    },
  });

  return queryClient;
};

Usage

import { dehydrate, HydrationBoundary } from '@tanstack/react-query';
import { QueryKeyFactory } from '../../../lib/keys';
import { prefetchQuery } from '../../../lib/server-actions-prefetch';
import { zsaServerSideAction } from './actions';

export default async function ExamplePage({
  params,
}: {
  params: { profileId: string };
}) {
  const queryClient = await prefetchQuery(zsaServerSideAction, {
    input: {
      id: parseInt(params.id),
    },
    queryKey: QueryKeyFactory.someKey(params.id),
  });
  return (
    <>
      <HydrationBoundary state={dehydrate(queryClient)}>
        <Component />
      </HydrationBoundary>
    </>
  );
}

[strangebot]

@davve5 I'd considered something similar after looking into the source code for the existing hooks, but you've done it. Great. I'll play around with this when I get a chance. Thanks.