Merge pull request #66 from jotaijs/kj/update-types

update types, move observer logic into baseAtomWithQuery

__tests__/atomWithQuery_spec.tsx

@@ -390,7 +390,10 @@ it('query with enabled (#500)', async () => {
 })
 
 it('query with initialData test', async () => {
-  const mockFetch = jest.fn((response) => ({ response }))
+  const mockFetch = jest.fn((response: { count: number }) => ({
+    response,
+  }))
+
   let resolve = () => {}
 
   const countAtom = atomWithQuery(() => ({
@@ -402,18 +405,14 @@ it('query with initialData test', async () => {
     initialData: { response: { count: 0 } },
   }))
   const Counter = () => {
-    const [countData] = useAtom(countAtom)
-    const { data, isPending, isError } = countData
-
-    if (isPending) {
-      return <>loading</>
-    }
-
-    if (isError) {
-      return <>errorred</>
-    }
+    const [
+      {
+        data: {
+          response: { count },
+        },
+      },
+    ] = useAtom(countAtom)
 
-    const count = data.response.count
     return (
       <>
         <div>count: {count}</div>

__tests__/atomWithSuspenseQuery_spec.tsx

@@ -18,7 +18,7 @@ afterEach(() => {
 it('suspense basic, suspends', async () => {
   let resolve = () => {}
   const countAtom = atomWithSuspenseQuery(() => ({
-    queryKey: ['test1'],
+    queryKey: ['test1', 'suspends'],
     queryFn: async () => {
       await new Promise<void>((r) => (resolve = r))
       return { response: { count: 0 } }

src/atomWithInfiniteQuery.ts

@@ -1,16 +1,21 @@
-import { InfiniteQueryObserver, QueryClient } from '@tanstack/query-core'
-import type {
+import {
   DefaultError,
-  DefaultedInfiniteQueryObserverOptions,
   InfiniteData,
-  InfiniteQueryObserverOptions,
-  InfiniteQueryObserverResult,
+  InfiniteQueryObserver,
+  QueryClient,
   QueryKey,
-  WithRequired,
+  QueryObserver,
 } from '@tanstack/query-core'
-import { Atom, type Getter, atom } from 'jotai/vanilla'
+import { Getter } from 'jotai'
 import { baseAtomWithQuery } from './baseAtomWithQuery'
 import { queryClientAtom } from './queryClientAtom'
+import {
+  AtomWithInfiniteQueryOptions,
+  AtomWithInfiniteQueryResult,
+  DefinedAtomWithInfiniteQueryResult,
+  DefinedInitialDataInfiniteOptions,
+  UndefinedInitialDataInfiniteOptions,
+} from './types'
 
 export function atomWithInfiniteQuery<
   TQueryFnData,
@@ -21,102 +26,59 @@ export function atomWithInfiniteQuery<
 >(
   getOptions: (
     get: Getter
-  ) => InfiniteQueryOptions<TQueryFnData, TError, TData, TQueryKey, TPageParam>,
-  getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
-): Atom<InfiniteQueryObserverResult<TData, TError>> {
-  const observerCacheAtom = atom(
-    () =>
-      new WeakMap<
-        QueryClient,
-        InfiniteQueryObserver<
-          TQueryFnData,
-          TError,
-          TData,
-          TQueryFnData,
-          TQueryKey,
-          TPageParam
-        >
-      >()
-  )
-  if (process.env.NODE_ENV !== 'production') {
-    observerCacheAtom.debugPrivate = true
-  }
-  const optionsAtom = atom((get) => {
-    const client = getQueryClient(get)
-    const options = getOptions(get)
-    const cache = get(observerCacheAtom)
-    const cachedObserver = cache.get(client)
-    const dOptions = client.defaultQueryOptions(
-      options
-    ) as DefaultedInfiniteQueryObserverOptions<
-      TQueryFnData,
-      TError,
-      TData,
-      TQueryFnData,
-      TQueryKey,
-      TPageParam
-    >
-
-    dOptions._optimisticResults = 'optimistic'
-
-    if (cachedObserver) {
-      cachedObserver.setOptions(dOptions, { listeners: false })
-    }
-
-    return dOptions
-  })
-  if (process.env.NODE_ENV !== 'production') {
-    optionsAtom.debugPrivate = true
-  }
-  const observerAtom = atom((get) => {
-    const options = get(optionsAtom)
-    const client = getQueryClient(get)
-
-    const observerCache = get(observerCacheAtom)
-
-    const cachedObserver = observerCache.get(client)
-    if (cachedObserver) return cachedObserver
-
-    const newObserver = new InfiniteQueryObserver(client, options)
-    observerCache.set(client, newObserver)
-
-    return newObserver
-  })
-
-  return baseAtomWithQuery<
+  ) => UndefinedInitialDataInfiniteOptions<
     TQueryFnData,
     TError,
     TData,
+    TQueryKey,
+    TPageParam
+  >,
+  getQueryClient?: (get: Getter) => QueryClient
+): AtomWithInfiniteQueryResult<TData, TError>
+export function atomWithInfiniteQuery<
+  TQueryFnData,
+  TError = DefaultError,
+  TData = InfiniteData<TQueryFnData>,
+  TQueryKey extends QueryKey = QueryKey,
+  TPageParam = unknown,
+>(
+  getOptions: (
+    get: Getter
+  ) => DefinedInitialDataInfiniteOptions<
     TQueryFnData,
+    TError,
+    TData,
     TQueryKey,
     TPageParam
-  >(
-    (get) => ({
-      ...get(optionsAtom),
-      suspense: false,
-    }),
-    (get) => get(observerAtom),
-    getQueryClient
-  )
-}
-
-interface InfiniteQueryOptions<
-  TQueryFnData = unknown,
+  >,
+  getQueryClient?: (get: Getter) => QueryClient
+): DefinedAtomWithInfiniteQueryResult<TData, TError>
+export function atomWithInfiniteQuery<
+  TQueryFnData,
   TError = DefaultError,
   TData = InfiniteData<TQueryFnData>,
   TQueryKey extends QueryKey = QueryKey,
   TPageParam = unknown,
-> extends WithRequired<
-    Omit<
-      InfiniteQueryObserverOptions<
-        TQueryFnData,
-        TError,
-        TData,
-        TQueryFnData,
-        TQueryKey,
-        TPageParam
-      >,
-      'suspense'
-    >,
-    'queryKey'
-  > {}
+>(
+  getOptions: (
+    get: Getter
+  ) => AtomWithInfiniteQueryOptions<
+    TQueryFnData,
+    TError,
+    TData,
+    TQueryFnData,
+    TQueryKey,
+    TPageParam
+  >,
+  getQueryClient?: (get: Getter) => QueryClient
+): AtomWithInfiniteQueryResult<TData, TError>
+export function atomWithInfiniteQuery(
+  getOptions: (get: Getter) => AtomWithInfiniteQueryOptions,
+  getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
+) {
+  return baseAtomWithQuery(
+    getOptions,
+    InfiniteQueryObserver as typeof QueryObserver,
+    getQueryClient
+  )
+}

src/atomWithMutation.ts

@@ -1,12 +1,13 @@
 import {
   MutationObserver,
-  type MutationObserverOptions,
   type MutationObserverResult,
+  MutationOptions,
   QueryClient,
 } from '@tanstack/query-core'
 import { Getter, atom } from 'jotai'
 import { make, pipe, toObservable } from 'wonka'
 import { queryClientAtom } from './queryClientAtom'
+import { AtomWithMutationResult, MutateFunction } from './types'
 import { shouldThrowError } from './utils'
 
 export function atomWithMutation<
@@ -17,9 +18,9 @@ export function atomWithMutation<
 >(
   getOptions: (
     get: Getter
-  ) => MutationObserverOptions<TData, TError, TVariables, TContext>,
+  ) => MutationOptions<TData, TError, TVariables, TContext>,
   getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
-) {
+): AtomWithMutationResult<TData, TError, TVariables, TContext> {
   const IN_RENDER = Symbol()
 
   const optionsAtom = atom((get) => {
@@ -113,9 +114,9 @@ export function atomWithMutation<
 
   const mutateAtom = atom((get) => {
     const observer = get(observerAtom)
-    const mutate = (
-      variables: TVariables,
-      options?: MutationObserverOptions<TData, TError, TVariables, TContext>
+    const mutate: MutateFunction<TData, TError, TVariables, TContext> = (
+      variables,
+      options
     ) => {
       observer.mutate(variables, options).catch(noop)
     }

src/atomWithQuery.ts

@@ -3,80 +3,54 @@ import {
   QueryClient,
   QueryKey,
   QueryObserver,
-  QueryObserverOptions,
-  QueryObserverResult,
 } from '@tanstack/query-core'
-import { Atom, Getter, atom } from 'jotai'
+import { Getter } from 'jotai'
 import { baseAtomWithQuery } from './baseAtomWithQuery'
 import { queryClientAtom } from './queryClientAtom'
+import {
+  AtomWithQueryOptions,
+  AtomWithQueryResult,
+  DefinedAtomWithQueryResult,
+  DefinedInitialDataOptions,
+  UndefinedInitialDataOptions,
+} from './types'
 
 export function atomWithQuery<
   TQueryFnData = unknown,
   TError = DefaultError,
   TData = TQueryFnData,
-  TQueryData = TQueryFnData,
   TQueryKey extends QueryKey = QueryKey,
 >(
   getOptions: (
     get: Getter
-  ) => Omit<
-    QueryObserverOptions<TQueryFnData, TError, TData, TQueryData, TQueryKey>,
-    'suspense'
-  >,
+  ) => UndefinedInitialDataOptions<TQueryFnData, TError, TData, TQueryKey>,
+  getQueryClient?: (get: Getter) => QueryClient
+): AtomWithQueryResult<TData, TError>
+export function atomWithQuery<
+  TQueryFnData = unknown,
+  TError = DefaultError,
+  TData = TQueryFnData,
+  TQueryKey extends QueryKey = QueryKey,
+>(
+  getOptions: (
+    get: Getter
+  ) => DefinedInitialDataOptions<TQueryFnData, TError, TData, TQueryKey>,
+  getQueryClient?: (get: Getter) => QueryClient
+): DefinedAtomWithQueryResult<TData, TError>
+export function atomWithQuery<
+  TQueryFnData = unknown,
+  TError = DefaultError,
+  TData = TQueryFnData,
+  TQueryKey extends QueryKey = QueryKey,
+>(
+  getOptions: (
+    get: Getter
+  ) => AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
+  getQueryClient?: (get: Getter) => QueryClient
+): AtomWithQueryResult<TData, TError>
+export function atomWithQuery(
+  getOptions: (get: Getter) => AtomWithQueryOptions,
   getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
-): Atom<QueryObserverResult<TData, TError>> {
-  const observerCacheAtom = atom(
-    () =>
-      new WeakMap<
-        QueryClient,
-        QueryObserver<TQueryFnData, TError, TData, TQueryData, TQueryKey>
-      >()
-  )
-  if (process.env.NODE_ENV !== 'production') {
-    observerCacheAtom.debugPrivate = true
-  }
-
-  const optionsAtom = atom((get) => {
-    const client = getQueryClient(get)
-    const options = getOptions(get)
-    const cache = get(observerCacheAtom)
-    const cachedObserver = cache.get(client)
-    const dOptions = client.defaultQueryOptions(options)
-
-    dOptions._optimisticResults = 'optimistic'
-
-    if (cachedObserver) {
-      cachedObserver.setOptions(dOptions, { listeners: false })
-    }
-
-    return dOptions
-  })
-  if (process.env.NODE_ENV !== 'production') {
-    optionsAtom.debugPrivate = true
-  }
-
-  const observerAtom = atom((get) => {
-    const options = get(optionsAtom)
-    const client = getQueryClient(get)
-
-    const observerCache = get(observerCacheAtom)
-
-    const cachedObserver = observerCache.get(client)
-
-    if (cachedObserver) return cachedObserver
-
-    const newObserver = new QueryObserver(client, options)
-    observerCache.set(client, newObserver)
-
-    return newObserver
-  })
-  if (process.env.NODE_ENV !== 'production') {
-    observerAtom.debugPrivate = true
-  }
-
-  return baseAtomWithQuery<TQueryFnData, TError, TData, TQueryData, TQueryKey>(
-    (get) => ({ ...get(optionsAtom), suspense: false }),
-    (get) => get(observerAtom),
-    getQueryClient
-  )
+) {
+  return baseAtomWithQuery(getOptions, QueryObserver, getQueryClient)
 }