@trpc/tanstack-react-query 0.0.0-alpha.0

package/src/internals/queryOptions.ts

@@ -0,0 +1,199 @@
+import type {
+  DataTag,
+  DefinedInitialDataOptions,
+  QueryClient,
+  QueryFunction,
+  UndefinedInitialDataOptions,
+  UnusedSkipTokenOptions,
+} from '@tanstack/react-query';
+import { queryOptions, skipToken, type SkipToken } from '@tanstack/react-query';
+import type { TRPCClientErrorLike, TRPCUntypedClient } from '@trpc/client';
+import type {
+  coerceAsyncIterableToArray,
+  DistributiveOmit,
+} from '@trpc/server/unstable-core-do-not-import';
+import { isAsyncIterable } from '@trpc/server/unstable-core-do-not-import';
+import type {
+  ResolverDef,
+  TRPCQueryBaseOptions,
+  TRPCQueryKey,
+  TRPCQueryOptionsResult,
+} from './types';
+import {
+  buildQueryFromAsyncIterable,
+  createTRPCOptionsResult,
+  getClientArgs,
+  unwrapLazyArg,
+} from './utils';
+
+type ReservedOptions = 'queryKey' | 'queryFn' | 'queryHashFn' | 'queryHash';
+
+interface UndefinedTRPCQueryOptionsIn<TQueryFnData, TData, TError>
+  extends DistributiveOmit<
+      UndefinedInitialDataOptions<
+        coerceAsyncIterableToArray<TQueryFnData>,
+        TError,
+        coerceAsyncIterableToArray<TData>,
+        TRPCQueryKey
+      >,
+      ReservedOptions
+    >,
+    TRPCQueryBaseOptions {}
+
+interface UndefinedTRPCQueryOptionsOut<TQueryFnData, TOutput, TError>
+  extends UndefinedInitialDataOptions<
+      coerceAsyncIterableToArray<TQueryFnData>,
+      TError,
+      coerceAsyncIterableToArray<TOutput>,
+      TRPCQueryKey
+    >,
+    TRPCQueryOptionsResult {
+  queryKey: DataTag<TRPCQueryKey, coerceAsyncIterableToArray<TOutput>>;
+}
+
+interface DefinedTRPCQueryOptionsIn<TQueryFnData, TData, TError>
+  extends DistributiveOmit<
+      DefinedInitialDataOptions<
+        coerceAsyncIterableToArray<TQueryFnData>,
+        TError,
+        coerceAsyncIterableToArray<TData>,
+        TRPCQueryKey
+      >,
+      ReservedOptions
+    >,
+    TRPCQueryBaseOptions {}
+
+interface DefinedTRPCQueryOptionsOut<TQueryFnData, TData, TError>
+  extends DefinedInitialDataOptions<
+      coerceAsyncIterableToArray<TQueryFnData>,
+      TError,
+      coerceAsyncIterableToArray<TData>,
+      TRPCQueryKey
+    >,
+    TRPCQueryOptionsResult {
+  queryKey: DataTag<TRPCQueryKey, coerceAsyncIterableToArray<TData>>;
+}
+
+interface UnusedSkipTokenTRPCQueryOptionsIn<TQueryFnData, TData, TError>
+  extends DistributiveOmit<
+      UnusedSkipTokenOptions<
+        coerceAsyncIterableToArray<TQueryFnData>,
+        TError,
+        coerceAsyncIterableToArray<TData>,
+        TRPCQueryKey
+      >,
+      ReservedOptions
+    >,
+    TRPCQueryBaseOptions {}
+
+interface UnusedSkipTokenTRPCQueryOptionsOut<TQueryFnData, TOutput, TError>
+  extends UnusedSkipTokenOptions<
+      coerceAsyncIterableToArray<TQueryFnData>,
+      TError,
+      coerceAsyncIterableToArray<TOutput>,
+      TRPCQueryKey
+    >,
+    TRPCQueryOptionsResult {
+  queryKey: DataTag<TRPCQueryKey, coerceAsyncIterableToArray<TOutput>>;
+}
+
+export interface TRPCQueryOptions<TDef extends ResolverDef> {
+  <TQueryFnData extends TDef['output'], TData = TQueryFnData>(
+    input: TDef['input'] | SkipToken,
+    opts: DefinedTRPCQueryOptionsIn<
+      TQueryFnData,
+      TData,
+      TRPCClientErrorLike<{
+        transformer: TDef['transformer'];
+        errorShape: TDef['errorShape'];
+      }>
+    >,
+  ): DefinedTRPCQueryOptionsOut<
+    TQueryFnData,
+    TData,
+    TRPCClientErrorLike<{
+      transformer: TDef['transformer'];
+      errorShape: TDef['errorShape'];
+    }>
+  >;
+  <TQueryFnData extends TDef['output'], TData = TQueryFnData>(
+    input: TDef['input'],
+    opts?: UnusedSkipTokenTRPCQueryOptionsIn<
+      TQueryFnData,
+      TData,
+      TRPCClientErrorLike<{
+        transformer: TDef['transformer'];
+        errorShape: TDef['errorShape'];
+      }>
+    >,
+  ): UnusedSkipTokenTRPCQueryOptionsOut<
+    TQueryFnData,
+    TData,
+    TRPCClientErrorLike<{
+      transformer: TDef['transformer'];
+      errorShape: TDef['errorShape'];
+    }>
+  >;
+  <TQueryFnData extends TDef['output'], TData = TQueryFnData>(
+    input: TDef['input'] | SkipToken,
+    opts?: UndefinedTRPCQueryOptionsIn<
+      TQueryFnData,
+      TData,
+      TRPCClientErrorLike<{
+        transformer: TDef['transformer'];
+        errorShape: TDef['errorShape'];
+      }>
+    >,
+  ): UndefinedTRPCQueryOptionsOut<
+    TQueryFnData,
+    TData,
+    TRPCClientErrorLike<{
+      transformer: TDef['transformer'];
+      errorShape: TDef['errorShape'];
+    }>
+  >;
+}
+
+export function trpcQueryOptions(args: {
+  query: typeof TRPCUntypedClient.prototype.query;
+  queryClient: QueryClient | (() => QueryClient);
+  path: readonly string[];
+  queryKey: TRPCQueryKey;
+  opts: UndefinedTRPCQueryOptionsIn<unknown, unknown, unknown>;
+}) {
+  const { query, path, queryKey, opts } = args;
+  const queryClient = unwrapLazyArg(args.queryClient);
+
+  const inputIsSkipToken = queryKey[1]?.input === skipToken;
+
+  const queryFn: QueryFunction<unknown, TRPCQueryKey> = async (
+    queryFnContext,
+  ) => {
+    const actualOpts = {
+      ...opts,
+      trpc: {
+        ...opts?.trpc,
+        ...(opts?.trpc?.abortOnUnmount
+          ? { signal: queryFnContext.signal }
+          : { signal: null }),
+      },
+    };
+
+    const result = await query(...getClientArgs(queryKey, actualOpts));
+
+    if (isAsyncIterable(result)) {
+      return buildQueryFromAsyncIterable(result, queryClient, queryKey);
+    }
+
+    return result;
+  };
+
+  return Object.assign(
+    queryOptions({
+      ...opts,
+      queryKey,
+      queryFn: inputIsSkipToken ? skipToken : queryFn,
+    }),
+    { trpc: createTRPCOptionsResult({ path }) },
+  );
+}

package/src/internals/subscriptionOptions.ts

@@ -0,0 +1,286 @@
+import { hashKey, skipToken, type SkipToken } from '@tanstack/react-query';
+import type { TRPCClientErrorLike, TRPCUntypedClient } from '@trpc/client';
+import type { TRPCConnectionState } from '@trpc/client/unstable-internals';
+import type { Unsubscribable } from '@trpc/server/observable';
+import type { inferAsyncIterableYield } from '@trpc/server/unstable-core-do-not-import';
+import * as React from 'react';
+import type {
+  ResolverDef,
+  TRPCQueryKey,
+  TRPCQueryOptionsResult,
+} from './types';
+import { createTRPCOptionsResult } from './utils';
+
+interface BaseTRPCSubscriptionOptionsIn<TOutput, TError> {
+  enabled?: boolean;
+  onStarted?: () => void;
+  onData: (data: inferAsyncIterableYield<TOutput>) => void;
+  onError?: (err: TError) => void;
+}
+
+interface UnusedSkipTokenTRPCSubscriptionOptionsIn<TOutput, TError> {
+  onStarted?: () => void;
+  onData: (data: inferAsyncIterableYield<TOutput>) => void;
+  onError?: (err: TError) => void;
+  onConnectionStateChange: (state: TRPCConnectionState<TError>) => void;
+}
+
+interface TRPCSubscriptionOptionsOut<TOutput, TError>
+  extends UnusedSkipTokenTRPCSubscriptionOptionsIn<TOutput, TError>,
+    TRPCQueryOptionsResult {
+  enabled: boolean;
+  queryKey: TRPCQueryKey;
+  subscribe: (
+    innerOpts: UnusedSkipTokenTRPCSubscriptionOptionsIn<TOutput, TError>,
+  ) => Unsubscribable;
+}
+
+export interface TRPCSubscriptionOptions<TDef extends ResolverDef> {
+  (
+    input: TDef['input'],
+    opts?: UnusedSkipTokenTRPCSubscriptionOptionsIn<
+      inferAsyncIterableYield<TDef['output']>,
+      TRPCClientErrorLike<TDef>
+    >,
+  ): TRPCSubscriptionOptionsOut<
+    inferAsyncIterableYield<TDef['output']>,
+    TRPCClientErrorLike<TDef>
+  >;
+  (
+    input: TDef['input'] | SkipToken,
+    opts?: BaseTRPCSubscriptionOptionsIn<
+      inferAsyncIterableYield<TDef['output']>,
+      TRPCClientErrorLike<TDef>
+    >,
+  ): TRPCSubscriptionOptionsOut<
+    inferAsyncIterableYield<TDef['output']>,
+    TRPCClientErrorLike<TDef>
+  >;
+}
+export type TRPCSubscriptionStatus =
+  | 'idle'
+  | 'connecting'
+  | 'pending'
+  | 'error';
+
+export interface TRPCSubscriptionBaseResult<TOutput, TError> {
+  status: TRPCSubscriptionStatus;
+  data: undefined | TOutput;
+  error: null | TError;
+  /**
+   * Reset the subscription
+   */
+  reset: () => void;
+}
+
+export interface TRPCSubscriptionIdleResult<TOutput>
+  extends TRPCSubscriptionBaseResult<TOutput, null> {
+  status: 'idle';
+  data: undefined;
+  error: null;
+}
+
+export interface TRPCSubscriptionConnectingResult<TOutput, TError>
+  extends TRPCSubscriptionBaseResult<TOutput, TError> {
+  status: 'connecting';
+  data: undefined | TOutput;
+  error: TError | null;
+}
+
+export interface TRPCSubscriptionPendingResult<TOutput>
+  extends TRPCSubscriptionBaseResult<TOutput, undefined> {
+  status: 'pending';
+  data: TOutput;
+  error: null;
+}
+
+export interface TRPCSubscriptionErrorResult<TOutput, TError>
+  extends TRPCSubscriptionBaseResult<TOutput, TError> {
+  status: 'error';
+  data: TOutput | undefined;
+  error: TError;
+}
+
+export type TRPCSubscriptionResult<TOutput, TError> =
+  | TRPCSubscriptionIdleResult<TOutput>
+  | TRPCSubscriptionConnectingResult<TOutput, TError>
+  | TRPCSubscriptionErrorResult<TOutput, TError>
+  | TRPCSubscriptionPendingResult<TOutput>;
+
+export const trpcSubscriptionOptions = (args: {
+  subscribe: typeof TRPCUntypedClient.prototype.subscription;
+  path: readonly string[];
+  queryKey: TRPCQueryKey;
+  opts: BaseTRPCSubscriptionOptionsIn<unknown, unknown> | undefined;
+}) => {
+  const { subscribe, path, queryKey, opts } = args;
+  const input = queryKey[1]?.input;
+  const enabled = opts?.enabled ?? input !== skipToken;
+
+  const _subscribe: ReturnType<TRPCSubscriptionOptions<any>>['subscribe'] = (
+    innerOpts,
+  ) => {
+    return subscribe(path.join('.'), input ?? undefined, innerOpts);
+  };
+
+  return {
+    ...opts,
+    enabled,
+    subscribe: _subscribe,
+    queryKey,
+    trpc: createTRPCOptionsResult({ path }),
+  };
+};
+
+export function useSubscription<TOutput, TError>(
+  opts: TRPCSubscriptionOptionsOut<TOutput, TError>,
+): TRPCSubscriptionResult<TOutput, TError> {
+  type $Result = TRPCSubscriptionResult<TOutput, TError>;
+
+  const optsRef = React.useRef(opts);
+  optsRef.current = opts;
+
+  const trackedProps = React.useRef(new Set<keyof $Result>([]));
+
+  const addTrackedProp = React.useCallback((key: keyof $Result) => {
+    trackedProps.current.add(key);
+  }, []);
+
+  type Unsubscribe = () => void;
+  const currentSubscriptionRef = React.useRef<Unsubscribe>();
+
+  const reset = React.useCallback((): void => {
+    // unsubscribe from the previous subscription
+    currentSubscriptionRef.current?.();
+
+    updateState(getInitialState);
+    if (!opts.enabled) {
+      return;
+    }
+
+    let isStopped = false;
+    const subscription = opts.subscribe({
+      onStarted: () => {
+        if (!isStopped) {
+          optsRef.current.onStarted?.();
+          updateState((prev) => ({
+            ...(prev as any),
+            status: 'pending',
+            error: null,
+          }));
+        }
+      },
+      onData: (data) => {
+        if (!isStopped) {
+          optsRef.current.onData?.(data);
+          updateState((prev) => ({
+            ...(prev as any),
+            status: 'pending',
+            data,
+            error: null,
+          }));
+        }
+      },
+      onError: (error) => {
+        if (!isStopped) {
+          optsRef.current.onError?.(error);
+          updateState((prev) => ({
+            ...(prev as any),
+            status: 'error',
+            error,
+          }));
+        }
+      },
+      onConnectionStateChange: (result) => {
+        const delta = {
+          status: result.state,
+          error: result.error,
+        } as $Result;
+
+        updateState((prev) => {
+          return {
+            ...prev,
+            ...delta,
+          };
+        });
+      },
+    });
+
+    currentSubscriptionRef.current = () => {
+      isStopped = true;
+      subscription.unsubscribe();
+    };
+
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [hashKey(opts.queryKey), opts.enabled]);
+
+  const getInitialState = React.useCallback((): $Result => {
+    return opts.enabled
+      ? {
+          data: undefined,
+          error: null,
+          status: 'connecting',
+          reset,
+        }
+      : {
+          data: undefined,
+          error: null,
+          status: 'idle',
+          reset,
+        };
+  }, [opts.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 (!opts.enabled) {
+      return;
+    }
+    reset();
+
+    return () => {
+      currentSubscriptionRef.current?.();
+    };
+  }, [reset, opts.enabled]);
+
+  return state;
+}
+
+function trackResult<T extends object>(
+  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;
+}

package/src/internals/types.ts

@@ -0,0 +1,47 @@
+import type { TRPCRequestOptions } from '@trpc/client';
+
+export type ResolverDef = {
+  input: any;
+  output: any;
+  transformer: boolean;
+  errorShape: any;
+};
+
+export type ExtractCursorType<TInput> = TInput extends { cursor?: any }
+  ? TInput['cursor']
+  : unknown;
+
+export interface TRPCReactRequestOptions
+  // For RQ, we use their internal AbortSignals instead of letting the user pass their own
+  extends Omit<TRPCRequestOptions, 'signal'> {
+  /**
+   * Opt out of SSR for this query by passing `ssr: false`
+   */
+  ssr?: boolean;
+  /**
+   * Opt out or into aborting request on unmount
+   */
+  abortOnUnmount?: boolean;
+}
+
+export interface TRPCQueryBaseOptions {
+  /**
+   * tRPC-related options
+   */
+  trpc?: TRPCReactRequestOptions;
+}
+
+export interface TRPCQueryOptionsResult {
+  trpc: {
+    path: string;
+  };
+}
+
+export type QueryType = 'any' | 'infinite' | 'query';
+
+export type TRPCQueryKey = [
+  readonly string[],
+  { input?: unknown; type?: Exclude<QueryType, 'any'> }?,
+];
+
+export type TRPCMutationKey = [readonly string[]]; // = [TRPCQueryKey[0]]

package/src/internals/utils.ts

@@ -0,0 +1,140 @@
+import { skipToken, type QueryClient } from '@tanstack/react-query';
+import { isFunction, isObject } from '@trpc/server/unstable-core-do-not-import';
+import type {
+  QueryType,
+  TRPCMutationKey,
+  TRPCQueryKey,
+  TRPCQueryOptionsResult,
+} from './types';
+
+/**
+ * @internal
+ */
+export function createTRPCOptionsResult(value: {
+  path: readonly string[];
+}): TRPCQueryOptionsResult['trpc'] {
+  const path = value.path.join('.');
+
+  return {
+    path,
+  };
+}
+
+/**
+ * @internal
+ */
+export function getClientArgs<TOptions>(
+  queryKey: TRPCQueryKey,
+  opts: TOptions,
+  infiniteParams?: {
+    pageParam: any;
+    direction: 'forward' | 'backward';
+  },
+) {
+  const path = queryKey[0];
+  let input = queryKey[1]?.input;
+  if (infiniteParams) {
+    input = {
+      ...(input ?? {}),
+      ...(infiniteParams.pageParam ? { cursor: infiniteParams.pageParam } : {}),
+      direction: infiniteParams.direction,
+    };
+  }
+  return [path.join('.'), input, (opts as any)?.trpc] as const;
+}
+
+/**
+ * @internal
+ */
+export async function buildQueryFromAsyncIterable(
+  asyncIterable: AsyncIterable<unknown>,
+  queryClient: QueryClient,
+  queryKey: TRPCQueryKey,
+) {
+  const queryCache = queryClient.getQueryCache();
+
+  const query = queryCache.build(queryClient, {
+    queryKey,
+  });
+
+  query.setState({
+    data: [],
+    status: 'success',
+  });
+
+  const aggregate: unknown[] = [];
+  for await (const value of asyncIterable) {
+    aggregate.push(value);
+
+    query.setState({
+      data: [...aggregate],
+    });
+  }
+  return aggregate;
+}
+
+/**
+ * To allow easy interactions with groups of related queries, such as
+ * invalidating all queries of a router, we use an array as the path when
+ * storing in tanstack query.
+ **/
+export function getQueryKeyInternal(
+  path: readonly string[],
+  input: unknown,
+  type: QueryType,
+): TRPCQueryKey {
+  // Construct a query key that is easy to destructure and flexible for
+  // partial selecting etc.
+  // https://github.com/trpc/trpc/issues/3128
+
+  // some parts of the path may be dot-separated, split them up
+  const splitPath = path.flatMap((part) => part.split('.'));
+
+  if (!input && (!type || type === 'any')) {
+    // this matches also all mutations (see `getMutationKeyInternal`)
+
+    // for `utils.invalidate()` to match all queries (including vanilla react-query)
+    // we don't want nested array if path is empty, i.e. `[]` instead of `[[]]`
+    return splitPath.length ? [splitPath] : ([] as unknown as TRPCQueryKey);
+  }
+
+  if (
+    type === 'infinite' &&
+    isObject(input) &&
+    ('direction' in input || 'cursor' in input)
+  ) {
+    const {
+      cursor: _,
+      direction: __,
+      ...inputWithoutCursorAndDirection
+    } = input;
+    return [
+      splitPath,
+      {
+        input: inputWithoutCursorAndDirection,
+        type: 'infinite',
+      },
+    ];
+  }
+  return [
+    splitPath,
+    {
+      ...(typeof input !== 'undefined' &&
+        input !== skipToken && { input: input }),
+      ...(type && type !== 'any' && { type: type }),
+    },
+  ];
+}
+
+export function getMutationKeyInternal(
+  path: readonly string[],
+): TRPCMutationKey {
+  // some parts of the path may be dot-separated, split them up
+  const splitPath = path.flatMap((part) => part.split('.'));
+
+  return splitPath.length ? [splitPath] : ([] as unknown as TRPCMutationKey);
+}
+
+export function unwrapLazyArg<T>(valueOrLazy: T | (() => T)): T {
+  return (isFunction(valueOrLazy) ? valueOrLazy() : valueOrLazy) as T;
+}