jotaijs · dai-shi · Apr 3, 2023 · Mar 22, 2023 · Mar 24, 2023 · Mar 24, 2023
diff --git a/__tests__/atomsWithQueryAsync_spec.tsx b/__tests__/atomsWithQueryAsync_spec.tsx
@@ -0,0 +1,51 @@
+import React, { StrictMode, Suspense } from 'react'
+import { render } from '@testing-library/react'
+import { atom, useAtom } from 'jotai'
+import { atomsWithQueryAsync } from '../src/index'
+
+it('async query basic test', async () => {
+  let resolve = () => {}
+  let resolveQuery = () => {}
+
+  const userIdAtom = atom(async () => {
+    await new Promise<void>((r) => (resolve = r))
+    return 2
+  })
+  const [userAtom] = atomsWithQueryAsync(async (get) => {
+    const userId = await get(userIdAtom)
+
+    return {
+      queryKey: ['idTest', userId],
+      queryFn: async ({ queryKey: [, id] }) => {
+        await new Promise<void>((r) => (resolveQuery = r))
+        return { response: { id: id as number } }
+      },
+    }
+  })
+  const User = () => {
+    const [
+      {
+        response: { id },
+      },
+    ] = useAtom(userAtom)
+
+    return (
+      <>
+        <div>id: {id}</div>
+      </>
+    )
+  }
+  const { findByText } = render(
+    <StrictMode>
+      <Suspense fallback="loading">
+        <User />
+      </Suspense>
+    </StrictMode>
+  )
+
+  await findByText('loading')
+  resolve()
+  await new Promise((r) => setTimeout(r)) // wait a tick
+  resolveQuery()
+  await findByText('id: 2')
+})
diff --git a/examples/05_async/package.json b/examples/05_async/package.json
@@ -0,0 +1,28 @@
+{
+  "name": "jotai-tanstack-query-example",
+  "version": "0.1.0",
+  "private": true,
+  "dependencies": {
+    "@tanstack/query-core": "latest",
+    "@types/react": "latest",
+    "@types/react-dom": "latest",
+    "jotai": "latest",
+    "jotai-tanstack-query": "latest",
+    "react": "latest",
+    "react-dom": "latest",
+    "react-scripts": "latest",
+    "typescript": "latest"
+  },
+  "scripts": {
+    "start": "react-scripts start",
+    "build": "react-scripts build",
+    "test": "react-scripts test",
+    "eject": "react-scripts eject"
+  },
+  "browserslist": [
+    ">0.2%",
+    "not dead",
+    "not ie <= 11",
+    "not op_mini all"
+  ]
+}
diff --git a/examples/05_async/public/index.html b/examples/05_async/public/index.html
@@ -0,0 +1,8 @@
+<html>
+  <head>
+    <title>jotai-tanstack-query example</title>
+  </head>
+  <body>
+    <div id="app"></div>
+  </body>
+</html>
diff --git a/examples/05_async/src/App.tsx b/examples/05_async/src/App.tsx
@@ -0,0 +1,32 @@
+import React from 'react'
+import { atom, useAtomValue } from 'jotai'
+import { atomsWithQueryAsync } from 'jotai-tanstack-query'
+
+const idAtom = atom(async () => {
+  await new Promise((resolve) => setTimeout(resolve, 2000))
+  return 2
+})
+
+const [userAtom] = atomsWithQueryAsync(async (get) => {
+  const id = await get(idAtom)
+  return {
+    queryKey: ['getUser', id],
+    queryFn: async () => {
+      const res = await fetch('https://reqres.in/api/users/' + id)
+      return res.json() as Promise<{ data: unknown }>
+    },
+  }
+})
+
+const UserData = () => {
+  const data = useAtomValue(userAtom)
+  return <div>{JSON.stringify(data)}</div>
+}
+
+const App = () => (
+  <React.Suspense fallback="Loading...">
+    <UserData />
+  </React.Suspense>
+)
+
+export default App
diff --git a/examples/05_async/src/index.tsx b/examples/05_async/src/index.tsx
@@ -0,0 +1,8 @@
+import React from 'react'
+import { createRoot } from 'react-dom/client'
+import App from './App'
+
+const ele = document.getElementById('app')
+if (ele) {
+  createRoot(ele).render(React.createElement(App))
+}
diff --git a/package.json b/package.json
@@ -35,7 +35,8 @@
     "examples:01_typescript": "DIR=01_typescript EXT=tsx webpack serve",
     "examples:02_refetch": "DIR=02_refetch EXT=tsx webpack serve",
     "examples:03_infinite": "DIR=03_infinite EXT=tsx webpack serve",
-    "examples:04_mutation": "DIR=04_mutation EXT=tsx webpack serve"
+    "examples:04_mutation": "DIR=04_mutation EXT=tsx webpack serve",
+    "examples:05_async": "DIR=05_async EXT=tsx webpack serve"
   },
   "jest": {
     "testEnvironment": "jsdom",

diff --git a/src/atomsWithQueryAsync.ts b/src/atomsWithQueryAsync.ts
@@ -0,0 +1,56 @@
+import type { QueryKey, QueryObserverOptions } from '@tanstack/query-core'
+import type { ExtractAtomValue, Getter, WritableAtom } from 'jotai'
-import type { ExtractAtomValue, Getter, WritableAtom } from 'jotai'
+import type { ExtractAtomArgs, ExtractAtomValue, Getter, WritableAtom } from 'jotai'
-import type { ExtractAtomValue, Getter, WritableAtom } from 'jotai'
+import type { ExtractAtomArgs, ExtractAtomValue, Getter, WritableAtom } from 'jotai'
+import { atom } from 'jotai'
+import { atomsWithQuery } from './atomsWithQuery'
+
+type ExtractAtomArgs<AtomType> = AtomType extends WritableAtom<
+  any,
+  infer Args,
+  any
+>
+  ? Args
+  : never
+
+export function atomsWithQueryAsync<
+  TQueryFnData = unknown,
+  TError = unknown,
+  TData = TQueryFnData,
+  TQueryData = TQueryFnData,
+  TQueryKey extends QueryKey = QueryKey
+>(
+  getOptions: (
+    get: Getter
+  ) => Promise<
+    QueryObserverOptions<TQueryFnData, TError, TData, TQueryData, TQueryKey>
+  >
+) {
+  const atomsAtom = atom(async (get) => {
+    const options = await getOptions(get)
+    return atomsWithQuery(() => options)
+  })
+
+  const dataAtom = atom(
+    async (get) => get((await get(atomsAtom))[0]),
+    async (
+      get,
+      set,
+      ...args: ExtractAtomArgs<Awaited<ExtractAtomValue<typeof atomsAtom>>[0]>
+    ) => {
+      const a = (await get(atomsAtom))[0]
+      return set(a, ...args)
+    }
+  )
+
+  const statusAtom = atom(
+    async (get) => get((await get(atomsAtom))[1]),
+    async (
+      get,
+      set,
+      ...args: ExtractAtomArgs<Awaited<ExtractAtomValue<typeof atomsAtom>>[1]>
+    ) => {
+      const a = (await get(atomsAtom))[1]
+      return set(a, ...args)
+    }
+  )
+  return [dataAtom, statusAtom] as const
+}
diff --git a/src/index.ts b/src/index.ts
@@ -2,3 +2,4 @@ export { queryClientAtom } from './queryClientAtom'
 export { atomsWithQuery } from './atomsWithQuery'
 export { atomsWithInfiniteQuery } from './atomsWithInfiniteQuery'
 export { atomsWithMutation } from './atomsWithMutation'
+export { atomsWithQueryAsync } from './atomsWithQueryAsync'