request for comment: atomFamilyWithQuery

[aulneau]

I use this in almost all of my projects and was curious if we should make it an official utility? this is an evolution of #412.

This version would need to handle initialData for SSR uses.

@dai-shi should i make a PR?

import { atomFamily } from 'jotai/utils';
import { atomWithQuery, queryClientAtom } from 'jotai/query';
import deepEqual from 'fast-deep-equal';

import type { Getter, WritableAtom } from 'jotai';
import type { QueryObserverOptions } from 'react-query';
import { atom } from 'jotai';

export const atomFamilyWithQuery = <Param, Data>(
  key: string,
  queryFn: (get: Getter, param: Param) => Data | Promise<Data>,
  options: {
    equalityFn?: (a: Data, b: Data) => boolean;
  } & QueryObserverOptions = {}
): ((param: Param) => WritableAtom<Data, null>) => {
  const { equalityFn = deepEqual, ...rest } = options;
  return atomFamily(param => {
    const queryKey = [key, param];
    const dataAtom = atomWithQuery(get => {
      return {
        queryKey,
        queryFn: () => queryFn(get, param),
        useErrorBoundary: true,
        keepPreviousData: true,
        refetchOnReconnect: true,
        refetchOnWindowFocus: true,
        refetchOnMount: true,
        ...(rest as any), // TODO: fix type cast
      };
    }, equalityFn);
    dataAtom.debugLabel = `atomFamilyWithQuery/dataAtom/${JSON.stringify(queryKey)}`;

    const anAtom = atom(
      get => get(dataAtom),
      async get => {
        const queryClient = get(queryClientAtom);
        await queryClient.refetchQueries({
          queryKey,
        });
      }
    );
    anAtom.debugLabel = `atomFamilyWithQuery/${JSON.stringify(queryKey)}`;
    return anAtom;
  }, deepEqual);
};

and an example of how this is used

// our query atom that takes params
export const accountInfoResponseState = atomFamilyWithQuery<
  [string, string],
  { balance: BN; nonce: number }
>(
  AccountQueryKeys.ACCOUNT_INFO_STATE, // key param
  async (get, [principal]) => {
    const client = get(accountsApiClientState);
    const data = await client.getAccountInfo({
      principal,
      proof: 0,
    });
    return {
      balance: new BN(data.balance.slice(2), 16),
      nonce: data.nonce,
    };
  },
  { refetchInterval: QueryRefreshRates.QUICK } // query atom options
);

// uses the query atom and passes these params to it
export const accountInfoState = atom(get => {
  const { network, principal } = get(
    waitForAll({
      network: currentNetworkState,
      principal: currentAccountStxAddressState,
    })
  );
  if (!address) return;
  return get(accountInfoResponseState([principal, network.url])); // implements the query atom defined above
});

and in a component

const Component = () => {
  const { principal, network } = useSomeOtherAtomThatProvidesTheseValues();
  const [data, refresh] = useAtom(accountInfoResponseState([principal, network.url]))
  return ...
}

[aulneau]

I wonder if an API for initialData could be something like

{ 
  initialData: (get) => get(someOtherAtom); // this would be pre-populated via Provider for SSR uses
}

[dai-shi]

I know this would be useful, but my feeling is this seems domain-specific application code.
To make it library code, we'd need to remove some assumptions. I'm not sure if you are happy with the result.
One thing I notice this util helps is, suggesting query key structure that allows refetching all. So, I'd probably suggest a recipe in the docs.

(btw, I think we should first implement "error handling" in atomWithQuery.)