Defer evaluation of the variables if query is paused, add atomWithLaz…

…yQuery atom (#15)

* fix: defer valuation of the variables if paused
feat: add atomWithLazyQuery that allows for promisifying of a singular query result the same way `atomWithMutation` does
* feat: add fetching state for mutations and lazy query
fix: make sure that lazy query gets reset after all listeners are unmounted

src/atomWithLazyQuery.ts

@@ -0,0 +1,70 @@
+import type { AnyVariables } from '@urql/core'
+import { Client, DocumentInput, OperationContext } from '@urql/core'
+import { WritableAtom, atom } from 'jotai/vanilla'
+import type { Getter } from 'jotai/vanilla'
+import { filter, pipe, subscribe } from 'wonka'
+import { clientAtom } from './clientAtom'
+import {
+  InitialOperationResultLazy,
+  urqlReactCompatibleInitialStateLazy,
+} from './common'
+
+export type AtomWithLazyQuery<
+  Data,
+  Variables extends AnyVariables
+> = WritableAtom<
+  InitialOperationResultLazy<Data, Variables>,
+  [Variables, Partial<OperationContext>] | [Variables],
+  Promise<InitialOperationResultLazy<Data, Variables>>
+>
+
+export function atomWithLazyQuery<
+  Data = unknown,
+  Variables extends AnyVariables = AnyVariables
+>(
+  query: DocumentInput<Data, Variables>,
+  getClient: (get: Getter) => Client = (get) => get(clientAtom)
+): AtomWithLazyQuery<Data, Variables> {
+  const atomDataBase = atom<InitialOperationResultLazy<Data, Variables>>(
+    urqlReactCompatibleInitialStateLazy
+  )
+  atomDataBase.onMount = (setAtom) => {
+    return () => {
+      // Clean up the atom cache on unmount
+      setAtom(urqlReactCompatibleInitialStateLazy)
+    }
+  }
+  const atomData = atom<
+    InitialOperationResultLazy<Data, Variables>,
+    [Variables, Partial<OperationContext>] | [Variables],
+    Promise<InitialOperationResultLazy<Data, Variables>>
+  >(
+    (get) => {
+      return get(atomDataBase)
+    },
+    (get, set, ...args) => {
+      const source = getClient(get).query(query, args[0], {
+        requestPolicy: 'network-only',
+        ...(args[1] ? args[1] : {}),
+      })
+      pipe(
+        source,
+        // This is needed so that the atom gets updated with loading states etc., but not with the final result that will be set by the promise
+        filter((result) => result?.data === undefined),
+        subscribe((result) => set(atomDataBase, { ...result, fetching: true }))
+      )
+
+      set(atomDataBase, {
+        ...urqlReactCompatibleInitialStateLazy,
+        fetching: true,
+      })
+      return source.toPromise().then((result) => {
+        const mergedResult = { ...result, fetching: false }
+        set(atomDataBase, mergedResult)
+        return mergedResult
+      })
+    }
+  )
+
+  return atomData
+}

src/atomWithMutation.ts

@@ -2,20 +2,20 @@ import { DocumentInput } from '@urql/core'
 import type { AnyVariables, Client, OperationContext } from '@urql/core'
 import { atom } from 'jotai/vanilla'
 import type { Getter, WritableAtom } from 'jotai/vanilla'
-import { pipe, subscribe } from 'wonka'
+import { filter, pipe, subscribe } from 'wonka'
 import { clientAtom } from './clientAtom'
 import {
-  type InitialOperationResult,
-  urqlReactCompatibleInitialState,
+  type InitialOperationResultLazy,
+  urqlReactCompatibleInitialStateLazy,
 } from './common'
 
 export type AtomWithMutation<
   Data,
   Variables extends AnyVariables
 > = WritableAtom<
-  InitialOperationResult<Data, Variables>,
+  InitialOperationResultLazy<Data, Variables>,
   [Variables, Partial<OperationContext>] | [Variables],
-  Promise<InitialOperationResult<Data, Variables>>
+  Promise<InitialOperationResultLazy<Data, Variables>>
 >
 
 export function atomWithMutation<
@@ -25,19 +25,19 @@ export function atomWithMutation<
   query: DocumentInput<Data, Variables>,
   getClient: (get: Getter) => Client = (get) => get(clientAtom)
 ): AtomWithMutation<Data, Variables> {
-  const atomDataBase = atom<InitialOperationResult<Data, Variables>>(
-    urqlReactCompatibleInitialState
+  const atomDataBase = atom<InitialOperationResultLazy<Data, Variables>>(
+    urqlReactCompatibleInitialStateLazy
   )
   atomDataBase.onMount = (setAtom) => {
     return () => {
       // Clean up the atom cache on unmount
-      setAtom(urqlReactCompatibleInitialState)
+      setAtom(urqlReactCompatibleInitialStateLazy)
     }
   }
   const atomData = atom<
-    InitialOperationResult<Data, Variables>,
+    InitialOperationResultLazy<Data, Variables>,
     [Variables, Partial<OperationContext>] | [Variables],
-    Promise<InitialOperationResult<Data, Variables>>
+    Promise<InitialOperationResultLazy<Data, Variables>>
   >(
     (get) => {
       return get(atomDataBase)
@@ -46,10 +46,20 @@ export function atomWithMutation<
       const source = getClient(get).mutation(query, args[0], args[1])
       pipe(
         source,
-        subscribe((result) => set(atomDataBase, result))
+        // This is needed so that the atom gets updated with loading states etc., but not with the final result that will be set by the promise
+        filter((result) => result?.data === undefined),
+        subscribe((result) => set(atomDataBase, { ...result, fetching: true }))
       )
 
-      return source.toPromise()
+      set(atomDataBase, {
+        ...urqlReactCompatibleInitialStateLazy,
+        fetching: true,
+      })
+      return source.toPromise().then((result) => {
+        const mergedResult = { ...result, fetching: false }
+        set(atomDataBase, mergedResult)
+        return mergedResult
+      })
     }
   )
 

src/common.ts

@@ -18,6 +18,25 @@ export type InitialOperationResult<Data, Variables extends AnyVariables> = Omit<
 > & {
   operation: Operation<Data, Variables> | undefined
 }
+
+export type InitialOperationResultLazy<
+  Data,
+  Variables extends AnyVariables
+> = Omit<OperationResult<Data, Variables>, 'operation'> & {
+  operation: Operation<Data, Variables> | undefined
+  fetching: boolean
+}
+export const urqlReactCompatibleInitialStateLazy = {
+  stale: false,
+  // Casting is needed to make typescript chill here as it tries here to be too smart
+  error: undefined as any,
+  data: undefined as any,
+  extensions: undefined as any,
+  hasNext: false,
+  operation: undefined,
+  fetching: false,
+} as InitialOperationResultLazy<any, any>
+
 // This is the same (aside from missing fetching and having hasNext) object shape as urql-react has by default while operation is yet to be triggered/yet to be fetched
 export const urqlReactCompatibleInitialState = {
   stale: false,
@@ -41,9 +60,9 @@ export const createAtoms = <Args, Result extends OperationResult, ActionResult>(
   )
 
   const baseStatusAtom = atom((get) => {
-    const args = getArgs(get)
+    const isPaused = getPause(get)
     const client = getClient(get)
-    const source = getPause(get) ? null : execute(client, args)
+    const source = isPaused ? null : execute(client, getArgs(get))
     if (!source) {
       return initialLoadAtom
     }

src/index.ts

@@ -1,6 +1,8 @@
 export { atomWithMutation } from './atomWithMutation'
 export type { AtomWithMutation } from './atomWithMutation'
 export { atomWithQuery } from './atomWithQuery'
+export { atomWithLazyQuery } from './atomWithLazyQuery'
+export type { AtomWithLazyQuery } from './atomWithLazyQuery'
 export { atomWithSubscription } from './atomWithSubscription'
 export { clientAtom } from './clientAtom'
 export { suspenseAtom } from './suspenseAtom'

tests/playwright-tests/smoke.test.ts

@@ -201,3 +201,72 @@ test('smoke - mutation', async ({ page }) => {
   await expect(page.getByTestId('query-table').getByText('25')).toBeVisible()
   await expect(page.getByTestId('mutation-table')).not.toBeVisible()
 })
+
+test('smoke - lazy query', async ({ page }) => {
+  await mockResponse(page, {
+    burgers: [
+      {
+        __typename: 'Burger',
+        id: '1',
+        name: 'Big Tasty',
+        price: 8,
+      },
+      {
+        __typename: 'Burger',
+        id: '2',
+        name: 'Big Mac',
+        price: 5,
+      },
+      {
+        __typename: 'Burger',
+        id: '3',
+        name: 'McChicken',
+        price: 3,
+      },
+    ],
+  })
+  await page.goto('/mutations')
+  // Suspended on initial loading
+  await expect(page.getByText('loading')).toBeVisible()
+  // Shows data after loading
+  await expect(page.getByText('Big Mac')).toBeVisible()
+
+  await mockResponse(page, {
+    burgers: [
+      {
+        __typename: 'Burger',
+        id: '1',
+        name: 'Big Tasty',
+        price: 8,
+      },
+      {
+        __typename: 'Burger',
+        id: '2',
+        name: 'Big Mac',
+        price: 25, // We update the data here as if it changed
+      },
+      {
+        __typename: 'Burger',
+        id: '3',
+        name: 'McChicken',
+        price: 3,
+      },
+    ],
+  })
+  await page.getByText('load lazy burgers').click()
+  await expect(page.getByText('loading burgers...')).toBeVisible()
+  // New one off result should affect query result updating burger with an id 2
+  await expect(page.getByTestId('query-table').getByText('25')).toBeVisible()
+  // Same result should be visible within lazy query table
+  await expect(
+    page.getByTestId('query-lazy-table').getByText('25')
+  ).toBeVisible()
+
+  // Rerouting to an empty page to check if lazy query state will reset (same behaviour as urql official react bindings)
+  await page.getByText('Home').click()
+  await expect(page.getByText('Empty')).toBeVisible()
+  // Rerouting back. Query results should stay the same (with mutation affecting them). Mutation results should be reset.
+  await page.getByText('Mutations').click()
+  await expect(page.getByTestId('query-table').getByText('25')).toBeVisible()
+  await expect(page.getByTestId('query-lazy-table')).not.toBeVisible()
+})

tests/test-app/src/CacheAndMutations.tsx

@@ -1,6 +1,10 @@
 import React, { Suspense } from 'react'
 import { useAtom } from 'jotai'
-import { atomWithMutation, atomWithQuery } from '../../../src'
+import {
+  atomWithLazyQuery,
+  atomWithMutation,
+  atomWithQuery,
+} from '../../../src'
 import { generateUrqlClient } from './client'
 
 const client = generateUrqlClient()
@@ -29,12 +33,24 @@ const burgerCreateAtom = atomWithMutation<{ burgerCreate: Burger }>(
   () => client
 )
 
+const burgerLazyBurgersAtom = atomWithLazyQuery<{ burgers: Burger[] }>(
+  `query Index_Burgers {
+  burgers {
+    id
+    name
+    price
+  }
+}`,
+  () => client
+)
+
 const Burgers = () => {
   const [opResult] = useAtom(burgersAtom)
   const [mutationOpResult, mutate] = useAtom(burgerCreateAtom)
+  const [burgersLazyOpResult, lazyLoading] = useAtom(burgerLazyBurgersAtom)
   const burgers = opResult?.data?.burgers
   const burger = mutationOpResult.data?.burgerCreate
-
+  const lazyBurgers = burgersLazyOpResult?.data?.burgers
   if (!burgers) throw new Error('No burgers loaded!')
 
   return (
@@ -50,6 +66,22 @@ const Burgers = () => {
           ))}
         </tbody>
       </table>
+      <h2>Lazy burgers</h2>
+      <button onClick={() => lazyLoading({})}>load lazy burgers</button>
+      {burgersLazyOpResult.fetching && <div>loading burgers...</div>}
+      {lazyBurgers && (
+        <table data-testid="query-lazy-table">
+          <tbody>
+            {lazyBurgers?.map((burger) => (
+              <tr key={burger.id}>
+                <td>{burger.id}</td>
+                <td>{burger.name}</td>
+                <td>{burger.price}</td>
+              </tr>
+            ))}
+          </tbody>
+        </table>
+      )}
       <button onClick={() => mutate({})}>mutate</button>
       {burger && (
         <table data-testid="mutation-table">