arbitrary change

__tests__/atomWithQuery_spec.tsx

@@ -1,4 +1,3 @@
-/* eslint-disable @typescript-eslint/no-unused-vars */
 import React, {
   Component,
   ReactNode,
@@ -16,7 +15,6 @@ beforeEach(() => {
 afterEach(() => {
   jest.runAllTimers()
   jest.useRealTimers()
-  jest.clearAllMocks()
 })
 
 it('query basic test', async () => {

__tests__/atomWithSuspenseQuery_spec.tsx

@@ -0,0 +1,198 @@
+import React, {
+  Component,
+  ReactNode,
+  StrictMode,
+  Suspense,
+  useState,
+} from 'react'
+import { QueryClient } from '@tanstack/query-core'
+import { fireEvent, render } from '@testing-library/react'
+import { Getter, atom, useAtom, useSetAtom } from 'jotai'
+import {
+  QueryErrorResetBoundary,
+  atomWithQuery,
+  atomWithSuspenseQuery,
+} from '../src'
+beforeEach(() => {
+  jest.useFakeTimers()
+})
+afterEach(() => {
+  jest.runAllTimers()
+  jest.useRealTimers()
+})
+
+// it('suspense basic, suspends', async () => {
+//   let resolve = () => {}
+//   const countAtom = atomWithSuspenseQuery(() => ({
+//     queryKey: ['test1'],
+//     queryFn: async () => {
+//       await new Promise<void>((r) => (resolve = r))
+//       return { response: { count: 0 } }
+//     },
+//   }))
+//   const Counter = () => {
+//     const [countData] = useAtom(countAtom)
+//     console.log(JSON.stringify({ countData }, null, 2))
+//     const { data } = countData
+//     return (
+//       <>
+//         <div>count: {data?.response.count}</div>
+//       </>
+//     )
+//   }
+
+//   const { findByText } = render(
+//     <StrictMode>
+//       <Suspense fallback="loading">
+//         <Counter />
+//       </Suspense>
+//     </StrictMode>
+//   )
+
+//   await findByText('loading')
+//   resolve()
+//   await findByText('count: 0')
+// })
+
+// it('query refetch', async () => {
+//   const mockFetch = jest.fn<
+//     { response: { message: string } },
+//     { message: string }[]
+//   >((response) => ({
+//     response,
+//   }))
+//   let resolve = () => {}
+//   const greetingAtom = atomWithSuspenseQuery(() => ({
+//     queryKey: ['test3'],
+//     queryFn: async () => {
+//       await new Promise<void>((r) => (resolve = r))
+//       const response = mockFetch({ message: 'helloWorld' })
+//       return response
+//     },
+//   }))
+//   const Greeting = () => {
+//     const [{ data, refetch }] = useAtom(greetingAtom)
+
+//     return (
+//       <>
+//         <div>message: {data?.response.message}</div>
+//         <button onClick={() => refetch?.()}>refetch</button>
+//       </>
+//     )
+//   }
+
+//   const { findByText, getByText } = render(
+//     <StrictMode>
+//       <Suspense fallback="loading">
+//         <Greeting />
+//       </Suspense>
+//     </StrictMode>
+//   )
+
+//   await findByText('loading')
+//   resolve()
+//   await findByText('message: helloWorld')
+//   expect(mockFetch).toBeCalledTimes(1)
+
+//   fireEvent.click(getByText('refetch'))
+//   await expect(() => findByText('loading')).rejects.toThrow() //refetch implementation in tanstack doesn't trigger loading state
+//   resolve()
+//   await findByText('message: helloWorld')
+//   expect(mockFetch).toBeCalledTimes(2) //this ensures we are actually running the query function again
+// })
+
+describe('intialData test', () => {
+  // it('query with initialData test', async () => {
+  //   const mockFetch = jest.fn((response) => ({ response }))
+  //   let resolve = () => {}
+
+  //   const countAtom = atomWithSuspenseQuery(() => ({
+  //     queryKey: ['initialData_count1'],
+  //     queryFn: async () => {
+  //       await new Promise<void>((r) => (resolve = r))
+  //       return mockFetch({ count: 10 })
+  //     },
+  //     initialData: { response: { count: 0 } },
+  //     staleTime: 0,
+  //   }))
+  //   const Counter = () => {
+  //     const [countData] = useAtom(countAtom)
+  //     const { data, isError } = countData
+
+  //     if (isError) {
+  //       return <>errorred</>
+  //     }
+
+  //     const count = data?.response.count
+  //     return (
+  //       <>
+  //         <div>count: {count}</div>
+  //       </>
+  //     )
+  //   }
+
+  //   const { findByText } = render(
+  //     <StrictMode>
+  //       <Suspense fallback="loading">
+  //         <Counter />
+  //       </Suspense>
+  //     </StrictMode>
+  //   )
+
+  //   // NOTE: the atom is never loading
+  //   await expect(() => findByText('loading')).rejects.toThrow()
+  //   await findByText('count: 0')
+  //   resolve()
+  //   await findByText('count: 10')
+  //   expect(mockFetch).toHaveBeenCalledTimes(1)
+  // })
+
+  it('query with initialData test with dependency', async () => {
+    const mockFetch = jest.fn((response) => ({ response }))
+    let resolve = () => {}
+    const numberAtom = atom(10)
+    const countAtom = atomWithSuspenseQuery((get) => ({
+      queryKey: ['initialData_count1', get(numberAtom)],
+      queryFn: async ({ queryKey: [, myNumber] }) => {
+        await new Promise<void>((r) => (resolve = r))
+        return mockFetch({ count: myNumber })
+      },
+      initialData: { response: { count: 0 } },
+      staleTime: 0,
+    }))
+    const Counter = () => {
+      const [countData] = useAtom(countAtom)
+      const { data, isError } = countData
+      if (isError) {
+        return <>errorred</>
+      }
+      const count = data?.response.count
+      return (
+        <>
+          <div>count: {count}</div>
+        </>
+      )
+    }
+
+    const Increment = () => {
+      const setNumber = useSetAtom(numberAtom)
+      return <button onClick={() => setNumber((n) => n + 1)}>increment</button>
+    }
+    const { findByText } = render(
+      <StrictMode>
+        <Suspense fallback="loading">
+          <Counter />
+        </Suspense>
+        <Increment />
+      </StrictMode>
+    )
+    // NOTE: the atom is never loading
+    await expect(() => findByText('loading')).rejects.toThrow()
+    await findByText('count: 0')
+    resolve()
+    await findByText('count: 10')
+    expect(mockFetch).toHaveBeenCalledTimes(1)
+    fireEvent.click(await findByText('increment'))
+    await findByText('count: 0')
+  })
+})

package.json

@@ -33,7 +33,7 @@
     "postcompile": "cp dist/index.modern.mjs dist/index.modern.js && cp dist/index.modern.mjs.map dist/index.modern.js.map",
     "test": "run-s eslint tsc-test jest",
     "eslint": "eslint --ext .js,.ts,.tsx .",
-    "jest": "jest",
+    "jest": "jest __tests__/atomWithSuspenseQuery_spec.tsx --watch",
     "tsc-test": "tsc --project . --noEmit",
     "examples:01_typescript": "DIR=01_typescript EXT=tsx webpack serve",
     "examples:02_refetch": "DIR=02_refetch EXT=tsx webpack serve",

src/atomWithSuspenseQuery.ts

@@ -3,17 +3,22 @@ import {
   type QueryKey,
   QueryObserver,
   type QueryObserverOptions,
+  type QueryObserverResult,
+  type QueryObserverSuccessResult,
 } from '@tanstack/query-core'
-import { Getter, atom } from 'jotai'
+import { Atom, Getter, atom } from 'jotai'
+import { filter, fromPromise, make, pipe, toObservable, toPromise } from 'wonka'
+import { isResetAtom } from './QueryAtomErrorResetBoundary'
 import { queryClientAtom } from './queryClientAtom'
+import { shouldSuspend } from './utils'
 
 export const atomWithSuspenseQuery = <
   TQueryFnData = unknown,
   TError = unknown,
   TData = TQueryFnData,
   TQueryData = TQueryFnData,
   TQueryKey extends QueryKey = QueryKey,
-  TSuspense extends boolean = false,
+  TInitialData = TData,
 >(
   getOptions: (
     get: Getter
@@ -23,18 +28,125 @@ export const atomWithSuspenseQuery = <
     TData,
     TQueryData,
     TQueryKey
-  > & { suspense?: TSuspense },
+  > & { initialData?: TInitialData },
   getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
 ) => {
-  return atom((get) => {
-    const options = { ...getOptions(get), enabled: true, suspense: true }
-    const queryClient = getQueryClient(get)
-    const defaultedQueryOptions = queryClient.defaultQueryOptions(options)
-    const observer = new QueryObserver(queryClient, defaultedQueryOptions)
-    // const errorResetBoundary = get(errorResetBoundaryAtom)
-
-    return observer.fetchOptimistic(defaultedQueryOptions).catch(() => {
-      //reset error boundary state
+  const IN_RENDER = Symbol()
+
+  const queryClientAtom = atom(getQueryClient)
+  const optionsAtom = atom((get) => {
+    const client = get(queryClientAtom)
+    const options = getOptions(get)
+    return client.defaultQueryOptions({
+      ...options,
+      enabled: true,
+      suspense: true,
+    })
+  })
+
+  const observerCacheAtom = atom(
+    () =>
+      new WeakMap<
+        QueryClient,
+        QueryObserver<TQueryFnData, TError, TData, TQueryData, TQueryKey>
+      >()
+  )
+
+  const observerAtom = atom((get) => {
+    const isReset = get(isResetAtom)
+
+    const options = get(optionsAtom)
+    const client = get(queryClientAtom)
+    const observerCache = get(observerCacheAtom)
+
+    const observer = observerCache.get(client)
+
+    if (isReset) {
+      if (observer) {
+        observerCache.delete(client)
+        observer.remove()
+      }
+      const newObserver = new QueryObserver(client, options)
+      observerCache.set(client, newObserver)
+      return newObserver
+    }
+
+    if (observer) {
+      ;(observer as any)[IN_RENDER] = true
+      observer.setOptions(options)
+      delete (observer as any)[IN_RENDER]
+
+      return observer
+    }
+
+    const newObserver = new QueryObserver(client, options)
+    observerCache.set(client, newObserver)
+
+    return newObserver
+  })
+
+  const observableAtom = atom((get) => {
+    const observer = get(observerAtom)
+    const source = make<QueryObserverResult<TData, TError>>(({ next }) => {
+      const callback = (result: QueryObserverResult<TData, TError>) => {
+        const notifyResult = () => next(result)
+
+        if ((observer as any)[IN_RENDER]) {
+          Promise.resolve().then(notifyResult)
+        } else {
+          notifyResult()
+        }
+      }
+
+      const unsubscribe = observer.subscribe(callback)
+      return () => unsubscribe()
     })
+    return pipe(
+      source,
+      filter((state) => !state.isFetching),
+      toObservable
+    )
+  })
+
+  const dataAtom = atom((get) => {
+    const observer = get(observerAtom)
+    const observable = get(observableAtom)
+
+    const currentResult = observer.getCurrentResult()
+    const resultAtom = atom(currentResult)
+
+    resultAtom.onMount = (set) => {
+      const { unsubscribe } = observable.subscribe((state) => {
+        set(state)
+      })
+      return () => unsubscribe()
+    }
+
+    return resultAtom
+  })
+
+  return atom((get) => {
+    const options = get(optionsAtom)
+    const observer = get(observerAtom)
+    const optimisticResult = observer.getOptimisticResult(options)
+    console.log({ defaultedOptions: options })
+    const suspend = shouldSuspend(options, optimisticResult, false)
+    console.log({ suspend })
+
+    if (suspend) {
+      return observer.fetchOptimistic(options)
+    }
+
+    const shouldThrowError =
+      optimisticResult.isError && !optimisticResult.isFetching
+
+    if (shouldThrowError) {
+      throw optimisticResult.error
+    }
+
+    const resultAtom = get(dataAtom)
+    const result = get(resultAtom)
+
+    return result
   })
 }

src/index.ts

@@ -4,4 +4,5 @@ export { atomsWithInfiniteQuery } from './atomsWithInfiniteQuery'
 export { atomsWithMutation } from './atomsWithMutation'
 export { atomsWithQueryAsync } from './atomsWithQueryAsync'
 export { atomWithQuery } from './atomWithQuery'
+export { atomWithSuspenseQuery } from './atomWithSuspenseQuery'
 export * from './QueryAtomErrorResetBoundary'