atomWithSuspenseQuery - implementation changes suggestion

[grzesiekgs]

Hello!

First of all, thanks for creating jotai-tanstack-query, it's a great alternative for @tanstack/react-query which allows us to manage data fetching directly via atoms.

Problem statement

Consider a query which did not manage to fetch data after first mount (therefore it throws an error and triggers ErrorBoundary).
After resetting ErrorBoundary, atomWithSuspenseQuery is still in error state, and it's not possible to "automatically" reset it.

I believe that this is same issue as #32.
I realize that theres an example which uses another atom, which can be updated in order to recalculate atomWithSuspenseQuery, but this approach does not scale well.

codesandbox with example of issue and proposed solution

Theres codesandbox which has two examples - BrokenExample and FixedExample.
First one is example of the issue, while second one is using modified implementation of atomWithSuspenseQuery and appears to work correctly.
You can switch between these examples in App.tsx file.

I've left bunch of comments in fixedAtomWithSuspenseQuery file, explaining the approach. This atom is also available in snippet below in order to allow easier discussion on github.

Proposed solution

fixedAtomWithSuspenseQuery snippet

// Code from jotai-tanstack-query
export const willFetch = (
  result: QueryObserverResult<any, any>,
  isRestoring: boolean,
) => result.isPending && !isRestoring;

export const shouldSuspend = (
  defaultedOptions:
    | DefaultedQueryObserverOptions<any, any, any, any, any>
    | undefined,
  result: QueryObserverResult<any, any>,
  isRestoring: boolean,
) => defaultedOptions?.suspense && willFetch(result, isRestoring);

export const getHasError = <
  TData,
  TError,
  TQueryFnData,
  TQueryData,
  TQueryKey extends QueryKey,
>({
  result,
  throwOnError,
  query,
}: {
  result: QueryObserverResult<TData, TError>;
  throwOnError:
    | ThrowOnError<TQueryFnData, TError, TQueryData, TQueryKey>
    | undefined;
  query: Query<TQueryFnData, TError, TQueryData, TQueryKey>;
}) =>
  result.isError &&
  !result.isFetching &&
  shouldThrowError(throwOnError, [result.error, query]);

export function shouldThrowError<T extends (...args: any[]) => boolean>(
  throwOnError: boolean | T | undefined,
  params: Parameters<T>,
): boolean {
  // Allow useErrorBoundary function to override throwing behavior on a per-error basis
  if (typeof throwOnError === "function") {
    return throwOnError(...params);
  }

  return !!throwOnError;
}

export const ensureStaleTime = (
  defaultedOptions: DefaultedQueryObserverOptions<any, any, any, any, any>,
) => {
  if (defaultedOptions.suspense) {
    if (typeof defaultedOptions.staleTime !== "number") {
      return {
        ...defaultedOptions,
        staleTime: 1000,
      };
    }
  }

  return defaultedOptions;
};

export function baseAtomWithQuery<
  TQueryFnData,
  TError,
  TData,
  TQueryData,
  TQueryKey extends QueryKey,
>(
  getOptions: (
    get: Getter,
  ) => BaseAtomWithQueryOptions<
    TQueryFnData,
    TError,
    TData,
    TQueryData,
    TQueryKey
  >,
  Observer: typeof QueryObserver,
  getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom),
): Atom<
  | QueryObserverResult<TData, TError>
  | Promise<QueryObserverResult<TData, TError>>
> {
  const clientAtom = atom(getQueryClient);
  const observerCacheAtom = atom(
    () =>
      new WeakMap<
        QueryClient,
        QueryObserver<TQueryFnData, TError, TData, TQueryData, TQueryKey>
      >(),
  );
  const defaultedOptionsAtom = atom((get) => {
    const client = get(clientAtom);
    const options = getOptions(get);
    const defaultedOptions = client.defaultQueryOptions(options);

    const cache = get(observerCacheAtom);
    const cachedObserver = cache.get(client);

    defaultedOptions._optimisticResults = "optimistic";

    if (cachedObserver) {
      cachedObserver.setOptions(defaultedOptions, {
        listeners: false,
      });
    }

    return ensureStaleTime(defaultedOptions);
  });
  const observerAtom = atom((get) => {
    const client = get(clientAtom);
    const defaultedOptions = get(defaultedOptionsAtom);

    const observerCache = get(observerCacheAtom);

    const cachedObserver = observerCache.get(client);

    if (cachedObserver) {
      return cachedObserver;
    }

    const newObserver = new Observer<
      TQueryFnData,
      TError,
      TData,
      TQueryData,
      TQueryKey
    >(client, defaultedOptions);
    observerCache.set(client, newObserver);

    return newObserver;
  });
  // Code different than jotai-tanstack-query
  const queryResultGetterAtom = atom((get) => {
    const observer = get(observerAtom);
    // When this atom has value other than null, it means that query result
    // has been obtained either by optimisticResultAtom or query subscription defined in onMount below.
    // Value is set to null when query is resetting.
    const queryResultAtom = atom<QueryObserverResult<TData, TError> | null>(
      null,
    );
    // Subscribe to query result updates. Note that this happens only when query did not throw promise or error
    // (therefore it has to have valid result)
    queryResultAtom.onMount = (set) => {
      // tanstack/react-query useBaseQuery is wrapping 'set' with
      // https://tanstack.com/query/latest/docs/reference/notifyManager#notifymanagerbatchcalls
      // not sure is it relevant here.
      const unsubscribe = observer.subscribe(set);
      // comment and line from tanstack/react-query useBaseQuery:
      // Update result to make sure we did not miss any query updates
      // between creating the observer and subscribing to it.
      observer.updateResult();

      return () => {
        // Is there any case where it makes sense to reset value on queryAtom unmount?
        // I think that it will get garbage collected due to WeakMap usage by jotai.
        // Doing this could result in returning outdated state in returnAtom.
        // set(null);
        unsubscribe();
      };
    };

    return queryResultAtom;
  });
  // Simple counter which allows to retrigger optimisticResultAtom calculation, after resetting atom.
  const optimisticResultResetAtom = atom(0);
  // Request optimistic result from observer. Decide whether to suspend or not.
  // This atom will recalculate only once per atom reset - see returnAtom.
  const optimisticResultAtom = atom<
    | Promise<QueryObserverResult<TData, TError>>
    | QueryObserverResult<TData, TError>,
    [QueryObserverResult<TData, TError>],
    void
  >(
    (get, { setSelf }) => {
      console.log("queryFetcherAtom");
      const observer = get(observerAtom);
      const defaultedOptions = get(defaultedOptionsAtom);
      // Recalculate atom when query is resetted.
      get(optimisticResultResetAtom);

      const result = observer.getOptimisticResult(defaultedOptions);

      if (!shouldSuspend(defaultedOptions, result, false)) {
        // Update queryResultAtom with sync result.
        Promise.resolve(result).then(setSelf);

        return result;
      }

      return (
        observer
          .fetchOptimistic(defaultedOptions)
          .then((succeedResult) => {
            setSelf(succeedResult);

            return succeedResult;
          })
          // useSuspenseQuery is catching fetchOptimistic error, and triggers error boundary.
          // Later, when error boundary is resetted, it means that useSuspenseQuery is mounting, and it's resetting the query.
          // (therefore for useSuspenseQuery, when error boundary is active then query.state.status === 'error'? Not sure is it like this)
          //
          // Jotai also has to catch error, because if error would be set as atom value,
          // then recovering from error is possible only by outside call - see https://github.com/jotaijs/jotai-tanstack-query/issues/32#issue-1684399582
          // This is different than jotai-tanstack-query current implementation.
          .catch(() => {
            // Since fetchOptimistic failed, error has to be thrown, but let's do it same way as tanstack/react-query useBaseQuery.
            // Error handling is done in returnAtom, but first, current query result must be obtained.
            //
            // observer.currentResult() has outdated value, and observer.updateResult() doesn't help,
            // (I'm guessing that's due how observer.getOptimisticResult behaves with suspense = true)
            // therefore get optimisticResult but without activating query (_optimisticResults: undefined [I think that's how it works])
            const erroredResult = observer.getOptimisticResult({
              ...defaultedOptions,
              _optimisticResults: undefined,
            });

            setSelf(erroredResult);

            return erroredResult;
          })
      );
    },
    (get, set, result) => {
      // Update queryResultAtom.
      set(get(queryResultGetterAtom), result);
    },
  );
  // Better name or just return without variable assignment?
  const returnAtom = atom<
    | QueryObserverResult<TData, TError>
    | Promise<QueryObserverResult<TData, TError>>,
    [],
    Promise<void>
  >(
    (get, { setSelf }) => {
      const result = get(get(queryResultGetterAtom));
      // If queryResultAtom has null value it means that value has not yet been obtained from observer.
      // Therefore read optimisticResultAtom which activates query and decides whether to suspend or not.
      if (result === null) {
        // Notice that this atom has valid (up to date) value only just after mounting or resetting queryAtom,
        // Later, valid value is held in queryResultAtom.
        return get(optimisticResultAtom);
      }
      // At this point, for atomWithSuspenseQuery it's known that fetchOptimistic promise has been resolved,
      // but it could be resolved with error, which is not yet handled.
      // Same goes for atomWithQuery (but it could throw error only by using custom throwOnError).
      // Therefore verify does query has error or not.

      // Theres potential to retrieve pre 0.8 feature of having async query options getter, because we only need options.throwOnError there
      // which is non-async function, but this only makes sense for atomWithSuspenseQuery - would have to think about it.
      const options = get(defaultedOptionsAtom);
      const query = get(observerAtom).getCurrentQuery();

      if (
        getHasError({
          result,
          query,
          throwOnError: options.throwOnError,
        })
      ) {
        // Atom reset has to be scheduled just before throwing error, because atom needs to somehow recover from error boundary,
        // and if left in error state after unmounting component, thats not possible without outside call.
        // If needed that outside call is still possible by using useSetAtom(yourAtomWithQuery),
        // but now it's easier to maintain and does not require manually resetting query.

        // While using Promise.resolve().then(setSelf) or wait(0), returnAtom was recalculating before error boundary caught error,
        // Maybe @dai-shi could check am I not doing some bad things out there.
        wait(1).then(setSelf);
        // eslint-disable-next-line @typescript-eslint/no-throw-literal
        throw result.error;
      }
      // Non-promise result without error - return it.
      return result;
    },
    async (get, set) => {
      // Not sure should it reset or validate query.
      get(observerAtom).getCurrentQuery().reset();
      // Reset to initial state and allow to recalculate optimisticResultAtom
      set(get(queryResultGetterAtom), null);
      set(optimisticResultResetAtom, get(optimisticResultResetAtom) + 1);
    },
  );

  return returnAtom;
}

export function atomWithSuspenseQueryFixed<
  TQueryFnData = unknown,
  TError = DefaultError,
  TData = TQueryFnData,
  TQueryKey extends QueryKey = QueryKey,
>(
  getOptions: (
    get: Getter,
  ) => Omit<
    AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
    "enabled" | "placeholderData"
  >,
  getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom),
): Atom<AtomWithSuspenseQueryResult<TData, TError>> {
  const optionsAtom = atom<
    Omit<
      AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
      "enabled" | "placeholderData"
    >
  >((get) => {
    const options = getOptions(get);
    return {
      // defaultThrowOnError - useQuery does not use it, useSuspenseQuery does.
      throwOnError: (_, query) => typeof query.state.data === "undefined",
      ...options,
      suspense: true,
      enabled: true,
    };
  });

  return baseAtomWithQuery(
    (get) => get(optionsAtom),
    QueryObserver,
    getQueryClient,
  ) as Atom<AtomWithSuspenseQueryResult<TData, TError>>;
}

export function atomWithQueryFixed<
  TQueryFnData = unknown,
  TError = DefaultError,
  TData = TQueryFnData,
  TQueryKey extends QueryKey = QueryKey,
>(
  getOptions: (
    get: Getter,
  ) => AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
  getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom),
): Atom<AtomWithQueryResult<TData, TError>> {
  const optionsAtom = atom<
    Omit<
      AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
      "enabled" | "placeholderData"
    >
  >((get) => {
    const options = getOptions(get);

    return {
      // defaultThrowOnError - useQuery does not use it, useSuspenseQuery does.
      throwOnError: (_, query) => typeof query.state.data === "undefined",
      ...options,
      suspense: false,
    };
  });
  return baseAtomWithQuery(
    (get) => get(optionsAtom),
    QueryObserver,
    getQueryClient,
  ) as Atom<AtomWithQueryResult<TData, TError>>;
}

Since english is not my main language, I suspect that some comments could be unclear, let me know if that's the case please.

In addition:

1. Theres possibility to retrieve pre 0.8 feature - async query options getter, Dependent Suspense Queries #65, I'm willing to work on it if You will consider my approach as valid.
2. Theres new functionality which allows to reset query by atom setter (not included in example).
3. I was playing around with throwOnError in order to resolve a specific case that I have. But my suggested change to useSuspenseQuery got rejected, therefore I am going to revert these changes. Ignore that for now please.
   useSuspenseQuery - throw error while having stale data TanStack/query#6960

[kalijonn]

You are amazing! Suspense Queries not being able to reset error boundaries has been one of the most annoying bugs I've ever seen. I've gone through the explanation in your comments and the working example you've linked. I haven't tested the code yet but I'll do that this week. Meanwhile if you want to raise a PR feel free to do so, and I can add my comments from my local testing if any.

[kalijonn]

Theres possibility to retrieve pre 0.8 feature - async query options getter, #65, I'm willing to work on it if You will consider my approach as valid.

I was hoping to not create a new atom and have it as a part of the existing atomWithSuspenseQuery if possible but I couldn't spend a lot of time on this. I would love to see your idea.

[feledori]

Just wanted to open an issue about the types being wrong as enabled is still allowed in a atomWithSuspenseQuery. Which it shouldn't be. But I see you fixed that as well, nice!

[grzesiek-ds]

@kalijonn sorry for the two weeks silence 🙈 I will create a PR in next few days.

[feledori]

@grzesiek-ds Sick!
@kalijonn If bugs like this are know, having a section in the readme outlining these issue would be a nice to have. You could use that as a platform for contributions.

[kalijonn]

enabled is still allowed in a atomWithSuspenseQuery

@feledori I think it's because of type widening. In some cases, typescript doesn't care about additional type parameters. Something like this: https://stackoverflow.com/questions/54977550/what-is-widening-of-a-function-return-type-in-typescript

Omit exists in the existing implementation but doesn't throw an error unless you cast it explicitly.

const userAtom = atomWithSuspenseQuery<object>((get) => ({
  queryKey: ['users', get(idAtom)],
  queryFn: async ({ queryKey: [, id] }) => {
    const res = await fetch(`https://jsonplaceholder.typicode.com/users/${id}`)
    return res.json()
  },
  enabled: false, //no error
}))

const userAtom2 = atomWithSuspenseQuery<object>((get) => {
  return {
    queryKey: ['users', get(idAtom)],
    queryFn: async ({ queryKey: [, id] }) => {
      const res = await fetch(
        `https://jsonplaceholder.typicode.com/users/${id}`
      )
      return res.json()
    },
    enabled: false, // error - Object literal may only specify known properties, and 'enabled' does not exist in type 'AtomWithSuspenseQueryOptions<unknown, Error, unknown, QueryKey>'
  } satisfies AtomWithSuspenseQueryOptions
})