Add datatag to ConnectQueryKey so queryClient can infer type info (#532)

By adding DataTag annotations to the ConnectQueryKey, QueryClient
methods can now infer the return type when calling things like
`getQueryData` or `invalidateQueries`.

This is at least a partial alternative to #468 which provides better
inference abilities while keeping our api surface smaller. There may
still be value down the line with very well defined query client methods
but I suspect this change covers 90% of the current use cases.

---------

Signed-off-by: Paul Sachs <psachs@buf.build>

Author: paul-sachs

.gitignore

@@ -3,3 +3,4 @@
 node_modules
 /packages/*/dist
 /packages/*/coverage
+tsconfig.vitest-temp.json

README.md

@@ -260,19 +260,17 @@ Any additional `options` you pass to `useMutation` will be merged with the optio
 
 ### `createConnectQueryKey`
 
-```ts
-function createConnectQueryKey<Desc extends DescMethod | DescService>(
-  params: KeyParams<Desc>,
-): ConnectQueryKey;
-```
-
-This function is used under the hood of `useQuery` and other hooks to compute a [`queryKey`](https://tanstack.com/query/v4/docs/react/guides/query-keys) for TanStack Query. You can use it to create (partial) keys yourself to filter queries.
+This function is used under the hood of `useQuery` and other hooks to compute a [`queryKey`](https://tanstack.com/query/v4/docs/react/guides/query-keys) for TanStack Query. You can use it to create keys yourself to filter queries.
 
 `useQuery` creates a query key with the following parameters:
 
 1. The qualified name of the RPC.
 2. The transport being used.
 3. The request message.
+4. The cardinality of the RPC (either "finite" or "infinite").
+5. Adds a DataTag which brands the key with the associated data type of the response.
+
+The DataTag type allows @tanstack/react-query functions to properly infer the type of the data returned by the query. This is useful for things like `QueryClient.setQueryData` and `QueryClient.getQueryData`.
 
 To create the same key manually, you simply provide the same parameters:
 
@@ -355,7 +353,7 @@ function callUnaryMethod<I extends DescMessage, O extends DescMessage>(
 
 This API allows you to directly call the method using the provided transport. Use this if you need to manually call a method outside of the context of a React component, or need to call it where you can't use hooks.
 
-### `createProtobufSafeUpdater`
+### `createProtobufSafeUpdater` (deprecated)
 
 Creates a typesafe updater that can be used to update data in a query cache. Used in combination with a queryClient.
 
@@ -387,6 +385,8 @@ queryClient.setQueryData(
 
 ```
 
+** Note: This API is deprecated and will be removed in a future version. `ConnectQueryKey` now contains type information to make it safer to use `setQueryData` directly. **
+
 ### `createQueryOptions`
 
 ```ts
@@ -599,21 +599,15 @@ Connect-Query does require React, but the core (`createConnectQueryKey` and `cal
 
 ### How do I do Prefetching?
 
-When you might not have access to React context, you can use the `create` series of functions and provide a transport directly. For example:
+When you might not have access to React context, you can use `createQueryOptions` and provide a transport directly. For example:
 
 ```ts
 import { say } from "./gen/eliza-ElizaService_connectquery";
 
 function prefetch() {
-  return queryClient.prefetchQuery({
-    queryKey: createConnectQueryKey({
-      schema: say,
-      transport: myTransport,
-      input: { sentence: "Hello" },
-      cardinality: "finite",
-    }),
-    queryFn: () => callUnaryMethod(myTransport, say, { sentence: "Hello" }),
-  });
+  return queryClient.prefetchQuery(
+    createQueryOptions(say, { sentence: "Hello" }, { transport: myTransport }),
+  );
 }
 ```
 

package-lock.json

@@ -44,7 +44,7 @@
   "peerDependencies": {
     "@bufbuild/protobuf": "2.x",
     "@connectrpc/connect": "^2.0.1",
-    "@tanstack/query-core": "5.x"
+    "@tanstack/query-core": ">=5.62.0"
   },
   "files": [
     "dist/**"

packages/connect-query-core/package.json

@@ -14,14 +14,20 @@
 
 import { create } from "@bufbuild/protobuf";
 import type { Transport } from "@connectrpc/connect";
-import { ElizaService, SayRequestSchema } from "test-utils/gen/eliza_pb.js";
+import {
+  ElizaService,
+  SayRequestSchema,
+  SayResponseSchema,
+  type SayResponse,
+} from "test-utils/gen/eliza_pb.js";
 import { ListRequestSchema, ListService } from "test-utils/gen/list_pb.js";
-import { describe, expect, it } from "vitest";
+import { describe, expect, expectTypeOf, it } from "vitest";
 
 import { createConnectQueryKey } from "./connect-query-key.js";
 import { skipToken } from "./index.js";
 import { createMessageKey } from "./message-key.js";
 import { createTransportKey } from "./transport-key.js";
+import { type InfiniteData, QueryClient } from "@tanstack/query-core";
 
 describe("createConnectQueryKey", () => {
   const fakeTransport: Transport = {
@@ -125,12 +131,96 @@ describe("createConnectQueryKey", () => {
 
   it("cannot except invalid input", () => {
     createConnectQueryKey({
+      // @ts-expect-error(2322) cannot create a key with invalid input
       schema: ElizaService.method.say,
       input: {
-        // @ts-expect-error(2322) cannot create a key with invalid input
         sentence: 1,
       },
       cardinality: undefined,
     });
   });
+
+  it("contains type hints to indicate the output type", () => {
+    const sampleQueryClient = new QueryClient();
+    const key = createConnectQueryKey({
+      schema: ElizaService.method.say,
+      input: create(SayRequestSchema, { sentence: "hi" }),
+      cardinality: "finite",
+    });
+    const data = sampleQueryClient.getQueryData(key);
+
+    expectTypeOf(data).toEqualTypeOf<SayResponse | undefined>();
+  });
+
+  it("supports typesafe data updaters", () => {
+    const sampleQueryClient = new QueryClient();
+    const key = createConnectQueryKey({
+      schema: ElizaService.method.say,
+      input: create(SayRequestSchema, { sentence: "hi" }),
+      cardinality: "finite",
+    });
+    // @ts-expect-error(2345) this is a test to check if the type is correct
+    sampleQueryClient.setQueryData(key, { sentence: 1 });
+    // @ts-expect-error(2345) $typename is required
+    sampleQueryClient.setQueryData(key, {
+      sentence: "a proper value but missing $typename",
+    });
+    sampleQueryClient.setQueryData(
+      key,
+      create(SayResponseSchema, { sentence: "a proper value" }),
+    );
+
+    sampleQueryClient.setQueryData(key, (prev) => {
+      expectTypeOf(prev).toEqualTypeOf<SayResponse | undefined>();
+      return create(SayResponseSchema, {
+        sentence: "a proper value",
+      });
+    });
+  });
+
+  describe("infinite queries", () => {
+    it("contains type hints to indicate the output type", () => {
+      const sampleQueryClient = new QueryClient();
+      const key = createConnectQueryKey({
+        schema: ElizaService.method.say,
+        input: create(SayRequestSchema, { sentence: "hi" }),
+        cardinality: "infinite",
+      });
+      const data = sampleQueryClient.getQueryData(key);
+
+      expectTypeOf(data).toEqualTypeOf<InfiniteData<SayResponse> | undefined>();
+    });
+
+    it("supports typesafe data updaters", () => {
+      const sampleQueryClient = new QueryClient();
+      const key = createConnectQueryKey({
+        schema: ElizaService.method.say,
+        input: create(SayRequestSchema, { sentence: "hi" }),
+        cardinality: "infinite",
+      });
+      sampleQueryClient.setQueryData(key, {
+        pages: [
+          // @ts-expect-error(2345) make sure the shape is as expected
+          { sentence: 1 },
+        ],
+      });
+      sampleQueryClient.setQueryData(key, {
+        // @ts-expect-error(2345) $typename is required
+        pages: [{ sentence: "a proper value but missing $typename" }],
+      });
+      sampleQueryClient.setQueryData(key, {
+        pageParams: [0],
+        pages: [create(SayResponseSchema, { sentence: "a proper value" })],
+      });
+      sampleQueryClient.setQueryData(key, (prev) => {
+        expectTypeOf(prev).toEqualTypeOf<
+          InfiniteData<SayResponse> | undefined
+        >();
+        return {
+          pageParams: [0],
+          pages: [create(SayResponseSchema, { sentence: "a proper value" })],
+        };
+      });
+    });
+  });
 });

packages/connect-query-core/src/connect-query-key.test.ts

@@ -18,13 +18,60 @@ import type {
   DescMethodUnary,
   DescService,
   MessageInitShape,
+  MessageShape,
 } from "@bufbuild/protobuf";
-import type { Transport } from "@connectrpc/connect";
-import type { SkipToken } from "@tanstack/query-core";
+import type { ConnectError, Transport } from "@connectrpc/connect";
+import type { DataTag, InfiniteData, SkipToken } from "@tanstack/query-core";
 
 import { createMessageKey } from "./message-key.js";
 import { createTransportKey } from "./transport-key.js";
 
+type SharedConnectQueryOptions = {
+  /**
+   * A key for a Transport reference, created with createTransportKey().
+   */
+  transport?: string;
+  /**
+   * The name of the service, e.g. connectrpc.eliza.v1.ElizaService
+   */
+  serviceName: string;
+  /**
+   * The name of the method, e.g. Say.
+   */
+  methodName?: string;
+  /**
+   * A key for the request message, created with createMessageKey(),
+   * or "skipped".
+   */
+  input?: Record<string, unknown> | "skipped";
+};
+
+type InfiniteConnectQueryKey<OutputMessage extends DescMessage = DescMessage> =
+  DataTag<
+    [
+      "connect-query",
+      SharedConnectQueryOptions & {
+        /** This data represents a infinite, paged result */
+        cardinality: "infinite";
+      },
+    ],
+    InfiniteData<MessageShape<OutputMessage>>,
+    ConnectError
+  >;
+
+type FiniteConnectQueryKey<OutputMessage extends DescMessage = DescMessage> =
+  DataTag<
+    [
+      "connect-query",
+      SharedConnectQueryOptions & {
+        /** This data represents a finite result */
+        cardinality: "finite";
+      },
+    ],
+    MessageShape<OutputMessage>,
+    ConnectError
+  >;
+
 /**
  * TanStack Query manages query caching for you based on query keys. `QueryKey`s in TanStack Query are arrays with arbitrary JSON-serializable data - typically handwritten for each endpoint.
  *
@@ -44,35 +91,15 @@ import { createTransportKey } from "./transport-key.js";
  *   }
  * ]
  */
-export type ConnectQueryKey = [
-  /**
-   * To distinguish Connect query keys from other query keys, they always start with the string "connect-query".
-   */
-  "connect-query",
-  {
-    /**
-     * A key for a Transport reference, created with createTransportKey().
-     */
-    transport?: string;
-    /**
-     * The name of the service, e.g. connectrpc.eliza.v1.ElizaService
-     */
-    serviceName: string;
-    /**
-     * The name of the method, e.g. Say.
-     */
-    methodName?: string;
-    /**
-     * A key for the request message, created with createMessageKey(),
-     * or "skipped".
-     */
-    input?: Record<string, unknown> | "skipped";
-    /**
-     * Whether this is an infinite query, or a regular one.
-     */
-    cardinality?: "infinite" | "finite" | undefined;
-  },
-];
+export type ConnectQueryKey<OutputMessage extends DescMessage = DescMessage> =
+  | InfiniteConnectQueryKey<OutputMessage>
+  | FiniteConnectQueryKey<OutputMessage>
+  | [
+      "connect-query",
+      SharedConnectQueryOptions & {
+        cardinality: undefined;
+      },
+    ];
 
 type KeyParamsForMethod<Desc extends DescMethod> = {
   /**
@@ -152,14 +179,48 @@ type KeyParamsForService<Desc extends DescService> = {
  *
  * @see ConnectQueryKey for information on the components of Connect-Query's keys.
  */
+export function createConnectQueryKey<
+  I extends DescMessage,
+  O extends DescMessage,
+>(
+  params: KeyParamsForMethod<DescMethodUnary<I, O>> & {
+    cardinality: "finite";
+  },
+): FiniteConnectQueryKey<O>;
+export function createConnectQueryKey<
+  I extends DescMessage,
+  O extends DescMessage,
+>(
+  params: KeyParamsForMethod<DescMethodUnary<I, O>> & {
+    cardinality: "infinite";
+  },
+): InfiniteConnectQueryKey<O>;
+export function createConnectQueryKey<
+  I extends DescMessage,
+  O extends DescMessage,
+>(
+  params: KeyParamsForMethod<DescMethodUnary<I, O>> & {
+    cardinality: undefined;
+  },
+): ConnectQueryKey<O>;
+export function createConnectQueryKey<
+  O extends DescMessage,
+  Desc extends DescService,
+>(params: KeyParamsForService<Desc>): ConnectQueryKey<O>;
 export function createConnectQueryKey<
   I extends DescMessage,
   O extends DescMessage,
   Desc extends DescService,
 >(
   params: KeyParamsForMethod<DescMethodUnary<I, O>> | KeyParamsForService<Desc>,
-): ConnectQueryKey {
-  const props: ConnectQueryKey[1] =
+): ConnectQueryKey<O> {
+  const props: {
+    serviceName: string;
+    methodName?: string;
+    transport?: string;
+    cardinality?: "finite" | "infinite";
+    input?: "skipped" | Record<string, unknown>;
+  } =
     params.schema.kind == "rpc"
       ? {
           serviceName: params.schema.parent.typeName,
@@ -185,5 +246,5 @@ export function createConnectQueryKey<
       );
     }
   }
-  return ["connect-query", props];
+  return ["connect-query", props] as ConnectQueryKey<O>;
 }