feat: expose querykey getter and fix some querykey stuff (trpc#3302)

packages/react-query/src/createTRPCReact.tsx

@@ -15,6 +15,7 @@ import {
   inferTransformedSubscriptionOutput,
 } from '@trpc/server/shared';
 import { useMemo } from 'react';
+import { QueryKey, QueryType } from './internals/getArrayQueryKey';
 import { TRPCUseQueries } from './internals/useQueries';
 import {
   CreateReactUtilsProxy,
@@ -48,6 +49,15 @@ export type DecorateProcedure<
   TPath extends string,
 > = TProcedure extends AnyQueryProcedure
   ? {
+      /**
+       * Method to extract the query key for a procedure
+       * @param type - defaults to `any`
+       * @link https://trpc.io/docs/useContext#-the-function-i-want-isnt-here
+       */
+      getQueryKey: (
+        input: inferProcedureInput<TProcedure>,
+        type?: QueryType,
+      ) => QueryKey;
       useQuery: <
         TQueryFnData = inferTransformedProcedureOutput<TProcedure>,
         TData = inferTransformedProcedureOutput<TProcedure>,
@@ -164,7 +174,9 @@ export type DecoratedProcedureRecord<
   TPath extends string = '',
 > = {
   [TKey in keyof TProcedures]: TProcedures[TKey] extends AnyRouter
-    ? DecoratedProcedureRecord<
+    ? {
+        getQueryKey: () => QueryKey;
+      } & DecoratedProcedureRecord<
         TProcedures[TKey]['_def']['record'],
         TFlags,
         `${TPath}${TKey & string}.`

packages/react-query/src/internals/context.tsx

@@ -160,13 +160,6 @@ export interface TRPCContextState<
     filters?: InvalidateQueryFilters,
     options?: InvalidateOptions,
   ): Promise<void>;
-  /**
-   * @link https://react-query.tanstack.com/guides/query-invalidation
-   */
-  invalidateQueries(
-    filters?: InvalidateQueryFilters,
-    options?: InvalidateOptions,
-  ): Promise<void>;
 
   /**
    * @link https://react-query.tanstack.com/reference/QueryClient#queryclientresetqueries

packages/react-query/src/internals/getArrayQueryKey.test.ts

@@ -1,6 +1,17 @@
 import { getArrayQueryKey } from './getArrayQueryKey';
 
 test('getArrayQueryKey', () => {
+  // empty path should not nest an extra array
+  expect(getArrayQueryKey('', 'any')).toMatchInlineSnapshot(`Array []`);
+
+  // should not nest an empty object
+  expect(getArrayQueryKey('foo', 'any')).toMatchInlineSnapshot(`
+    Array [
+      Array [
+        "foo",
+      ],
+    ]
+  `);
   expect(getArrayQueryKey('foo', 'query')).toMatchInlineSnapshot(`
     Array [
       Array [

packages/react-query/src/internals/getArrayQueryKey.ts

@@ -1,16 +1,21 @@
 export type QueryType = 'query' | 'infinite' | 'any';
 
+export type QueryKey = [
+  string[],
+  { input?: unknown; type?: Exclude<QueryType, 'any'> }?,
+];
+
 /**
  * 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. This function converts from the `.` separated
  * path passed around internally by both the legacy and proxy implementation.
  * https://github.com/trpc/trpc/issues/2611
- */
+ **/
 export function getArrayQueryKey(
   queryKey: string | [string] | [string, ...unknown[]] | unknown[],
   type: QueryType,
-): [string[], { input?: unknown; type?: Exclude<QueryType, 'any'> }] {
+): QueryKey {
   const queryKeyArrayed = Array.isArray(queryKey) ? queryKey : [queryKey];
   const [path, input] = queryKeyArrayed;
 
@@ -20,6 +25,10 @@ export function getArrayQueryKey(
   // Construct a query key that is easy to destructure and flexible for
   // partial selecting etc.
   // https://github.com/trpc/trpc/issues/3128
+  if (!input && (!type || type === 'any'))
+    // 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 arrayPath.length ? [arrayPath] : ([] as unknown as QueryKey);
   return [
     arrayPath,
     {

packages/react-query/src/internals/getQueryKey.ts

@@ -6,5 +6,6 @@ export function getQueryKey(
   path: string,
   input: unknown,
 ): [string] | [string, unknown] {
-  return input === undefined ? [path] : [path, input];
+  if (path.length) return input === undefined ? [path] : [path, input];
+  return [] as unknown as [string];
 }

packages/react-query/src/shared/proxy/decorationProxy.ts

@@ -1,5 +1,6 @@
 import { AnyRouter } from '@trpc/server';
 import { createRecursiveProxy } from '@trpc/server/shared';
+import { getArrayQueryKey } from '../../internals/getArrayQueryKey';
 import { getQueryKey } from '../../internals/getQueryKey';
 import { CreateReactQueryHooks } from '../hooks/createHooksInternal';
 
@@ -28,6 +29,12 @@ export function createReactProxyDecoration<
     const [input, ...rest] = args;
 
     const queryKey = getQueryKey(path, input);
+
+    // Expose queryKey helper
+    if (lastArg === 'getQueryKey') {
+      return getArrayQueryKey(queryKey, (rest[0] as any) ?? 'any');
+    }
+
     if (lastArg.startsWith('useSuspense')) {
       const opts = rest[0] || {};
       const fn =

packages/react-query/src/shared/proxy/utilsProxy.ts

@@ -171,6 +171,7 @@ type DecorateRouter = {
    * @link https://react-query.tanstack.com/guides/query-invalidation
    */
   invalidate(
+    input?: undefined,
     filters?: InvalidateQueryFilters,
     options?: InvalidateOptions,
   ): Promise<void>;

packages/tests/server/interop/react/invalidateQueries.test.tsx

@@ -192,12 +192,15 @@ describe('invalidateQueries()', () => {
           <button
             data-testid="invalidate-5-predicate"
             onClick={() => {
-              utils.invalidateQueries({
+              utils.invalidateQueries(undefined, {
                 predicate(opts) {
                   const { queryKey } = opts;
-                  const [path, input] = queryKey;
+                  const [path, rest] = queryKey;
 
-                  return path === 'count' && input === 'test';
+                  return (
+                    JSON.stringify(path) === JSON.stringify(['count']) &&
+                    (rest as any)?.input === 'test'
+                  );
                 },
               });
             }}

packages/tests/server/react/getQueryKey.test.tsx

@@ -0,0 +1,178 @@
+import { getServerAndReactClient } from './__reactHelpers';
+import { useIsFetching } from '@tanstack/react-query';
+import { render, waitFor } from '@testing-library/react';
+import { initTRPC } from '@trpc/server';
+import { konn } from 'konn/dist-cjs';
+import React from 'react';
+import { z } from 'zod';
+
+type Post = {
+  id: number;
+  text: string;
+};
+
+const defaultPost = { id: 0, text: 'new post' };
+const ctx = konn()
+  .beforeEach(() => {
+    const t = initTRPC.create();
+
+    const posts: Post[] = [defaultPost];
+
+    const appRouter = t.router({
+      post: t.router({
+        byId: t.procedure
+          .input(
+            z.object({
+              id: z.number(),
+            }),
+          )
+          .query(({ input }) => posts.find((post) => post.id === input.id)),
+        all: t.procedure.query(() => posts),
+      }),
+    });
+
+    return getServerAndReactClient(appRouter);
+  })
+  .afterEach(async (ctx) => {
+    await ctx?.close?.();
+  })
+  .done();
+
+describe('getQueryKeys', () => {
+  test('no input', async () => {
+    const { proxy, App } = ctx;
+
+    function MyComponent() {
+      const happy1 = proxy.post.all.getQueryKey(undefined, 'query');
+      const happy2 = proxy.post.all.getQueryKey();
+
+      // @ts-expect-error - post.all has no input
+      const sad = proxy.post.all.getQueryKey('foo');
+
+      return (
+        <>
+          <pre data-testid="qKey1">{JSON.stringify(happy1)}</pre>
+          <pre data-testid="qKey2">{JSON.stringify(happy2)}</pre>
+        </>
+      );
+    }
+
+    const utils = render(
+      <App>
+        <MyComponent />
+      </App>,
+    );
+
+    await waitFor(() => {
+      expect(utils.getByTestId('qKey1')).toHaveTextContent(
+        JSON.stringify([['post', 'all'], { type: 'query' }]),
+      );
+      expect(utils.getByTestId('qKey2')).toHaveTextContent(
+        JSON.stringify([['post', 'all']]),
+      );
+    });
+  });
+
+  test('with input', async () => {
+    const { proxy, App } = ctx;
+
+    function MyComponent() {
+      const happy1 = proxy.post.byId.getQueryKey({ id: 1 }, 'query');
+
+      // doesn't really make sense but should still work
+      const happyIsh = proxy.post.byId.getQueryKey({ id: 1 });
+
+      // @ts-expect-error - post.byId has required input
+      const sad = proxy.post.byId.getQueryKey(undefined, 'query');
+
+      return (
+        <>
+          <pre data-testid="qKey1">{JSON.stringify(happy1)}</pre>
+          <pre data-testid="qKey2">{JSON.stringify(happyIsh)}</pre>
+        </>
+      );
+    }
+
+    const utils = render(
+      <App>
+        <MyComponent />
+      </App>,
+    );
+
+    await waitFor(() => {
+      expect(utils.getByTestId('qKey1')).toHaveTextContent(
+        JSON.stringify([['post', 'byId'], { input: { id: 1 }, type: 'query' }]),
+      );
+      expect(utils.getByTestId('qKey2')).toHaveTextContent(
+        JSON.stringify([['post', 'byId'], { input: { id: 1 } }]),
+      );
+    });
+  });
+
+  test('on router', async () => {
+    const { proxy, App } = ctx;
+
+    function MyComponent() {
+      const happy = proxy.post.getQueryKey();
+
+      // @ts-expect-error - router has no input
+      const sad = proxy.post.getQueryKey('foo');
+
+      return (
+        <div>
+          <pre data-testid="qKey">{JSON.stringify(happy)}</pre>
+        </div>
+      );
+    }
+
+    const utils = render(
+      <App>
+        <MyComponent />
+      </App>,
+    );
+
+    await waitFor(() => {
+      expect(utils.getByTestId('qKey')).toHaveTextContent(
+        JSON.stringify([['post']]),
+      );
+    });
+  });
+
+  test('forwarded to a real method', async () => {
+    const { proxy, App } = ctx;
+
+    function MyComponent() {
+      proxy.post.all.useQuery();
+
+      const qKey = proxy.post.all.getQueryKey(undefined, 'query');
+      const isFetching = useIsFetching(qKey);
+
+      return <div>{isFetching}</div>;
+    }
+
+    const utils = render(
+      <App>
+        <MyComponent />
+      </App>,
+    );
+
+    // should be fetching initially, and then not
+    expect(utils.container).toHaveTextContent('1');
+    await waitFor(() => {
+      expect(utils.container).toHaveTextContent('0');
+    });
+  });
+
+  test('outside of the react context', () => {
+    const { proxy } = ctx;
+
+    const all = proxy.post.all.getQueryKey(undefined, 'query');
+    const byId = proxy.post.byId.getQueryKey({ id: 1 }, 'query');
+
+    expect(all).toEqual([['post', 'all'], { type: 'query' }]);
+    expect(byId).toEqual([
+      ['post', 'byId'],
+      { input: { id: 1 }, type: 'query' },
+    ]);
+  });
+});

packages/tests/server/react/invalidateQueries.test.tsx

@@ -186,7 +186,7 @@ describe('invalidateQueries()', () => {
           <button
             data-testid="invalidate-4-predicate"
             onClick={() => {
-              utils.invalidate({
+              utils.invalidate(undefined, {
                 predicate(opts) {
                   const { queryKey } = opts;
                   const [path, input] = queryKey;

packages/tests/server/react/useContext.test.tsx

@@ -61,6 +61,10 @@ const ctx = konn()
             return newPost;
           }),
       }),
+
+      greeting: t.router({
+        get: t.procedure.query(() => 'hello'),
+      }),
     });
 
     return getServerAndReactClient(appRouter);