From 1929e157890cdf73d49d944f290106abaff066e3 Mon Sep 17 00:00:00 2001
From: Misha Bergal <misha.bergal@gmail.com>
Date: Thu, 10 Nov 2022 12:45:51 -0600
Subject: [PATCH 1/2] remember result promise for query

---
 packages/utils/src/useSubscription.ts | 59 +++++++++++++++++----------
 1 file changed, 37 insertions(+), 22 deletions(-)

diff --git a/packages/utils/src/useSubscription.ts b/packages/utils/src/useSubscription.ts
index ce8ef66..30aef00 100644
--- a/packages/utils/src/useSubscription.ts
+++ b/packages/utils/src/useSubscription.ts
@@ -30,16 +30,21 @@ import {
 type Unsubscribe = AuthUnsubscribe | FirestoreUnsubscribe | DatabaseUnsubscribe;
 
 const firestoreUnsubscribes: Record<string, any> = {};
-const queryCacheUnsubscribes: Record<string, () => void> = {};
+const queryCacheSubscribes: Record<
+  string,
+  { result: Promise<any>; unsubscribe: () => void }
+> = {};
 const eventCount: Record<string, number> = {};
 
 interface CancellablePromise<T = void> extends Promise<T> {
   cancel?: () => void;
 }
 
-type UseSubscriptionOptions<TData, TError, R> = UseQueryOptions<TData,
+type UseSubscriptionOptions<TData, TError, R> = UseQueryOptions<
+  TData,
   TError,
-  R> & {
+  R
+> & {
   onlyOnce?: boolean;
   fetchFn?: () => Promise<TData | null>;
 };
@@ -54,13 +59,14 @@ function firestoreUnsubscribe(subscriptionHash: string) {
 }
 
 function queryCacheUnsubscribe(subscriptionHash: string) {
-  const queryCacheUnsubscribe = queryCacheUnsubscribes[subscriptionHash];
+  const queryCacheUnsubscribe = queryCacheSubscribes[subscriptionHash];
   if (queryCacheUnsubscribe) {
-    queryCacheUnsubscribe();
-    delete queryCacheUnsubscribes[subscriptionHash];
+    queryCacheUnsubscribe.unsubscribe();
+    delete queryCacheSubscribes[subscriptionHash];
   }
 }
 
+let a = 1;
 /**
  * Utility hook to subscribe to events, given a function that returns an observer callback.
  * @param queryKey The react-query queryKey
@@ -79,17 +85,18 @@ export function useSubscription<TData, TError, R = TData>(
   const subscriptionHash = hashFn(subscriptionKey);
   const queryClient = useQueryClient();
 
-
   let resolvePromise: (data: TData | null) => void = () => null;
   let rejectPromise: (err: any) => void = () => null;
 
-  const result: CancellablePromise<TData | null> = new Promise<TData | null>(
+  let result: CancellablePromise<TData | null> = new Promise<TData | null>(
     (resolve, reject) => {
-      resolvePromise = resolve;
+      resolvePromise = (x) => {
+        resolve(x);
+      };
       rejectPromise = reject;
     }
   );
-
+  (result as any).a = a++;
   result.cancel = () => {
     queryClient.invalidateQueries(queryKey);
   };
@@ -109,10 +116,10 @@ export function useSubscription<TData, TError, R = TData>(
       }
     }
   } else {
-    const subscribedToQueryCache = !!queryCacheUnsubscribes[subscriptionHash];
+    const subscribedToQueryCache = !!queryCacheSubscribes[subscriptionHash];
     if (!subscribedToQueryCache) {
       const queryCache = queryClient.getQueryCache();
-      queryCacheUnsubscribes[subscriptionHash] = queryCache.subscribe((event) => {
+      const unsubscribe = queryCache.subscribe((event) => {
         if (!event || event.query.queryHash !== hashFn(queryKey)) {
           return;
         }
@@ -127,28 +134,36 @@ export function useSubscription<TData, TError, R = TData>(
           if (observersCount === 0) {
             firestoreUnsubscribe(subscriptionHash);
           } else {
-            const isSubscribedToFirestore = !!firestoreUnsubscribes[subscriptionHash];
+            const isSubscribedToFirestore =
+              !!firestoreUnsubscribes[subscriptionHash];
             if (isSubscribedToFirestore) {
-              const cachedData = queryClient.getQueryData<TData | null>(queryKey);
+              const cachedData = queryClient.getQueryData<TData | null>(
+                queryKey
+              );
               const hasData = !!eventCount[subscriptionHash];
 
               if (hasData) {
                 resolvePromise(cachedData ?? null);
               }
             } else {
-              firestoreUnsubscribes[subscriptionHash] = subscribeFn(async (data) => {
-                eventCount[subscriptionHash] ??= 0;
-                eventCount[subscriptionHash]++;
-                if (eventCount[subscriptionHash] === 1) {
-                  resolvePromise(data || null);
-                } else {
-                  queryClient.setQueryData(queryKey, data);
+              firestoreUnsubscribes[subscriptionHash] = subscribeFn(
+                async (data) => {
+                  eventCount[subscriptionHash] ??= 0;
+                  eventCount[subscriptionHash]++;
+                  if (eventCount[subscriptionHash] === 1) {
+                    resolvePromise(data || null);
+                  } else {
+                    queryClient.setQueryData(queryKey, data);
+                  }
                 }
-              });
+              );
             }
           }
         }
       });
+      queryCacheSubscribes[subscriptionHash] = { result, unsubscribe };
+    } else {
+      result = queryCacheSubscribes[subscriptionHash].result;
     }
   }
 

From f240e145c01b43d2f86e472c7ba20ed6025e9e44 Mon Sep 17 00:00:00 2001
From: Misha Bergal <misha.bergal@gmail.com>
Date: Thu, 10 Nov 2022 12:49:46 -0600
Subject: [PATCH 2/2] Cleanup

---
 packages/utils/src/useSubscription.ts | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/packages/utils/src/useSubscription.ts b/packages/utils/src/useSubscription.ts
index 30aef00..8518cb4 100644
--- a/packages/utils/src/useSubscription.ts
+++ b/packages/utils/src/useSubscription.ts
@@ -66,7 +66,6 @@ function queryCacheUnsubscribe(subscriptionHash: string) {
   }
 }
 
-let a = 1;
 /**
  * Utility hook to subscribe to events, given a function that returns an observer callback.
  * @param queryKey The react-query queryKey
@@ -90,13 +89,10 @@ export function useSubscription<TData, TError, R = TData>(
 
   let result: CancellablePromise<TData | null> = new Promise<TData | null>(
     (resolve, reject) => {
-      resolvePromise = (x) => {
-        resolve(x);
-      };
+      resolvePromise = resolve;
       rejectPromise = reject;
     }
   );
-  (result as any).a = a++;
   result.cancel = () => {
     queryClient.invalidateQueries(queryKey);
   };