From 8ec8aafb8b6f07bf2bae8e864b47d5c10d25e123 Mon Sep 17 00:00:00 2001
From: kalijonn <43421621+kalijonn@users.noreply.github.com>
Date: Sun, 18 Aug 2024 08:24:22 +0530
Subject: [PATCH] add notify manager, cleanup unused code

---
 src/atomWithMutation.ts  |  7 ++-----
 src/baseAtomWithQuery.ts | 39 +++++++++++++--------------------------
 2 files changed, 15 insertions(+), 31 deletions(-)

diff --git a/src/atomWithMutation.ts b/src/atomWithMutation.ts
index 722ec2f..db91c7e 100644
--- a/src/atomWithMutation.ts
+++ b/src/atomWithMutation.ts
@@ -2,6 +2,7 @@ import {
   MutationObserver,
   MutationOptions,
   QueryClient,
+  notifyManager,
 } from '@tanstack/query-core'
 import { Atom, Getter, atom } from 'jotai'
 import { queryClientAtom } from './queryClientAtom'
@@ -67,17 +68,13 @@ export function atomWithMutation<
 
   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 = observer.subscribe((state) => {
-        set(state)
-      })
+      observer.subscribe(notifyManager.batchCalls(set))
       return () => {
-        unsubscribe
         observer.reset()
       }
     }
diff --git a/src/baseAtomWithQuery.ts b/src/baseAtomWithQuery.ts
index 112ba9b..a6d254e 100644
--- a/src/baseAtomWithQuery.ts
+++ b/src/baseAtomWithQuery.ts
@@ -3,6 +3,7 @@ import {
   QueryKey,
   QueryObserver,
   QueryObserverResult,
+  notifyManager,
 } from '@tanstack/query-core'
 import { Atom, Getter, atom } from 'jotai'
 import { queryClientAtom } from './queryClientAtom'
@@ -31,11 +32,6 @@ export function baseAtomWithQuery<
   | QueryObserverResult<TData, TError>
   | Promise<QueryObserverResult<TData, TError>>
 > {
-  const resetAtom = atom(0)
-  if (process.env.NODE_ENV !== 'production') {
-    resetAtom.debugPrivate = true
-  }
-
   const clientAtom = atom(getQueryClient)
   if (process.env.NODE_ENV !== 'production') {
     clientAtom.debugPrivate = true
@@ -94,20 +90,24 @@ export function baseAtomWithQuery<
   }
 
   const dataAtom = atom((get) => {
+    const client = getQueryClient(get)
     const observer = get(observerAtom)
+    const defaultedOptions = get(defaultedOptionsAtom)
+    const result = observer.getOptimisticResult(defaultedOptions)
 
-    const currentResult = observer.getCurrentResult()
-
-    const resultAtom = atom(currentResult)
+    const resultAtom = atom(result)
     if (process.env.NODE_ENV !== 'production') {
       resultAtom.debugPrivate = true
     }
 
     resultAtom.onMount = (set) => {
-      const unsubscribe = observer.subscribe((state) => {
-        set(state)
-      })
-      return unsubscribe
+      const unsubscribe = observer.subscribe(notifyManager.batchCalls(set))
+      return () => {
+        if (observer.getCurrentResult().isError) {
+          client.resetQueries({ queryKey: observer.getCurrentQuery().queryKey })
+        }
+        unsubscribe()
+      }
     }
 
     return resultAtom
@@ -120,20 +120,7 @@ export function baseAtomWithQuery<
     const observer = get(observerAtom)
     const defaultedOptions = get(defaultedOptionsAtom)
 
-    const client = getQueryClient(get)
-
-    resetAtom.onMount = () => {
-      return () => {
-        if (observer.getCurrentResult().isError) {
-          client.resetQueries({ queryKey: observer.getCurrentQuery().queryKey })
-        }
-      }
-    }
-
-    get(resetAtom)
-    get(get(dataAtom))
-
-    const result = observer.getOptimisticResult(defaultedOptions)
+    const result = get(get(dataAtom))
 
     if (shouldSuspend(defaultedOptions, result, false)) {
       return observer.fetchOptimistic(defaultedOptions)