feat: expose querykey getter and fix some querykey stuff

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;
 

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

@@ -29,6 +29,11 @@ import {
   TRPCFetchQueryOptions,
   contextProps,
 } from '../../internals/context';
+import {
+  QueryKey,
+  QueryType,
+  getArrayQueryKey,
+} from '../../internals/getArrayQueryKey';
 import { getQueryKey } from '../../internals/getQueryKey';
 
 type DecorateProcedure<
@@ -158,6 +163,15 @@ type DecorateProcedure<
   getInfiniteData(
     input?: inferProcedureInput<TProcedure>,
   ): InfiniteData<inferTransformedProcedureOutput<TProcedure>> | undefined;
+
+  /**
+   * Method to extract the query key for a procedure
+   * @link https://trpc.io/docs/useContext#-the-function-i-want-isnt-here
+   */
+  getQueryKey(
+    input: inferProcedureInput<TProcedure>,
+    type: QueryType,
+  ): QueryKey;
 };
 
 /**
@@ -174,6 +188,12 @@ type DecorateRouter = {
     filters?: InvalidateQueryFilters,
     options?: InvalidateOptions,
   ): Promise<void>;
+
+  /**
+   * Method to extract the query key for a router
+   * @link https://trpc.io/docs/useContext#-the-function-i-want-isnt-here
+   */
+  getQueryKey(input: undefined, type: 'any'): QueryKey;
 };
 
 /**
@@ -265,6 +285,7 @@ export function createReactQueryUtilsProxy<
           context.setInfiniteQueryData(queryKey, updater, ...rest),
         getData: () => context.getQueryData(queryKey),
         getInfiniteData: () => context.getInfiniteQueryData(queryKey),
+        getQueryKey: () => getArrayQueryKey(queryKey, rest[0]),
       };
 
       return contextMap[utilName]();

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

@@ -1,5 +1,6 @@
 /* eslint-disable react-hooks/exhaustive-deps */
 import { getServerAndReactClient } from './__reactHelpers';
+import { useIsFetching } from '@tanstack/react-query';
 import { render, waitFor } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import { initTRPC } from '@trpc/server/src/core';
@@ -798,4 +799,121 @@ describe('query keys are stored separtely', () => {
     `);
     expect(data.infinite).toBeUndefined();
   });
+
+  describe('getQueryKeys', () => {
+    test('no input', async () => {
+      const { proxy, App } = ctx;
+
+      function MyComponent() {
+        const utils = proxy.useContext();
+        const happy = utils.post.all.getQueryKey(undefined, 'query');
+
+        // @ts-expect-error - post.all has no input
+        const sad1 = utils.post.all.getQueryKey('foo');
+        // @ts-expect-error - need to specify type
+        const sad2 = utils.post.all.getQueryKey(undefined);
+
+        return <pre data-testid="qKey">{JSON.stringify(happy)}</pre>;
+      }
+
+      const utils = render(
+        <App>
+          <MyComponent />
+        </App>,
+      );
+
+      await waitFor(() => {
+        expect(utils.getByTestId('qKey')).toHaveTextContent(
+          JSON.stringify([['post', 'all'], { type: 'query' }]),
+        );
+      });
+    });
+
+    test('with input', async () => {
+      const { proxy, App } = ctx;
+
+      function MyComponent() {
+        const utils = proxy.useContext();
+        const happy = utils.post.byId.getQueryKey({ id: 1 }, 'query');
+
+        // @ts-expect-error - post.byId has required input
+        const sad1 = utils.post.byId.getQueryKey(undefined, 'query');
+        // @ts-expect-error - need to specify type
+        const sad2 = utils.post.byId.getQueryKey({ id: 1 });
+
+        return <pre data-testid="qKey">{JSON.stringify(happy)}</pre>;
+      }
+
+      const utils = render(
+        <App>
+          <MyComponent />
+        </App>,
+      );
+
+      await waitFor(() => {
+        expect(utils.getByTestId('qKey')).toHaveTextContent(
+          JSON.stringify([
+            ['post', 'byId'],
+            { input: { id: 1 }, type: 'query' },
+          ]),
+        );
+      });
+    });
+
+    test('on router', async () => {
+      const { proxy, App } = ctx;
+
+      function MyComponent() {
+        const utils = proxy.useContext();
+        const happy = utils.post.getQueryKey(undefined, 'any');
+
+        // @ts-expect-error - router has no input
+        const sad = utils.post.getQueryKey('foo', 'any');
+
+        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 utils = proxy.useContext();
+        const qKey = utils.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');
+      });
+    });
+  });
 });