@trpc/react-query 11.0.0-alpha-tmp-subscription-connection-state.488 → 11.0.0-alpha-tmp-12-06-react.665

package/src/shared/hooks/createHooksInternal.tsx

@@ -1,14 +1,15 @@
 import {
   useInfiniteQuery as __useInfiniteQuery,
   useMutation as __useMutation,
+  usePrefetchInfiniteQuery as __usePrefetchInfiniteQuery,
   useQueries as __useQueries,
   useQuery as __useQuery,
   useSuspenseInfiniteQuery as __useSuspenseInfiniteQuery,
   useSuspenseQueries as __useSuspenseQueries,
   useSuspenseQuery as __useSuspenseQuery,
+  usePrefetchQuery as _usePrefetchQuery,
   hashKey,
   skipToken,
-  useQueryClient,
 } from '@tanstack/react-query';
 import type { TRPCClientErrorLike } from '@trpc/client';
 import { createTRPCUntypedClient } from '@trpc/client';
@@ -23,7 +24,10 @@ import {
   getMutationKeyInternal,
   getQueryKeyInternal,
 } from '../../internals/getQueryKey';
-import { useHookResult } from '../../internals/useHookResult';
+import {
+  buildQueryFromAsyncIterable,
+  useHookResult,
+} from '../../internals/trpcResult';
 import type {
   TRPCUseQueries,
   TRPCUseSuspenseQueries,
@@ -31,33 +35,27 @@ import type {
 import { createUtilityFunctions } from '../../utils/createUtilityFunctions';
 import { createUseQueries } from '../proxy/useQueriesProxy';
 import type { CreateTRPCReactOptions, UseMutationOverride } from '../types';
-// import type { restartSubscriptionFn } from './types';
-import {
-  getConnectingResult,
-  getErrorResult,
-  getIdleResult,
-  getPendingResult,
-  getStartingResult,
-  type CreateClient,
-  type TRPCProvider,
-  type TRPCQueryOptions,
-  type UseTRPCInfiniteQueryOptions,
-  type UseTRPCInfiniteQueryResult,
-  type UseTRPCMutationOptions,
-  type UseTRPCMutationResult,
-  type UseTRPCQueryOptions,
-  type UseTRPCQueryResult,
-  type UseTRPCSubscriptionOptions,
-  type UseTRPCSubscriptionResult,
-  type UseTRPCSuspenseInfiniteQueryOptions,
-  type UseTRPCSuspenseInfiniteQueryResult,
-  type UseTRPCSuspenseQueryOptions,
-  type UseTRPCSuspenseQueryResult,
+import type {
+  CreateClient,
+  TRPCProvider,
+  TRPCQueryOptions,
+  TRPCSubscriptionResult,
+  UseTRPCInfiniteQueryOptions,
+  UseTRPCInfiniteQueryResult,
+  UseTRPCMutationOptions,
+  UseTRPCMutationResult,
+  UseTRPCPrefetchInfiniteQueryOptions,
+  UseTRPCPrefetchQueryOptions,
+  UseTRPCQueryOptions,
+  UseTRPCQueryResult,
+  UseTRPCSubscriptionOptions,
+  UseTRPCSuspenseInfiniteQueryOptions,
+  UseTRPCSuspenseInfiniteQueryResult,
+  UseTRPCSuspenseQueryOptions,
+  UseTRPCSuspenseQueryResult,
 } from './types';
 
-const trackResult = <
-  T extends UseTRPCSubscriptionResult<unknown, unknown, unknown>,
->(
+const trackResult = <T extends object>(
   result: T,
   onTrackResult: (key: keyof T) => void,
 ): T => {
@@ -143,7 +141,7 @@ export function createRootHooks<
 
   /**
    * Hack to make sure errors return `status`='error` when doing SSR
-   * @link https://github.com/trpc/trpc/pull/1645
+   * @see https://github.com/trpc/trpc/pull/1645
    */
   function useSSRQueryOptionsIfNeeded<
     TOptions extends { retryOnMount?: boolean } | undefined,
@@ -213,26 +211,11 @@ export function createRootHooks<
               );
 
               if (isAsyncIterable(result)) {
-                const queryCache = queryClient.getQueryCache();
-
-                const query = queryCache.build(queryFunctionContext.queryKey, {
+                return buildQueryFromAsyncIterable(
+                  result,
+                  queryClient,
                   queryKey,
-                });
-
-                query.setState({
-                  data: [],
-                  status: 'success',
-                });
-
-                const aggregate: unknown[] = [];
-                for await (const value of result) {
-                  aggregate.push(value);
-
-                  query.setState({
-                    data: [...aggregate],
-                  });
-                }
-                return aggregate;
+                );
               }
               return result;
             },
@@ -247,6 +230,41 @@ export function createRootHooks<
     return hook;
   }
 
+  function usePrefetchQuery(
+    path: string[],
+    input: unknown,
+    opts?: UseTRPCPrefetchQueryOptions<unknown, unknown, TError>,
+  ): void {
+    const context = useContext();
+    const queryKey = getQueryKeyInternal(path, input, 'query');
+
+    const isInputSkipToken = input === skipToken;
+
+    const shouldAbortOnUnmount =
+      opts?.trpc?.abortOnUnmount ??
+      config?.abortOnUnmount ??
+      context.abortOnUnmount;
+
+    _usePrefetchQuery({
+      ...opts,
+      queryKey: queryKey as any,
+      queryFn: isInputSkipToken
+        ? input
+        : (queryFunctionContext) => {
+            const actualOpts = {
+              trpc: {
+                ...opts?.trpc,
+                ...(shouldAbortOnUnmount
+                  ? { signal: queryFunctionContext.signal }
+                  : {}),
+              },
+            };
+
+            return context.client.query(...getClientArgs(queryKey, actualOpts));
+          },
+    });
+  }
+
   function useSuspenseQuery(
     path: readonly string[],
     input: unknown,
@@ -266,7 +284,9 @@ export function createRootHooks<
         queryKey: queryKey as any,
         queryFn: (queryFunctionContext) => {
           const actualOpts = {
+            ...opts,
             trpc: {
+              ...opts?.trpc,
               ...(shouldAbortOnUnmount
                 ? { signal: queryFunctionContext.signal }
                 : { signal: null }),
@@ -290,8 +310,7 @@ export function createRootHooks<
     path: readonly string[],
     opts?: UseTRPCMutationOptions<unknown, TError, unknown, unknown>,
   ): UseTRPCMutationResult<unknown, TError, unknown, unknown> {
-    const { client } = useContext();
-    const queryClient = useQueryClient();
+    const { client, queryClient } = useContext();
 
     const mutationKey = getMutationKeyInternal(path);
 
@@ -332,80 +351,35 @@ export function createRootHooks<
     path: readonly string[],
     input: unknown,
     opts: UseTRPCSubscriptionOptions<unknown, TError>,
-  ): UseTRPCSubscriptionResult<unknown, unknown, TError> {
+  ) {
     const enabled = opts?.enabled ?? input !== skipToken;
     const queryKey = hashKey(getQueryKeyInternal(path, input, 'any'));
+    const { client } = useContext();
 
-    const trackedProps = React.useRef(
-      new Set<keyof UseTRPCSubscriptionResult<unknown, unknown, TError>>([]),
-    );
-
-    const addTrackedProp = React.useCallback(
-      (key: keyof UseTRPCSubscriptionResult<unknown, unknown, TError>) => {
-        trackedProps.current.add(key);
-      },
-      [],
-    );
-
-    // const restart = React.useRef<restartSubscriptionFn<unknown>>(() => {
-    //   throw new Error('not implemented');
-    // });
-
-    const currentResult = React.useRef<
-      UseTRPCSubscriptionResult<unknown, unknown, TError>
-    >(
-      enabled
-        ? getStartingResult(/* restart.current */)
-        : getIdleResult(/* restart.current */),
-    );
-
-    const [subscriptionState, setSubscriptionState] = React.useState(
-      trackResult(currentResult.current, addTrackedProp),
-    );
-
-    const updateSubscriptionState = React.useCallback(
-      (
-        opts:
-          | UseTRPCSubscriptionResult<unknown, unknown, TError>
-          | ((
-              prev: UseTRPCSubscriptionResult<unknown, unknown, TError>,
-            ) => UseTRPCSubscriptionResult<unknown, unknown, TError>),
-      ) => {
-        const oldResult = currentResult.current;
-
-        const newResult =
-          typeof opts === 'function' ? opts(currentResult.current) : opts;
+    const optsRef = React.useRef<typeof opts>(opts);
+    optsRef.current = opts;
 
-        currentResult.current = newResult;
+    type $Result = TRPCSubscriptionResult<unknown, TError>;
 
-        let shouldUpdate = false;
+    const trackedProps = React.useRef(new Set<keyof $Result>([]));
 
-        for (const key of trackedProps.current) {
-          if (oldResult[key] !== newResult[key]) {
-            shouldUpdate = true;
-            break;
-          }
-        }
-
-        if (shouldUpdate) {
-          setSubscriptionState(trackResult(newResult, addTrackedProp));
-          return;
-        }
-      },
-      [addTrackedProp],
-    );
+    const addTrackedProp = React.useCallback((key: keyof $Result) => {
+      trackedProps.current.add(key);
+    }, []);
 
-    const { client } = useContext();
+    type Unsubscribe = () => void;
+    const currentSubscriptionRef = React.useRef<Unsubscribe | null>(null);
 
-    const optsRef = React.useRef<typeof opts>(opts);
-    optsRef.current = opts;
+    const reset = React.useCallback((): void => {
+      // unsubscribe from the previous subscription
+      currentSubscriptionRef.current?.();
 
-    React.useEffect(() => {
+      updateState(getInitialState);
       if (!enabled) {
         return;
       }
-      let isStopped = false;
 
+      let isStopped = false;
       const subscription = client.subscription(
         path.join('.'),
         input ?? undefined,
@@ -413,67 +387,113 @@ export function createRootHooks<
           onStarted: () => {
             if (!isStopped) {
               optsRef.current.onStarted?.();
+              updateState((prev) => ({
+                ...prev,
+                status: 'pending',
+                error: null,
+              }));
             }
           },
           onData: (data) => {
             if (!isStopped) {
-              optsRef.current.onData(data);
-
-              updateSubscriptionState((prev) => {
-                if (prev.isPending) {
-                  return getPendingResult(prev, data);
-                }
-
-                return prev;
-              });
+              optsRef.current.onData?.(data);
+              updateState((prev) => ({
+                ...prev,
+                status: 'pending',
+                data,
+                error: null,
+              }));
             }
           },
-          onError: (err) => {
+          onError: (error) => {
             if (!isStopped) {
-              optsRef.current.onError?.(err);
+              optsRef.current.onError?.(error);
+              updateState((prev) => ({
+                ...prev,
+                status: 'error',
+                error,
+              }));
             }
           },
-          onStateChange: (state) => {
-            if (state.state === 'idle') {
-              updateSubscriptionState(getIdleResult(/* restart.current */));
-
-              return;
-            }
-
-            if (state.state === 'connecting') {
-              updateSubscriptionState((prev) => {
-                return getConnectingResult(prev, state.data ?? null);
-              });
-
-              return;
-            }
-
-            if (state.state === 'pending') {
-              updateSubscriptionState((prev) => getPendingResult(prev));
-            }
-
-            if (state.state === 'error') {
-              updateSubscriptionState((prev) => {
-                return getErrorResult(prev, state.data);
-              });
-            }
+          onConnectionStateChange: (result) => {
+            const delta = {
+              status: result.state,
+              error: result.error,
+            } as $Result;
+
+            updateState((prev) => {
+              return {
+                ...prev,
+                ...delta,
+              };
+            });
           },
         },
       );
 
-      // const effectRestart = restart.current;
-
-      return () => {
+      currentSubscriptionRef.current = () => {
         isStopped = true;
         subscription.unsubscribe();
-
-        updateSubscriptionState(getIdleResult(/* effectRestart */));
       };
 
       // eslint-disable-next-line react-hooks/exhaustive-deps
     }, [queryKey, enabled]);
 
-    return subscriptionState;
+    const getInitialState = React.useCallback((): $Result => {
+      return enabled
+        ? {
+            data: undefined,
+            error: null,
+            status: 'connecting',
+            reset,
+          }
+        : {
+            data: undefined,
+            error: null,
+            status: 'idle',
+            reset,
+          };
+    }, [enabled, reset]);
+
+    const resultRef = React.useRef<$Result>(getInitialState());
+
+    const [state, setState] = React.useState<$Result>(
+      trackResult(resultRef.current, addTrackedProp),
+    );
+
+    state.reset = reset;
+
+    const updateState = React.useCallback(
+      (callback: (prevState: $Result) => $Result) => {
+        const prev = resultRef.current;
+        const next = (resultRef.current = callback(prev));
+
+        let shouldUpdate = false;
+        for (const key of trackedProps.current) {
+          if (prev[key] !== next[key]) {
+            shouldUpdate = true;
+            break;
+          }
+        }
+        if (shouldUpdate) {
+          setState(trackResult(next, addTrackedProp));
+        }
+      },
+      [addTrackedProp],
+    );
+
+    React.useEffect(() => {
+      if (!enabled) {
+        return;
+      }
+      reset();
+
+      return () => {
+        currentSubscriptionRef.current?.();
+      };
+    }, [reset, enabled]);
+
+    return state;
   }
 
   function useInfiniteQuery(
@@ -550,6 +570,54 @@ export function createRootHooks<
     return hook;
   }
 
+  function usePrefetchInfiniteQuery(
+    path: string[],
+    input: unknown,
+    opts: UseTRPCPrefetchInfiniteQueryOptions<unknown, unknown, TError>,
+  ): void {
+    const context = useContext();
+    const queryKey = getQueryKeyInternal(path, input, 'infinite');
+
+    const defaultOpts = context.queryClient.getQueryDefaults(queryKey);
+
+    const isInputSkipToken = input === skipToken;
+
+    const ssrOpts = useSSRQueryOptionsIfNeeded(queryKey, {
+      ...defaultOpts,
+      ...opts,
+    });
+
+    // request option should take priority over global
+    const shouldAbortOnUnmount =
+      opts?.trpc?.abortOnUnmount ?? context.abortOnUnmount;
+
+    __usePrefetchInfiniteQuery({
+      ...opts,
+      initialPageParam: opts.initialCursor ?? null,
+      queryKey,
+      queryFn: isInputSkipToken
+        ? input
+        : (queryFunctionContext) => {
+            const actualOpts = {
+              ...ssrOpts,
+              trpc: {
+                ...ssrOpts?.trpc,
+                ...(shouldAbortOnUnmount
+                  ? { signal: queryFunctionContext.signal }
+                  : {}),
+              },
+            };
+
+            return context.client.query(
+              ...getClientArgs(queryKey, actualOpts, {
+                pageParam: queryFunctionContext.pageParam ?? opts.initialCursor,
+                direction: queryFunctionContext.direction,
+              }),
+            );
+          },
+    });
+  }
+
   function useSuspenseInfiniteQuery(
     path: readonly string[],
     input: unknown,
@@ -647,6 +715,7 @@ export function createRootHooks<
       {
         queries: queries.map((query) => ({
           ...query,
+          queryFn: query.queryFn,
           queryKey: (query as TRPCQueryOptions<any, any>).queryKey,
         })),
       },
@@ -662,12 +731,14 @@ export function createRootHooks<
     useContext,
     useUtils: useContext,
     useQuery,
+    usePrefetchQuery,
     useSuspenseQuery,
     useQueries,
     useSuspenseQueries,
     useMutation,
     useSubscription,
     useInfiniteQuery,
+    usePrefetchInfiniteQuery,
     useSuspenseInfiniteQuery,
   };
 }