@trpc/react-query 11.0.0-alpha-tmp-issues-5851-take-two.499 → 11.0.0-alpha-tmp-subscription-connection-state.485

package/src/shared/hooks/createHooksInternal.tsx

@@ -19,7 +19,10 @@ import type { SSRState, TRPCContextState } from '../../internals/context';
 import { TRPCContext } from '../../internals/context';
 import { getClientArgs } from '../../internals/getClientArgs';
 import type { TRPCQueryKey } from '../../internals/getQueryKey';
-import { getQueryKeyInternal } from '../../internals/getQueryKey';
+import {
+  getMutationKeyInternal,
+  getQueryKeyInternal,
+} from '../../internals/getQueryKey';
 import { useHookResult } from '../../internals/useHookResult';
 import type {
   TRPCUseQueries,
@@ -28,23 +31,46 @@ import type {
 import { createUtilityFunctions } from '../../utils/createUtilityFunctions';
 import { createUseQueries } from '../proxy/useQueriesProxy';
 import type { CreateTRPCReactOptions, UseMutationOverride } from '../types';
-import type {
-  CreateClient,
-  TRPCProvider,
-  TRPCQueryOptions,
-  UseTRPCInfiniteQueryOptions,
-  UseTRPCInfiniteQueryResult,
-  UseTRPCMutationOptions,
-  UseTRPCMutationResult,
-  UseTRPCQueryOptions,
-  UseTRPCQueryResult,
-  UseTRPCSubscriptionOptions,
-  UseTRPCSuspenseInfiniteQueryOptions,
-  UseTRPCSuspenseInfiniteQueryResult,
-  UseTRPCSuspenseQueryOptions,
-  UseTRPCSuspenseQueryResult,
+// 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,
 } from './types';
 
+const trackResult = <
+  T extends UseTRPCSubscriptionResult<unknown, unknown, unknown>,
+>(
+  result: T,
+  onTrackResult: (key: keyof T) => void,
+): T => {
+  const trackedResult = new Proxy(result, {
+    get(target, prop) {
+      onTrackResult(prop as keyof T);
+      return target[prop as keyof T];
+    },
+  });
+
+  return trackedResult;
+};
+
 /**
  * @internal
  */
@@ -267,7 +293,7 @@ export function createRootHooks<
     const { client } = useContext();
     const queryClient = useQueryClient();
 
-    const mutationKey = [path];
+    const mutationKey = getMutationKeyInternal(path);
 
     const defaultOpts = queryClient.defaultMutationOptions(
       queryClient.getMutationDefaults(mutationKey),
@@ -306,9 +332,69 @@ 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 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;
+
+        currentResult.current = newResult;
+
+        let shouldUpdate = false;
+
+        for (const key of trackedProps.current) {
+          if (oldResult[key] !== newResult[key]) {
+            shouldUpdate = true;
+            break;
+          }
+        }
+
+        if (shouldUpdate) {
+          setSubscriptionState(trackResult(newResult, addTrackedProp));
+          return;
+        }
+      },
+      [addTrackedProp],
+    );
+
     const { client } = useContext();
 
     const optsRef = React.useRef<typeof opts>(opts);
@@ -319,6 +405,7 @@ export function createRootHooks<
         return;
       }
       let isStopped = false;
+
       const subscription = client.subscription(
         path.join('.'),
         input ?? undefined,
@@ -331,6 +418,14 @@ export function createRootHooks<
           onData: (data) => {
             if (!isStopped) {
               optsRef.current.onData(data);
+
+              updateSubscriptionState((prev) => {
+                if (prev.isPending) {
+                  return getPendingResult(prev, data);
+                }
+
+                return prev;
+              });
             }
           },
           onError: (err) => {
@@ -338,14 +433,47 @@ export function createRootHooks<
               optsRef.current.onError?.(err);
             }
           },
+          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);
+              });
+            }
+          },
         },
       );
+
+      // const effectRestart = restart.current;
+
       return () => {
         isStopped = true;
         subscription.unsubscribe();
+
+        updateSubscriptionState(getIdleResult(/* effectRestart */));
       };
+
       // eslint-disable-next-line react-hooks/exhaustive-deps
     }, [queryKey, enabled]);
+
+    return subscriptionState;
   }
 
   function useInfiniteQuery(

package/src/shared/hooks/types.ts

@@ -18,7 +18,9 @@ import type {
   UseSuspenseQueryResult,
 } from '@tanstack/react-query';
 import type {
+  ConnectionState,
   CreateTRPCClientOptions,
+  TRPCConnectionStateMessage,
   TRPCRequestOptions,
   TRPCUntypedClient,
 } from '@trpc/client';
@@ -139,9 +141,371 @@ export interface UseTRPCMutationOptions<
 export interface UseTRPCSubscriptionOptions<TOutput, TError> {
   enabled?: boolean;
   onStarted?: () => void;
-  onData: (data: TOutput) => void;
+  /**
+   * @deprecated use onStateChange instead
+   */
   onError?: (err: TError) => void;
+  onData: (data: TOutput) => void;
+  onStateChange?: (state: TRPCConnectionStateMessage<TError>) => void;
+}
+
+export interface restartSubscriptionOptionsBase {
+  /**
+   * Cancel the current subscription and re-establish a new one
+   * - Defaults to `true`
+   * - Set to `false` no new subscription will be established if there is already an active subscription
+   */
+  cancelSubscription?: boolean;
+}
+
+export interface restartSubscriptionOptionsWithLastEventId
+  extends restartSubscriptionOptionsBase {
+  /**
+   * Defaults to `true` in case of failure or if the subscription is still active
+   * - When the susbscription has successfully completed, it will be set to `false`
+   */
+  sendLastEventId?: boolean;
 }
+
+export type restartSubscriptionOptions<TInput> = TInput extends {
+  lastEventId?: string | null;
+} | void
+  ? restartSubscriptionOptionsWithLastEventId
+  : restartSubscriptionOptionsBase;
+
+export type restartSubscriptionFn<TInput> = (
+  options?: restartSubscriptionOptions<TInput>,
+) => void;
+
+export interface TRPCSubscriptionBaseResult<_TInput, TOutput, TError> {
+  /**
+   * The last data received from the subscription
+   */
+  data: TOutput | null;
+  /**
+   * The timestamp for when the connection was last (re-)established
+   */
+  connectionStartedAt: number;
+  /**
+   * The timestamp for when the connection was initially established
+   */
+  initialConnectionStartedAt: number;
+  /**
+   * The error that caused the subscription to stop
+   * - Defaults to `null`
+   * - Resets to `null` after the subscription is restarted and the connection is re-established
+   */
+  error: TError | null;
+  /**
+   * The timestamp for when the last error was captured
+   */
+  errorUpdatedAt: number;
+  /**
+   * The reason for the reconnection
+   * - Resets to `null` after the subscription is restarted and the connection is re-established
+   */
+  connectionError: TError | null;
+  /**
+   * Reconnection attempts since last successful connection
+   */
+  connectionAttemptCount: number;
+  /**
+   * The timestamp for when the last reconnection error was captured
+   */
+  connectionErrorUpdatedAt: number;
+  /**
+   * Is `true` when the subscription is establishing the initial connection
+   */
+  isStarting: boolean;
+  /**
+   * Is `true` when the subscription has successfully connected at least once
+   */
+  isStarted: boolean;
+  /**
+   * Is `true` if the subscription is (re-)establishing a connection
+   * - Alias for `status === 'connecting'`
+   */
+  isConnecting: boolean;
+  /**
+   * Is `true` when the subscription is connected and listening for data
+   * - Alias for `status === 'pending'`
+   */
+  isPending: boolean;
+  /**
+   * Is `true` if the subscription is re-establishing a connection
+   * - Alias for `status === 'connecting' && isStarted === true`
+   */
+  isReconnecting: boolean;
+  /**
+   * Is `true` if the subscription ended in an error state
+   * - Alias for `status === 'error'`
+   */
+  isError: boolean;
+  /**
+   * The current state of the subscription
+   * - Will be:
+   *  - `'idle'` when the subscription is not enabled
+   *  - `'connecting'` when the subscription is (re-)establishing the connection
+   *  - `'pending'` when the subscription is connected and receiving data
+   *  - `'error'` when the subscription has stopped due to an error
+   */
+  status: ConnectionState;
+  /**
+   * Restart the subscription
+   */
+  // restart: restartSubscriptionFn<TInput>;
+}
+
+export interface TRPCSubscriptionIdleResult<TInput, TOutput, TError>
+  extends TRPCSubscriptionBaseResult<TInput, TOutput, TError> {
+  data: null;
+  error: null;
+  errorUpdatedAt: 0;
+  connectionError: null;
+  isStarting: false;
+  isStarted: false;
+  isConnecting: false;
+  isPending: false;
+  isReconnecting: false;
+  isError: false;
+  connectionAttemptCount: 0;
+  connectionErrorUpdatedAt: 0;
+  connectionStartedAt: 0;
+  initialConnectionStartedAt: 0;
+  status: 'idle';
+}
+
+const defaultIdleResult: Omit<
+  TRPCSubscriptionIdleResult<unknown, unknown, unknown>,
+  'restart'
+> = {
+  data: null,
+  connectionStartedAt: 0,
+  initialConnectionStartedAt: 0,
+  error: null,
+  errorUpdatedAt: 0,
+  connectionError: null,
+  connectionAttemptCount: 0,
+  isStarting: false,
+  isStarted: false,
+  isConnecting: false,
+  isPending: false,
+  isReconnecting: false,
+  isError: false,
+  connectionErrorUpdatedAt: 0,
+  status: 'idle',
+};
+
+export const getIdleResult = <
+  TInput,
+  TOutput,
+  TError,
+>(): //  restart: restartSubscriptionFn<TInput>,
+TRPCSubscriptionIdleResult<TInput, TOutput, TError> => {
+  return {
+    ...defaultIdleResult,
+    data: null,
+    //    restart,
+  };
+};
+
+export interface TRPCSubscriptionStartingResult<TInput, TOutput, TError>
+  extends TRPCSubscriptionBaseResult<TInput, TOutput, TError> {
+  connectionStartedAt: 0;
+  initialConnectionStartedAt: 0;
+  error: null;
+  data: null;
+  connectionError: TError | null;
+  isStarting: true;
+  isStarted: false;
+  isConnecting: true;
+  isPending: false;
+  isReconnecting: false;
+  isError: false;
+  status: 'connecting';
+}
+
+export const getStartingResult = <TInput, TOutput, TError>(
+  // restart: restartSubscriptionFn<TInput>,
+  previous?:
+    | TRPCSubscriptionIdleResult<TInput, TOutput, TError>
+    | TRPCSubscriptionErrorResult<TInput, TOutput, TError>
+    | TRPCSubscriptionStartingResult<TInput, TOutput, TError>,
+  error?: TError | null,
+): TRPCSubscriptionStartingResult<TInput, TOutput, TError> => {
+  const now = Date.now();
+
+  if (previous) {
+    return {
+      ...defaultIdleResult,
+      ...previous,
+      data: null,
+      connectionError: error ?? null,
+      isStarting: true,
+      isConnecting: true,
+      errorUpdatedAt: 0,
+      connectionAttemptCount: previous.connectionAttemptCount + 1,
+      connectionErrorUpdatedAt: error ? now : 0,
+      connectionStartedAt: 0,
+      initialConnectionStartedAt: 0,
+      error: null,
+      isStarted: false,
+      isError: false,
+      status: 'connecting',
+    };
+  }
+
+  return {
+    ...getIdleResult(/*restart*/),
+    connectionError: error ?? null,
+    isStarting: true,
+    isConnecting: true,
+    errorUpdatedAt: 0,
+    connectionAttemptCount: 0,
+    connectionErrorUpdatedAt: error ? now : 0,
+    status: 'connecting',
+  };
+};
+
+export interface TRPCSubscriptionPendingResult<TInput, TOutput, TError>
+  extends TRPCSubscriptionBaseResult<TInput, TOutput, TError> {
+  error: null;
+  connectionError: null;
+  connectionAttemptCount: 0;
+  isStarting: false;
+  isStarted: true;
+  isConnecting: false;
+  isPending: true;
+  isReconnecting: false;
+  isError: false;
+  status: 'pending';
+}
+
+export const getPendingResult = <TInput, TOutput, TError>(
+  previous: UseTRPCSubscriptionResult<TInput, TOutput, TError>,
+  data?: TOutput,
+): TRPCSubscriptionPendingResult<TInput, TOutput, TError> => {
+  const time = Date.now();
+
+  if (previous.isStarting) {
+    return {
+      ...previous,
+      error: null,
+      isStarting: false,
+      isStarted: true,
+      isConnecting: false,
+      isReconnecting: false,
+      isPending: true,
+      status: 'pending',
+      connectionStartedAt: time,
+      connectionAttemptCount: 0,
+      connectionError: null,
+      initialConnectionStartedAt: time,
+    };
+  }
+
+  return {
+    ...previous,
+    data: data ?? previous.data,
+    error: null,
+    isError: false,
+    isStarting: false,
+    isStarted: true,
+    isConnecting: false,
+    isReconnecting: false,
+    isPending: true,
+    status: 'pending',
+    connectionStartedAt: time,
+    connectionAttemptCount: 0,
+    connectionError: null,
+  };
+};
+
+export interface TRPCSubscriptionReconnectingResult<TInput, TOutput, TError>
+  extends TRPCSubscriptionBaseResult<TInput, TOutput, TError> {
+  error: null;
+  connectionError: TError;
+  isStarting: false;
+  isStarted: true;
+  isConnecting: true;
+  isPending: false;
+  isReconnecting: true;
+  isError: false;
+  status: 'connecting';
+}
+
+export const getReconnectingResult = <TInput, TOutput, TError>(
+  previous:
+    | TRPCSubscriptionPendingResult<TInput, TOutput, TError>
+    | TRPCSubscriptionReconnectingResult<TInput, TOutput, TError>,
+  error: TError,
+): TRPCSubscriptionReconnectingResult<TInput, TOutput, TError> => {
+  return {
+    ...previous,
+    isStarting: false,
+    isStarted: true,
+    isConnecting: true,
+    isReconnecting: true,
+    isPending: false,
+    status: 'connecting',
+    connectionError: error,
+    connectionAttemptCount: previous.connectionAttemptCount + 1,
+  };
+};
+
+export const getConnectingResult = <TInput, TOutput, TError>(
+  previous: UseTRPCSubscriptionResult<TInput, TOutput, TError>,
+  error: TError | null,
+): TRPCSubscriptionConnectingResult<TInput, TOutput, TError> => {
+  if (previous.isReconnecting || previous.isPending) {
+    if (!error) throw new Error('Reconnecting without error?');
+
+    return getReconnectingResult(previous, error);
+  }
+
+  return getStartingResult(/*previous.restart, */ previous, error);
+};
+
+export interface TRPCSubscriptionErrorResult<TInput, TOutput, TError>
+  extends TRPCSubscriptionBaseResult<TInput, TOutput, TError> {
+  error: TError;
+  isStarting: false;
+  isStarted: true; // Not sure about this one
+  isConnecting: false;
+  isPending: false;
+  isReconnecting: false;
+  isError: true;
+  status: 'error';
+}
+
+export const getErrorResult = <TInput, TOutput, TError>(
+  previous: UseTRPCSubscriptionResult<TInput, TOutput, TError>,
+  error: TError,
+): TRPCSubscriptionErrorResult<TInput, TOutput, TError> => {
+  return {
+    ...previous,
+    isStarting: false,
+    isStarted: true,
+    isConnecting: false,
+    isReconnecting: false,
+    isPending: false,
+    isError: true,
+    status: 'error',
+    error,
+    errorUpdatedAt: Date.now(),
+  };
+};
+
+export type TRPCSubscriptionConnectingResult<TInput, TOutput, TError> =
+  | TRPCSubscriptionStartingResult<TInput, TOutput, TError>
+  | TRPCSubscriptionReconnectingResult<TInput, TOutput, TError>;
+
+export type UseTRPCSubscriptionResult<TInput, TOutput, TError> =
+  | TRPCSubscriptionIdleResult<TInput, TOutput, TError>
+  | TRPCSubscriptionPendingResult<TInput, TOutput, TError>
+  | TRPCSubscriptionConnectingResult<TInput, TOutput, TError>
+  | TRPCSubscriptionErrorResult<TInput, TOutput, TError>;
+
 export interface TRPCProviderProps<TRouter extends AnyRouter, TSSRContext>
   extends TRPCContextProps<TRouter, TSSRContext> {
   children: ReactNode;

package/src/shared/proxy/utilsProxy.ts

@@ -15,6 +15,7 @@ import type {
 import type { TRPCClientError } from '@trpc/client';
 import { createTRPCClientProxy } from '@trpc/client';
 import type {
+  AnyMutationProcedure,
   AnyQueryProcedure,
   AnyRootTypes,
   AnyRouter,
@@ -37,10 +38,14 @@ import type {
 } from '../../internals/context';
 import { contextProps } from '../../internals/context';
 import type { QueryKeyKnown, QueryType } from '../../internals/getQueryKey';
-import { getQueryKeyInternal } from '../../internals/getQueryKey';
+import {
+  getMutationKeyInternal,
+  getQueryKeyInternal,
+} from '../../internals/getQueryKey';
+import type { InferMutationOptions } from '../../utils/inferReactQueryProcedure';
 import type { ExtractCursorType } from '../hooks/types';
 
-type DecorateProcedure<
+type DecorateQueryProcedure<
   TRoot extends AnyRootTypes,
   TProcedure extends AnyQueryProcedure,
 > = {
@@ -225,6 +230,25 @@ type DecorateProcedure<
     | undefined;
 };
 
+type DecorateMutationProcedure<
+  TRoot extends AnyRootTypes,
+  TProcedure extends AnyMutationProcedure,
+> = {
+  setMutationDefaults(
+    options:
+      | InferMutationOptions<TRoot, TProcedure>
+      | ((args: {
+          canonicalMutationFn: NonNullable<
+            InferMutationOptions<TRoot, TProcedure>['mutationFn']
+          >;
+        }) => InferMutationOptions<TRoot, TProcedure>),
+  ): void;
+
+  getMutationDefaults(): InferMutationOptions<TRoot, TProcedure> | undefined;
+
+  isMutating(): number;
+};
+
 /**
  * this is the type that is used to add in procedures that can be used on
  * an entire router
@@ -252,14 +276,16 @@ export type DecoratedProcedureUtilsRecord<
   [TKey in keyof TRecord]: TRecord[TKey] extends infer $Value
     ? $Value extends RouterRecord
       ? DecoratedProcedureUtilsRecord<TRoot, $Value> & DecorateRouter
-      : // utils only apply to queries
-      $Value extends AnyQueryProcedure
-      ? DecorateProcedure<TRoot, $Value>
+      : $Value extends AnyQueryProcedure
+      ? DecorateQueryProcedure<TRoot, $Value>
+      : $Value extends AnyMutationProcedure
+      ? DecorateMutationProcedure<TRoot, $Value>
       : never
     : never;
 }; // Add functions that should be available at utils root
 
-type AnyDecoratedProcedure = DecorateProcedure<any, any>;
+type AnyDecoratedProcedure = DecorateQueryProcedure<any, any> &
+  DecorateMutationProcedure<any, any>;
 
 export type CreateReactUtils<
   TRouter extends AnyRouter,
@@ -296,6 +322,9 @@ export const getQueryType = (
     case 'setInfiniteData':
       return 'infinite';
 
+    case 'setMutationDefaults':
+    case 'getMutationDefaults':
+    case 'isMutating':
     case 'cancel':
     case 'invalidate':
     case 'refetch':
@@ -321,6 +350,9 @@ function createRecursiveUtilsProxy<TRouter extends AnyRouter>(
     const queryKey = getQueryKeyInternal(path, input, queryType);
 
     const contextMap: Record<keyof AnyDecoratedProcedure, () => unknown> = {
+      /**
+       * DecorateQueryProcedure
+       */
       fetch: () => context.fetchQuery(queryKey, ...args),
       fetchInfinite: () => context.fetchInfiniteQuery(queryKey, args[0]),
       prefetch: () => context.prefetchQuery(queryKey, ...args),
@@ -340,6 +372,15 @@ function createRecursiveUtilsProxy<TRouter extends AnyRouter>(
       },
       getData: () => context.getQueryData(queryKey),
       getInfiniteData: () => context.getInfiniteQueryData(queryKey),
+      /**
+       * DecorateMutationProcedure
+       */
+      setMutationDefaults: () =>
+        context.setMutationDefaults(getMutationKeyInternal(path), input),
+      getMutationDefaults: () =>
+        context.getMutationDefaults(getMutationKeyInternal(path)),
+      isMutating: () =>
+        context.isMutating({ mutationKey: getMutationKeyInternal(path) }),
     };
 
     return contextMap[utilName]();