fix(*): fix calling subscription fn multiple times when `subscription…

…Key` is not a string
kaciakmaciak · Oct 16, 2022 · a7f32f3 · a7f32f3
1 parent 7c1ef77
commit a7f32f3
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 29 deletions.
diff --git a/src/__tests__/subscription-storage.spec.ts b/src/__tests__/subscription-storage.spec.ts
@@ -38,14 +38,14 @@ describe('subscription storage', () => {
     const observerFn = jest.fn();
     const { subscription, next } = subscriptionFactory(observerFn);
 
-    storeSubscription(queryClient, ['test'], subscription);
+    storeSubscription(queryClient, 'test', subscription);
 
     next('value-1');
     expect(observerFn).toHaveBeenCalledTimes(1);
     expect(observerFn).toHaveBeenCalledWith('value-1');
     observerFn.mockClear();
 
-    cleanupSubscription(queryClient, ['test']);
+    cleanupSubscription(queryClient, 'test');
 
     next('value-2');
     expect(observerFn).not.toHaveBeenCalled();
@@ -59,8 +59,8 @@ describe('subscription storage', () => {
     const { subscription: subscriptionB, next: nextB } =
       subscriptionFactory(observerBFn);
 
-    storeSubscription(queryClient, ['testA'], subscriptionA);
-    storeSubscription(queryClient, ['testB'], subscriptionB);
+    storeSubscription(queryClient, 'testA', subscriptionA);
+    storeSubscription(queryClient, 'testB', subscriptionB);
 
     nextA('A1');
     expect(observerAFn).toHaveBeenCalledTimes(1);
@@ -71,7 +71,7 @@ describe('subscription storage', () => {
     expect(observerBFn).toHaveBeenCalledWith('B1');
     observerBFn.mockClear();
 
-    cleanupSubscription(queryClient, ['testA']);
+    cleanupSubscription(queryClient, 'testA');
 
     nextA('A2');
     expect(observerAFn).not.toHaveBeenCalled();
@@ -82,7 +82,7 @@ describe('subscription storage', () => {
 
   it('should not fail when key does not exist', () => {
     expect(() =>
-      cleanupSubscription(queryClient, ['test-non-existing'])
+      cleanupSubscription(queryClient, 'test-non-existing')
     ).not.toThrow();
   });
 });
diff --git a/src/subscription-storage.ts b/src/subscription-storage.ts
@@ -1,4 +1,4 @@
-import { QueryKey, hashQueryKey, QueryClient } from 'react-query';
+import type { QueryClient } from 'react-query';
 import { Subscription } from 'rxjs';
 
 const clientCacheSubscriptionsKey = ['__activeSubscriptions__'];
@@ -10,16 +10,15 @@ type SubscriptionStorage = Map<string, SubscriptionStorageItem>;
 /**
  * Stores subscription by its key and `pageParam` in the clientCache.
  */
-export function storeSubscription<TSubscriptionKey extends QueryKey = QueryKey>(
+export function storeSubscription(
   queryClient: QueryClient,
-  subscriptionKey: TSubscriptionKey,
+  hashedSubscriptionKey: string,
   subscription: Subscription,
   pageParam?: string
 ) {
   const activeSubscriptions: SubscriptionStorage =
     queryClient.getQueryData(clientCacheSubscriptionsKey) || new Map();
 
-  const hashedSubscriptionKey = hashQueryKey(subscriptionKey);
   const previousSubscription = activeSubscriptions.get(hashedSubscriptionKey);
 
   let newSubscriptionValue: SubscriptionStorageItem;
@@ -38,17 +37,14 @@ export function storeSubscription<TSubscriptionKey extends QueryKey = QueryKey>(
 /**
  * Removes stored subscription by its key and `pageParam` from the clientCache.
  */
-export function cleanupSubscription<
-  TSubscriptionKey extends QueryKey = QueryKey
->(
+export function cleanupSubscription(
   queryClient: QueryClient,
-  subscriptionKey: TSubscriptionKey,
+  hashedSubscriptionKey: string,
   pageParam?: string
 ) {
   const activeSubscriptions: SubscriptionStorage =
     queryClient.getQueryData(clientCacheSubscriptionsKey) || new Map();
 
-  const hashedSubscriptionKey = hashQueryKey(subscriptionKey);
   const subscription = activeSubscriptions.get(hashedSubscriptionKey);
 
   if (!subscription) return;

diff --git a/src/use-infinite-subscription.ts b/src/use-infinite-subscription.ts
@@ -1,5 +1,5 @@
 import { useEffect } from 'react';
-import { useInfiniteQuery, useQueryClient } from 'react-query';
+import { useInfiniteQuery, useQueryClient, hashQueryKey } from 'react-query';
 import type {
   QueryKey,
   UseInfiniteQueryResult,
@@ -161,6 +161,8 @@ export function useInfiniteSubscription<
     TSubscriptionKey
   > = {}
 ): UseInfiniteSubscriptionResult<TData, TError> {
+  const hashedSubscriptionKey = hashQueryKey(subscriptionKey);
+
   const { queryFn, clearErrors } = useObservableQueryFn(
     subscriptionFn,
     (data, previousData, pageParam): InfiniteData<TSubscriptionFnData> => {
@@ -225,10 +227,12 @@ export function useInfiniteSubscription<
         ?.getObserversCount();
 
       if (activeObserversCount === 0) {
-        cleanupSubscription(queryClient, subscriptionKey);
+        cleanupSubscription(queryClient, hashedSubscriptionKey);
       }
     };
-  }, [queryClient, subscriptionKey]);
+    // This is safe as `hashedSubscriptionKey` is derived from `subscriptionKey`.
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [queryClient, hashedSubscriptionKey]);
 
   return queryResult;
 }
diff --git a/src/use-observable-query-fn.ts b/src/use-observable-query-fn.ts
@@ -1,5 +1,5 @@
 import { useRef } from 'react';
-import { useQueryClient } from 'react-query';
+import { useQueryClient, hashQueryKey } from 'react-query';
 import type {
   QueryFunction,
   QueryKey,
@@ -43,7 +43,8 @@ export function useObservableQueryFn<
   const queryFn: QueryFunction<TSubscriptionFnData, TSubscriptionKey> = (
     context
   ) => {
-    const { queryKey, pageParam, signal } = context;
+    const { queryKey: subscriptionKey, pageParam, signal } = context;
+    const hashedSubscriptionKey = hashQueryKey(subscriptionKey);
 
     if (failRefetchWith.current) {
       throw failRefetchWith.current;
@@ -58,7 +59,7 @@ export function useObservableQueryFn<
     // If we do not invalidate the query, the hook will never re-subscribe,
     // as data are otherwise marked as fresh.
     function cancel() {
-      queryClient.invalidateQueries(queryKey, undefined, {
+      queryClient.invalidateQueries(subscriptionKey, undefined, {
         cancelRefetch: false,
       });
     }
@@ -75,27 +76,31 @@ export function useObservableQueryFn<
     }
 
     // @todo: Skip subscription for SSR
-    cleanupSubscription(queryClient, queryKey, pageParam ?? undefined);
+    cleanupSubscription(
+      queryClient,
+      hashedSubscriptionKey,
+      pageParam ?? undefined
+    );
 
     const subscription = stream$
       .pipe(
         skip(1),
         tap((data) => {
-          queryClient.setQueryData(queryKey, (previousData) =>
+          queryClient.setQueryData(subscriptionKey, (previousData) =>
             dataUpdater(data, previousData, pageParam)
           );
         }),
         catchError((error) => {
           failRefetchWith.current = error;
-          queryClient.setQueryData(queryKey, (data) => data, {
+          queryClient.setQueryData(subscriptionKey, (data) => data, {
             // To make the retryOnMount work
             // @see: https://github.com/tannerlinsley/react-query/blob/9e414e8b4f3118b571cf83121881804c0b58a814/src/core/queryObserver.ts#L727
             updatedAt: 0,
           });
           return of(undefined);
         }),
         finalize(() => {
-          queryClient.invalidateQueries(queryKey, undefined, {
+          queryClient.invalidateQueries(subscriptionKey, undefined, {
             cancelRefetch: false,
           });
         })
@@ -106,7 +111,7 @@ export function useObservableQueryFn<
     // see `cleanup` fn for more info
     storeSubscription(
       queryClient,
-      queryKey,
+      hashedSubscriptionKey,
       subscription,
       pageParam ?? undefined
     );

diff --git a/src/use-subscription.ts b/src/use-subscription.ts
@@ -1,5 +1,5 @@
 import { useEffect } from 'react';
-import { useQuery, useQueryClient } from 'react-query';
+import { useQuery, useQueryClient, hashQueryKey } from 'react-query';
 import type {
   QueryKey,
   UseQueryResult,
@@ -115,6 +115,8 @@ export function useSubscription<
     TSubscriptionKey
   > = {}
 ): UseSubscriptionResult<TData, TError> {
+  const hashedSubscriptionKey = hashQueryKey(subscriptionKey);
+
   const { queryFn, clearErrors } = useObservableQueryFn(
     subscriptionFn,
     (data) => data
@@ -161,10 +163,12 @@ export function useSubscription<
         ?.getObserversCount();
 
       if (activeObserversCount === 0) {
-        cleanupSubscription(queryClient, subscriptionKey);
+        cleanupSubscription(queryClient, hashedSubscriptionKey);
       }
     };
-  }, [queryClient, subscriptionKey]);
+    // This is safe as `hashedSubscriptionKey` is derived from `subscriptionKey`.
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [queryClient, hashedSubscriptionKey]);
 
   return queryResult;
 }