From 3629b1ca7f205e73f5f862ba23ac73e677d7cbfd Mon Sep 17 00:00:00 2001
From: kalijonn <43421621+kalijonn@users.noreply.github.com>
Date: Mon, 29 Jan 2024 03:15:06 +0530
Subject: [PATCH] update types, move some of the observer logic into
 baseAtomQuery

---
 __tests__/atomWithQuery_spec.tsx         |  23 ++-
 __tests__/atomWithSuspenseQuery_spec.tsx |   2 +-
 src/atomWithInfiniteQuery.ts             | 154 ++++++---------
 src/atomWithMutation.ts                  |  13 +-
 src/atomWithQuery.ts                     | 100 ++++------
 src/atomWithSuspenseInfiniteQuery.ts     | 134 +++----------
 src/atomWithSuspenseQuery.ts             |  91 ++-------
 src/baseAtomWithQuery.ts                 | 216 ++++++++-------------
 src/types.ts                             | 230 +++++++++++++++++++++++
 src/utils.ts                             |  26 +++
 10 files changed, 490 insertions(+), 499 deletions(-)
 create mode 100644 src/types.ts

diff --git a/__tests__/atomWithQuery_spec.tsx b/__tests__/atomWithQuery_spec.tsx
index f337963..86de580 100644
--- a/__tests__/atomWithQuery_spec.tsx
+++ b/__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>
diff --git a/__tests__/atomWithSuspenseQuery_spec.tsx b/__tests__/atomWithSuspenseQuery_spec.tsx
index f7bcfac..21db26b 100644
--- a/__tests__/atomWithSuspenseQuery_spec.tsx
+++ b/__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 } }
diff --git a/src/atomWithInfiniteQuery.ts b/src/atomWithInfiniteQuery.ts
index d3189bf..daffd4a 100644
--- a/src/atomWithInfiniteQuery.ts
+++ b/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
+  )
+}
diff --git a/src/atomWithMutation.ts b/src/atomWithMutation.ts
index c15e85c..5f052f1 100644
--- a/src/atomWithMutation.ts
+++ b/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)
     }
diff --git a/src/atomWithQuery.ts b/src/atomWithQuery.ts
index c74251a..ff5edef 100644
--- a/src/atomWithQuery.ts
+++ b/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)
 }
diff --git a/src/atomWithSuspenseInfiniteQuery.ts b/src/atomWithSuspenseInfiniteQuery.ts
index 3033b9d..68a24da 100644
--- a/src/atomWithSuspenseInfiniteQuery.ts
+++ b/src/atomWithSuspenseInfiniteQuery.ts
@@ -1,19 +1,22 @@
 import {
   DefaultError,
-  DefaultedInfiniteQueryObserverOptions,
   InfiniteData,
   InfiniteQueryObserver,
-  InfiniteQueryObserverOptions,
-  InfiniteQueryObserverSuccessResult,
   QueryClient,
-  type QueryKey,
+  QueryKey,
+  QueryObserver,
 } from '@tanstack/query-core'
-import { Atom, Getter, atom } from 'jotai'
+import { Getter, atom } from 'jotai'
 import { baseAtomWithQuery } from './baseAtomWithQuery'
 import { queryClientAtom } from './queryClientAtom'
+import {
+  AtomWithSuspenseInfiniteQueryOptions,
+  AtomWithSuspenseInfiniteQueryResult,
+} from './types'
+import { defaultThrowOnError } from './utils'
 
-export const atomWithSuspenseInfiniteQuery = <
-  TQueryFnData = unknown,
+export function atomWithSuspenseInfiniteQuery<
+  TQueryFnData,
   TError = DefaultError,
   TData = InfiniteData<TQueryFnData>,
   TQueryKey extends QueryKey = QueryKey,
@@ -21,116 +24,29 @@ export const atomWithSuspenseInfiniteQuery = <
 >(
   getOptions: (
     get: Getter
-  ) => SuspenseInfiniteQueryOptions<
+  ) => AtomWithSuspenseInfiniteQueryOptions<
     TQueryFnData,
     TError,
-    TPageParam,
     TData,
-    TQueryKey
+    TQueryFnData,
+    TQueryKey,
+    TPageParam
   >,
   getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
-): Atom<Promise<InfiniteQueryObserverSuccessResult<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)
+): AtomWithSuspenseInfiniteQueryResult<TData, TError> {
+  const suspenseOptionsAtom = atom((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 (typeof dOptions.staleTime !== 'number') {
-      dOptions.staleTime = 1000
-    }
-
-    if (cachedObserver) {
-      cachedObserver.setOptions(dOptions, { listeners: false })
+    return {
+      ...options,
+      enabled: true,
+      suspense: true,
+      throwOnError: defaultThrowOnError,
     }
-
-    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
   })
-  if (process.env.NODE_ENV !== 'production') {
-    observerAtom.debugPrivate = true
-  }
 
-  return baseAtomWithQuery<
-    TQueryFnData,
-    TError,
-    TData,
-    TQueryFnData,
-    TQueryKey,
-    TPageParam
-  >(
-    (get) => ({
-      ...get(optionsAtom),
-      suspense: true,
-      enabled: true,
-    }),
-    (get) => get(observerAtom),
+  return baseAtomWithQuery(
+    (get: Getter) => get(suspenseOptionsAtom),
+    InfiniteQueryObserver as typeof QueryObserver,
     getQueryClient
-  )
+  ) as unknown as AtomWithSuspenseInfiniteQueryResult<TData, TError>
 }
-
-interface SuspenseInfiniteQueryOptions<
-  TQueryFnData = unknown,
-  TError = DefaultError,
-  TPageParam = unknown,
-  TData = InfiniteData<TQueryFnData>,
-  TQueryKey extends QueryKey = QueryKey,
-> extends Omit<
-    InfiniteQueryObserverOptions<
-      TQueryFnData,
-      TError,
-      TData,
-      TQueryFnData,
-      TQueryKey,
-      TPageParam
-    >,
-    'enabled' | 'throwOnError' | 'placeholderData'
-  > {}
diff --git a/src/atomWithSuspenseQuery.ts b/src/atomWithSuspenseQuery.ts
index b20e65a..e9e0763 100644
--- a/src/atomWithSuspenseQuery.ts
+++ b/src/atomWithSuspenseQuery.ts
@@ -1,93 +1,42 @@
 import {
   DefaultError,
-  type DefinedQueryObserverResult,
   QueryClient,
-  type QueryKey,
+  QueryKey,
   QueryObserver,
-  type QueryObserverOptions,
 } from '@tanstack/query-core'
-import { Atom, Getter, atom } from 'jotai'
+import { Getter, atom } from 'jotai'
 import { baseAtomWithQuery } from './baseAtomWithQuery'
 import { queryClientAtom } from './queryClientAtom'
+import {
+  AtomWithSuspenseQueryOptions,
+  AtomWithSuspenseQueryResult,
+} from './types'
+import { defaultThrowOnError } from './utils'
 
-export const atomWithSuspenseQuery = <
+export function atomWithSuspenseQuery<
   TQueryFnData = unknown,
   TError = DefaultError,
   TData = TQueryFnData,
-  TQueryData = TQueryFnData,
   TQueryKey extends QueryKey = QueryKey,
 >(
   getOptions: (
     get: Getter
-  ) => Omit<
-    QueryObserverOptions<TQueryFnData, TError, TData, TQueryData, TQueryKey>,
-    'suspense' | 'enabled'
-  >,
+  ) => AtomWithSuspenseQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
   getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
-): Atom<
-  | DefinedQueryObserverResult<TData, TError>
-  | Promise<DefinedQueryObserverResult<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)
+): AtomWithSuspenseQueryResult<TData, TError> {
+  const suspenseOptions = atom((get) => {
     const options = getOptions(get)
-    const cache = get(observerCacheAtom)
-    const cachedObserver = cache.get(client)
-    const dOptions = client.defaultQueryOptions(options)
-
-    dOptions._optimisticResults = 'optimistic'
-
-    if (typeof dOptions.staleTime !== 'number') {
-      dOptions.staleTime = 1000
-    }
-
-    if (cachedObserver) {
-      cachedObserver.setOptions(dOptions, { listeners: false })
+    return {
+      ...options,
+      suspense: true,
+      enabled: true,
+      throwOnError: defaultThrowOnError,
     }
-
-    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: true,
-      enabled: true,
-    }),
-    (get) => get(observerAtom),
+  return baseAtomWithQuery(
+    (get: Getter) => get(suspenseOptions),
+    QueryObserver,
     getQueryClient
-  )
+  ) as AtomWithSuspenseQueryResult<TData, TError>
 }
diff --git a/src/baseAtomWithQuery.ts b/src/baseAtomWithQuery.ts
index ba2419a..112ba9b 100644
--- a/src/baseAtomWithQuery.ts
+++ b/src/baseAtomWithQuery.ts
@@ -1,179 +1,113 @@
 import {
-  DefaultError,
-  DefaultedInfiniteQueryObserverOptions,
-  DefaultedQueryObserverOptions,
-  InfiniteQueryObserver,
-  InfiniteQueryObserverResult,
-  InfiniteQueryObserverSuccessResult,
   QueryClient,
   QueryKey,
   QueryObserver,
   QueryObserverResult,
-  QueryObserverSuccessResult,
 } from '@tanstack/query-core'
 import { Atom, Getter, atom } from 'jotai'
-import { make, pipe, toObservable } from 'wonka'
 import { queryClientAtom } from './queryClientAtom'
-import { getHasError, shouldSuspend } from './utils'
+import { BaseAtomWithQueryOptions } from './types'
+import { ensureStaleTime, getHasError, shouldSuspend } from './utils'
 
 export function baseAtomWithQuery<
-  TQueryFnData = unknown,
-  TError = DefaultError,
-  TData = TQueryFnData,
-  TQueryData = TQueryFnData,
-  TQueryKey extends QueryKey = QueryKey,
+  TQueryFnData,
+  TError,
+  TData,
+  TQueryData,
+  TQueryKey extends QueryKey,
 >(
   getOptions: (
     get: Getter
-  ) => DefaultedQueryObserverOptions<
+  ) => BaseAtomWithQueryOptions<
     TQueryFnData,
     TError,
     TData,
     TQueryData,
     TQueryKey
-  > & { suspense: true; enabled: true },
-  getObserver: (
-    get: Getter
-  ) => QueryObserver<TQueryFnData, TError, TData, TQueryData, TQueryKey>,
-  getQueryClient?: (get: Getter) => QueryClient
-): Atom<Promise<QueryObserverSuccessResult<TData, TError>>>
-export function baseAtomWithQuery<
-  TQueryFnData = unknown,
-  TError = DefaultError,
-  TData = TQueryFnData,
-  TQueryData = TQueryFnData,
-  TQueryKey extends QueryKey = QueryKey,
->(
-  getOptions: (
-    get: Getter
-  ) => DefaultedQueryObserverOptions<
-    TQueryFnData,
-    TError,
-    TData,
-    TQueryData,
-    TQueryKey
-  > & { suspense: false },
-  getObserver: (
-    get: Getter
-  ) => QueryObserver<TQueryFnData, TError, TData, TQueryData, TQueryKey>,
-  getQueryClient?: (get: Getter) => QueryClient
-): Atom<QueryObserverResult<TData, TError>>
-export function baseAtomWithQuery<
-  TQueryFnData = unknown,
-  TError = DefaultError,
-  TData = TQueryFnData,
-  TQueryData = TQueryFnData,
-  TQueryKey extends QueryKey = QueryKey,
-  TPageParam = unknown,
->(
-  getOptions: (
-    get: Getter
-  ) => DefaultedInfiniteQueryObserverOptions<
-    TQueryFnData,
-    TError,
-    TData,
-    TQueryData,
-    TQueryKey,
-    TPageParam
-  > & { suspense: true; enabled: true },
-  getObserver: (
-    get: Getter
-  ) => InfiniteQueryObserver<
-    TQueryFnData,
-    TError,
-    TData,
-    TQueryFnData,
-    TQueryKey,
-    TPageParam
   >,
-  getQueryClient?: (get: Getter) => QueryClient
-): Atom<Promise<InfiniteQueryObserverSuccessResult<TData, TError>>>
-export function baseAtomWithQuery<
-  TQueryFnData = unknown,
-  TError = DefaultError,
-  TData = TQueryFnData,
-  TQueryData = TQueryFnData,
-  TQueryKey extends QueryKey = QueryKey,
-  TPageParam = unknown,
->(
-  getOptions: (
-    get: Getter
-  ) => DefaultedInfiniteQueryObserverOptions<
-    TQueryFnData,
-    TError,
-    TData,
-    TQueryData,
-    TQueryKey,
-    TPageParam
-  > & { suspense: false },
-  getObserver: (
-    get: Getter
-  ) => InfiniteQueryObserver<
-    TQueryFnData,
-    TError,
-    TData,
-    TQueryFnData,
-    TQueryKey,
-    TPageParam
-  >,
-  getQueryClient?: (get: Getter) => QueryClient
-): Atom<InfiniteQueryObserverResult<TData, TError>>
-export function baseAtomWithQuery<
-  TQueryFnData = unknown,
-  TError = DefaultError,
-  TData = TQueryFnData,
-  TQueryData = TQueryFnData,
-  TQueryKey extends QueryKey = QueryKey,
->(
-  getOptions: (
-    get: Getter
-  ) => DefaultedQueryObserverOptions<
-    TQueryFnData,
-    TError,
-    TData,
-    TQueryData,
-    TQueryKey
-  >,
-  getObserver: (
-    get: Getter
-  ) => QueryObserver<TQueryFnData, TError, TData, TQueryData, TQueryKey>,
+  Observer: typeof QueryObserver,
   getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
-) {
+): Atom<
+  | QueryObserverResult<TData, TError>
+  | Promise<QueryObserverResult<TData, TError>>
+> {
   const resetAtom = atom(0)
   if (process.env.NODE_ENV !== 'production') {
     resetAtom.debugPrivate = true
   }
 
-  const observableAtom = atom((get) => {
-    const observer = getObserver(get)
-    const source = make<QueryObserverResult<TData, TError>>(({ next }) => {
-      const callback = (result: QueryObserverResult<TData, TError>) => {
-        next(result)
-      }
+  const clientAtom = atom(getQueryClient)
+  if (process.env.NODE_ENV !== 'production') {
+    clientAtom.debugPrivate = true
+  }
+
+  const observerCacheAtom = atom(
+    () =>
+      new WeakMap<
+        QueryClient,
+        QueryObserver<TQueryFnData, TError, TData, TQueryData, TQueryKey>
+      >()
+  )
+  if (process.env.NODE_ENV !== 'production') {
+    observerCacheAtom.debugPrivate = true
+  }
 
-      return observer.subscribe(callback)
-    })
-    return pipe(source, toObservable)
+  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)
+  })
+  if (process.env.NODE_ENV !== 'production') {
+    defaultedOptionsAtom.debugPrivate = true
+  }
+
+  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(client, defaultedOptions)
+    observerCache.set(client, newObserver)
+
+    return newObserver
   })
   if (process.env.NODE_ENV !== 'production') {
-    observableAtom.debugPrivate = true
+    observerAtom.debugPrivate = true
   }
 
   const dataAtom = atom((get) => {
-    const observer = getObserver(get)
-    const observable = get(observableAtom)
+    const observer = get(observerAtom)
 
     const currentResult = observer.getCurrentResult()
+
     const resultAtom = atom(currentResult)
     if (process.env.NODE_ENV !== 'production') {
       resultAtom.debugPrivate = true
     }
 
     resultAtom.onMount = (set) => {
-      const { unsubscribe } = observable.subscribe((state) => {
+      const unsubscribe = observer.subscribe((state) => {
         set(state)
       })
-      return () => unsubscribe()
+      return unsubscribe
     }
 
     return resultAtom
@@ -183,8 +117,8 @@ export function baseAtomWithQuery<
   }
 
   return atom((get) => {
-    const observer = getObserver(get)
-    const options = getOptions(get)
+    const observer = get(observerAtom)
+    const defaultedOptions = get(defaultedOptionsAtom)
 
     const client = getQueryClient(get)
 
@@ -199,17 +133,17 @@ export function baseAtomWithQuery<
     get(resetAtom)
     get(get(dataAtom))
 
-    const result = observer.getOptimisticResult(options)
+    const result = observer.getOptimisticResult(defaultedOptions)
 
-    if (shouldSuspend(options, result, false)) {
-      return observer.fetchOptimistic(options)
+    if (shouldSuspend(defaultedOptions, result, false)) {
+      return observer.fetchOptimistic(defaultedOptions)
     }
 
     if (
       getHasError({
         result,
         query: observer.getCurrentQuery(),
-        throwOnError: options.throwOnError,
+        throwOnError: defaultedOptions.throwOnError,
       })
     ) {
       throw result.error
diff --git a/src/types.ts b/src/types.ts
new file mode 100644
index 0000000..c0c6be4
--- /dev/null
+++ b/src/types.ts
@@ -0,0 +1,230 @@
+import {
+  DefaultError,
+  DefinedInfiniteQueryObserverResult,
+  DefinedQueryObserverResult,
+  InfiniteData,
+  InfiniteQueryObserverOptions,
+  InfiniteQueryObserverResult,
+  MutationObserverOptions,
+  MutationObserverResult,
+  QueryKey,
+  MutateFunction as QueryMutateFunction,
+  QueryObserverOptions,
+  QueryObserverResult,
+  WithRequired,
+} from '@tanstack/query-core'
+import { Atom } from 'jotai'
+
+type Override<A, B> = { [K in keyof A]: K extends keyof B ? B[K] : A[K] }
+
+export type MutateFunction<
+  TData = unknown,
+  TError = DefaultError,
+  TVariables = void,
+  TContext = unknown,
+> = (
+  ...args: Parameters<QueryMutateFunction<TData, TError, TVariables, TContext>>
+) => void
+
+export type MutateAsyncFunction<
+  TData = unknown,
+  TError = DefaultError,
+  TVariables = void,
+  TContext = unknown,
+> = QueryMutateFunction<TData, TError, TVariables, TContext>
+
+export type AtomWithMutationResult<TData, TError, TVariables, TContext> = Atom<
+  Override<
+    MutationObserverResult<TData, TError, TVariables, TContext>,
+    { mutate: MutateFunction<TData, TError, TVariables, TContext> }
+  > & {
+    mutateAsync: MutateAsyncFunction<TData, TError, TVariables, TContext>
+  }
+>
+
+export interface MutationOptions<
+  TData = unknown,
+  TError = DefaultError,
+  TVariables = void,
+  TContext = unknown,
+> extends Omit<
+    MutationObserverOptions<TData, TError, TVariables, TContext>,
+    '_defaulted' | 'variables'
+  > {}
+
+export interface BaseAtomWithQueryOptions<
+  TQueryFnData = unknown,
+  TError = DefaultError,
+  TData = TQueryFnData,
+  TQueryData = TQueryFnData,
+  TQueryKey extends QueryKey = QueryKey,
+> extends WithRequired<
+    QueryObserverOptions<TQueryFnData, TError, TData, TQueryData, TQueryKey>,
+    'queryKey'
+  > {}
+
+export interface AtomWithQueryOptions<
+  TQueryFnData = unknown,
+  TError = DefaultError,
+  TData = TQueryFnData,
+  TQueryKey extends QueryKey = QueryKey,
+> extends Omit<
+    WithRequired<
+      BaseAtomWithQueryOptions<
+        TQueryFnData,
+        TError,
+        TData,
+        TQueryFnData,
+        TQueryKey
+      >,
+      'queryKey'
+    >,
+    'suspense'
+  > {}
+
+export interface AtomWithSuspenseQueryOptions<
+  TQueryFnData = unknown,
+  TError = DefaultError,
+  TData = TQueryFnData,
+  TQueryKey extends QueryKey = QueryKey,
+> extends Omit<
+    AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
+    'enabled' | 'throwOnError' | 'placeholderData'
+  > {}
+
+export interface AtomWithInfiniteQueryOptions<
+  TQueryFnData = unknown,
+  TError = DefaultError,
+  TData = TQueryFnData,
+  TQueryData = TQueryFnData,
+  TQueryKey extends QueryKey = QueryKey,
+  TPageParam = unknown,
+> extends WithRequired<
+    Omit<
+      InfiniteQueryObserverOptions<
+        TQueryFnData,
+        TError,
+        TData,
+        TQueryData,
+        TQueryKey,
+        TPageParam
+      >,
+      'suspense'
+    >,
+    'queryKey'
+  > {}
+
+export interface AtomWithSuspenseInfiniteQueryOptions<
+  TQueryFnData = unknown,
+  TError = DefaultError,
+  TData = TQueryFnData,
+  TQueryData = TQueryFnData,
+  TQueryKey extends QueryKey = QueryKey,
+  TPageParam = unknown,
+> extends Omit<
+    AtomWithInfiniteQueryOptions<
+      TQueryFnData,
+      TError,
+      TData,
+      TQueryData,
+      TQueryKey,
+      TPageParam
+    >,
+    'enabled' | 'throwOnError' | 'placeholderData'
+  > {}
+
+export type AtomWithQueryResult<TData = unknown, TError = DefaultError> = Atom<
+  QueryObserverResult<TData, TError>
+>
+
+export type DefinedAtomWithQueryResult<
+  TData = unknown,
+  TError = DefaultError,
+> = Atom<DefinedQueryObserverResult<TData, TError>>
+
+export type AtomWithSuspenseQueryResult<
+  TData = unknown,
+  TError = DefaultError,
+> = Atom<
+  | Omit<DefinedQueryObserverResult<TData, TError>, 'isPlaceholderData'>
+  | Promise<
+      Omit<DefinedQueryObserverResult<TData, TError>, 'isPlaceholderData'>
+    >
+>
+
+export type AtomWithInfiniteQueryResult<
+  TData = unknown,
+  TError = DefaultError,
+> = Atom<InfiniteQueryObserverResult<TData, TError>>
+
+export type DefinedAtomWithInfiniteQueryResult<
+  TData = unknown,
+  TError = DefaultError,
+> = Atom<DefinedInfiniteQueryObserverResult<TData, TError>>
+
+export type AtomWithSuspenseInfiniteQueryResult<
+  TData = unknown,
+  TError = DefaultError,
+> = Atom<
+  Promise<
+    Omit<DefinedInfiniteQueryObserverResult<TData, TError>, 'isPlaceholderData'>
+  >
+>
+
+export type UndefinedInitialDataOptions<
+  TQueryFnData = unknown,
+  TError = DefaultError,
+  TData = TQueryFnData,
+  TQueryKey extends QueryKey = QueryKey,
+> = AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryKey> & {
+  initialData?: undefined
+}
+
+type NonUndefinedGuard<T> = T extends undefined ? never : T
+
+export type DefinedInitialDataOptions<
+  TQueryFnData = unknown,
+  TError = DefaultError,
+  TData = TQueryFnData,
+  TQueryKey extends QueryKey = QueryKey,
+> = AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryKey> & {
+  initialData:
+    | NonUndefinedGuard<TQueryFnData>
+    | (() => NonUndefinedGuard<TQueryFnData>)
+}
+
+export type UndefinedInitialDataInfiniteOptions<
+  TQueryFnData,
+  TError = DefaultError,
+  TData = InfiniteData<TQueryFnData>,
+  TQueryKey extends QueryKey = QueryKey,
+  TPageParam = unknown,
+> = AtomWithInfiniteQueryOptions<
+  TQueryFnData,
+  TError,
+  TData,
+  TQueryFnData,
+  TQueryKey,
+  TPageParam
+> & {
+  initialData?: undefined
+}
+
+export type DefinedInitialDataInfiniteOptions<
+  TQueryFnData,
+  TError = DefaultError,
+  TData = InfiniteData<TQueryFnData>,
+  TQueryKey extends QueryKey = QueryKey,
+  TPageParam = unknown,
+> = AtomWithInfiniteQueryOptions<
+  TQueryFnData,
+  TError,
+  TData,
+  TQueryFnData,
+  TQueryKey,
+  TPageParam
+> & {
+  initialData:
+    | NonUndefinedGuard<InfiniteData<TQueryFnData, TPageParam>>
+    | (() => NonUndefinedGuard<InfiniteData<TQueryFnData, TPageParam>>)
+}
diff --git a/src/utils.ts b/src/utils.ts
index 49bf13b..a42a721 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -1,4 +1,5 @@
 import type {
+  DefaultError,
   DefaultedQueryObserverOptions,
   Query,
   QueryKey,
@@ -54,3 +55,28 @@ export function shouldThrowError<T extends (...args: any[]) => boolean>(
 
   return !!throwOnError
 }
+
+export const defaultThrowOnError = <
+  TQueryFnData = unknown,
+  TError = DefaultError,
+  TData = TQueryFnData,
+  TQueryKey extends QueryKey = QueryKey,
+>(
+  _error: TError,
+  query: Query<TQueryFnData, TError, TData, TQueryKey>
+) => typeof query.state.data === 'undefined'
+
+export const ensureStaleTime = (
+  defaultedOptions: DefaultedQueryObserverOptions<any, any, any, any, any>
+) => {
+  if (defaultedOptions.suspense) {
+    if (typeof defaultedOptions.staleTime !== 'number') {
+      return {
+        ...defaultedOptions,
+        staleTime: 1000,
+      }
+    }
+  }
+
+  return defaultedOptions
+}