trpc · Nick-Lucas · Nov 14, 2024 · Nov 14, 2024 · Nov 14, 2024 · Nov 14, 2024
diff --git a/packages/tanstack-react-query/src/index.ts b/packages/tanstack-react-query/src/index.ts
@@ -1,4 +1,8 @@
 export { createTRPCContext } from './internals/Context';
-export type { TRPCOptionsProxy } from './internals/createOptionsProxy';
+export type {
+  TRPCOptionsProxy,
+  InferInput,
+  InferOutput,
+} from './internals/createOptionsProxy';
 export { createTRPCOptionsProxy } from './internals/createOptionsProxy';
 export { useSubscription } from './internals/subscriptionOptions';
diff --git a/packages/tanstack-react-query/src/internals/createOptionsProxy.ts b/packages/tanstack-react-query/src/internals/createOptionsProxy.ts
@@ -1,4 +1,4 @@
-import { type QueryClient } from '@tanstack/react-query';
+import { type QueryClient, type QueryFilters } from '@tanstack/react-query';
 import {
   getUntypedClient,
   TRPCUntypedClient,
@@ -33,10 +33,48 @@ import {
   trpcSubscriptionOptions,
   type TRPCSubscriptionOptions,
 } from './subscriptionOptions';
-import type { QueryType, ResolverDef } from './types';
+import type { QueryType, ResolverDef, TRPCQueryKey } from './types';
 import { getQueryKeyInternal } from './utils';
 
+export interface DecorateQueryKeyable {
+  /**
+   * Calculate the Tanstack Query Key for a Route
+   *
+   * @see https://tanstack.com/query/latest/docs/framework/react/guides/query-keys
+   */
+  queryKey: () => TRPCQueryKey;
+
+  /**
+   * Calculate a Tanstack Query Filter for a Route
+   *
+   * @see https://tanstack.com/query/latest/docs/framework/react/guides/filters
+   */
+  queryFilter: () => QueryFilters;
+}
+
+export type InferInput<
+  TProcedure extends
+    | DecorateQueryProcedure<any>
+    | DecorateMutationProcedure<any>,
+> = TProcedure['~input'];
+
+export type InferOutput<
+  TProcedure extends
+    | DecorateQueryProcedure<any>
+    | DecorateMutationProcedure<any>,
+> = TProcedure['~output'];
+
 export interface DecorateQueryProcedure<TDef extends ResolverDef> {
+  /**
+   * @internal prefer using InferInput to access input type
+   */
+  '~input': TDef['input'];
+
+  /**
+   * @internal prefer using InferOutput to access input type
+   */
+  '~output': TDef['output'];
+
   /**
    * @see https://tanstack.com/query/latest/docs/framework/react/reference/queryOptions#queryoptions
    */
@@ -46,9 +84,33 @@ export interface DecorateQueryProcedure<TDef extends ResolverDef> {
    * @see https://tanstack.com/query/latest/docs/framework/react/reference/infiniteQueryOptions#infinitequeryoptions
    */
   infiniteQueryOptions: TRPCInfiniteQueryOptions<TDef>;
+
+  /**
+   * Calculate the Tanstack Query Key for a Query Procedure
+   *
+   * @see https://tanstack.com/query/latest/docs/framework/react/guides/query-keys
+   */
+  queryKey: (input?: TDef['input']) => TRPCQueryKey;
+
+  /**
+   * Calculate a Tanstack Query Filter for a Query Procedure
+   *
+   * @see https://tanstack.com/query/latest/docs/framework/react/guides/filters
+   */
+  queryFilter: (input?: TDef['input']) => QueryFilters;
 }
 
 export interface DecorateMutationProcedure<TDef extends ResolverDef> {
+  /**
+   * @internal prefer using InferInput to access input type
+   */
+  '~input': TDef['input'];
+
+  /**
+   * @internal prefer using InferOutput to access input type
+   */
+  '~output': TDef['output'];
+
   /**
    * @see
    */
@@ -82,7 +144,7 @@ export type DecoratedProcedureUtilsRecord<
 > = {
   [TKey in keyof TRecord]: TRecord[TKey] extends infer $Value
     ? $Value extends RouterRecord
-      ? DecoratedProcedureUtilsRecord<TRoot, $Value>
+      ? DecoratedProcedureUtilsRecord<TRoot, $Value> & DecorateQueryKeyable
       : $Value extends AnyProcedure
       ? DecorateProcedure<
           $Value['_def']['type'],
@@ -101,7 +163,8 @@ export type TRPCOptionsProxy<TRouter extends AnyRouter> =
   DecoratedProcedureUtilsRecord<
     TRouter['_def']['_config']['$types'],
     TRouter['_def']['record']
-  >;
+  > &
+    DecorateQueryKeyable;
 
 export interface TRPCOptionsProxyOptionsBase {
   queryClient: QueryClient;
@@ -132,12 +195,14 @@ type UtilsMethods =
   | keyof DecorateSubscriptionProcedure<any>;
 
 function getQueryType(method: UtilsMethods) {
-  return {
+  const map: Partial<Record<UtilsMethods, QueryType>> = {
     queryOptions: 'query',
     infiniteQueryOptions: 'infinite',
     subscriptionOptions: 'any',
     mutationOptions: 'any',
-  }[method] as QueryType;
+  };
+
+  return map[method];
 }
 
 export function createTRPCOptionsProxy<TRouter extends AnyRouter>(
@@ -172,23 +237,32 @@ export function createTRPCOptionsProxy<TRouter extends AnyRouter>(
   return createRecursiveProxy(({ args, path: _path }) => {
     const path = [..._path];
     const utilName = path.pop() as UtilsMethods;
-    const [input, userOptions] = args as any[];
+    const [arg1, arg2] = args as any[];
 
     const queryType = getQueryType(utilName);
-    const queryKey = getQueryKeyInternal(path, input, queryType);
+    const queryKey = getQueryKeyInternal(path, arg1, queryType ?? 'any');
 
     const contextMap: Record<UtilsMethods, () => unknown> = {
-      infiniteQueryOptions: () =>
-        trpcInfiniteQueryOptions({
-          opts: userOptions,
+      '~input': undefined as any,
+      '~output': undefined as any,
+      queryKey: () => queryKey,
+      queryFilter: (): QueryFilters => {
+        return {
+          queryKey: queryKey,
+        };
+      },
+      infiniteQueryOptions: () => {
+        return trpcInfiniteQueryOptions({
+          opts: arg2,
           path,
           queryClient: opts.queryClient,
-          queryKey,
+          queryKey: queryKey,
           query: callIt('query'),
-        }),
+        });
+      },
       queryOptions: () => {
         return trpcQueryOptions({
-          opts: userOptions,
+          opts: arg2,
           path,
           queryClient: opts.queryClient,
           queryKey: queryKey,
@@ -197,7 +271,7 @@ export function createTRPCOptionsProxy<TRouter extends AnyRouter>(
       },
       mutationOptions: () => {
         return trpcMutationOptions({
-          opts: userOptions,
+          opts: arg1,
           path,
           queryClient: opts.queryClient,
           mutate: callIt('mutation'),
@@ -206,9 +280,9 @@ export function createTRPCOptionsProxy<TRouter extends AnyRouter>(
       },
       subscriptionOptions: () => {
         return trpcSubscriptionOptions({
-          opts: userOptions,
+          opts: arg2,
           path,
-          queryKey,
+          queryKey: queryKey,
           subscribe: callIt('subscription'),
         });
       },

diff --git a/packages/tanstack-react-query/test/mutationOptions.test.tsx b/packages/tanstack-react-query/test/mutationOptions.test.tsx
@@ -55,9 +55,26 @@ const ctx = konn()
 describe('mutationOptions', () => {
   test('useMutation', async () => {
     const { App, useTRPC } = ctx;
+
+    const calls: string[] = [];
+
     function MyComponent() {
       const trpc = useTRPC();
-      const options = trpc.post.create.mutationOptions({});
+
+      const options = trpc.post.create.mutationOptions({
+        onMutate(variables) {
+          calls.push('onMutate');
+        },
+        onSettled(variables) {
+          calls.push('onSettled');
+        },
+        onError(variables) {
+          calls.push('onError');
+        },
+        onSuccess(variables) {
+          calls.push('onSuccess');
+        },
+      });
       expect(options.trpc.path).toBe('post.create');
 
       const mutation = useMutation(options);
@@ -86,5 +103,7 @@ describe('mutationOptions', () => {
     await waitFor(() => {
       expect(utils.container).toHaveTextContent(`__mutationResult`);
     });
+
+    expect(calls).toEqual(['onMutate', 'onSuccess', 'onSettled']);
   });
 });
diff --git a/packages/tanstack-react-query/test/polymorphism.common.tsx b/packages/tanstack-react-query/test/polymorphism.common.tsx
@@ -0,0 +1,5 @@
+import { initTRPC } from '@trpc/server';
+
+export const t = initTRPC.create();
+
+export type $RootTypes = (typeof t)['_config']['$types'];
diff --git a/packages/tanstack-react-query/test/polymorphism.factory.tsx b/packages/tanstack-react-query/test/polymorphism.factory.tsx
@@ -0,0 +1,104 @@
+//
+// This file contains a useful pattern in tRPC,
+//  building factories which can produce common functionality over a homologous data source.
+//
+import { TRPCError } from '@trpc/server';
+import type {
+  AnyRootTypes,
+  createBuilder,
+} from '@trpc/server/unstable-core-do-not-import';
+import z from 'zod';
+import type { TRPCOptionsProxy } from '../src';
+import type { $RootTypes } from './polymorphism.common';
+import { t } from './polymorphism.common';
+
+//
+// DTOs
+//
+
+export const FileExportRequest = z.object({
+  name: z.string().min(0),
+  filter: z.string().min(0),
+});
+
+export const FileExportStatus = z.object({
+  id: z.number().min(0),
+  name: z.string().min(0),
+  downloadUri: z.string().optional(),
+  createdAt: z.date(),
+});
+export type FileExportStatusType = z.infer<typeof FileExportStatus>;
+
+//
+// Dependencies
+//
+
+type BaseProcedure<TRoot extends AnyRootTypes> = ReturnType<
+  typeof createBuilder<TRoot['ctx'], TRoot['meta']>
+>;
+
+export type DataProvider = FileExportStatusType[];
+
+//
+// Set up a route factory which can be re-used for different data sources.
+// In this case just with a simple array data source a POC
+//
+
+let COUNTER = 1;
+
+export function createExportRoute<
+  TBaseProcedure extends BaseProcedure<$RootTypes>,
+>(baseProcedure: TBaseProcedure, dataProvider: DataProvider) {
+  return t.router({
+    start: baseProcedure
+      .input(FileExportRequest)
+      .output(FileExportStatus)
+      .mutation(async (opts) => {
+        const exportInstance: FileExportStatusType = {
+          id: COUNTER++,
+          name: opts.input.name,
+          createdAt: new Date(),
+          downloadUri: undefined,
+        };
+
+        dataProvider.push(exportInstance);
+
+        return exportInstance;
+      }),
+    list: baseProcedure.output(z.array(FileExportStatus)).query(async () => {
+      return dataProvider;
+    }),
+    status: baseProcedure
+      .input(z.object({ id: z.number().min(0) }))
+      .output(FileExportStatus)
+      .query(async (opts) => {
+        const index = dataProvider.findIndex(
+          (item) => item.id === opts.input.id,
+        );
+
+        const exportInstance = dataProvider[index];
+
+        if (!exportInstance) {
+          throw new TRPCError({
+            code: 'NOT_FOUND',
+          });
+        }
+
+        // When status is polled a second time the download should be ready
+        dataProvider[index] = {
+          ...exportInstance,
+          downloadUri: `example.com/export-${exportInstance.name}.csv`,
+        };
+
+        return exportInstance;
+      }),
+  });
+}
+
+//
+// Generate abstract types which can be used by the client
+//
+
+type ExportRouteType = ReturnType<typeof createExportRoute>;
+
+export type ExportRouteLike = TRPCOptionsProxy<ExportRouteType>;