Add datatag to ConnectQueryKey so queryClient can infer type info

.gitignore

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

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

@@ -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.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>;
 }

packages/connect-query-core/src/create-infinite-query-options.ts

@@ -67,7 +67,7 @@ function createUnaryInfiniteQueryFn<
   },
 ): QueryFunction<
   MessageShape<O>,
-  ConnectQueryKey,
+  ConnectQueryKey<O>,
   MessageInitShape<I>[ParamKey]
 > {
   return async (context) => {
@@ -104,10 +104,10 @@ export function createInfiniteQueryOptions<
     O,
     ParamKey
   >["getNextPageParam"];
-  queryKey: ConnectQueryKey;
+  queryKey: ConnectQueryKey<O>;
   queryFn: QueryFunction<
     MessageShape<O>,
-    ConnectQueryKey,
+    ConnectQueryKey<O>,
     MessageInitShape<I>[ParamKey]
   >;
   structuralSharing: (oldData: unknown, newData: unknown) => unknown;
@@ -131,7 +131,7 @@ export function createInfiniteQueryOptions<
     O,
     ParamKey
   >["getNextPageParam"];
-  queryKey: ConnectQueryKey;
+  queryKey: ConnectQueryKey<O>;
   queryFn: SkipToken;
   structuralSharing: (oldData: unknown, newData: unknown) => unknown;
   initialPageParam: MessageInitShape<I>[ParamKey];
@@ -156,11 +156,11 @@ export function createInfiniteQueryOptions<
     O,
     ParamKey
   >["getNextPageParam"];
-  queryKey: ConnectQueryKey;
+  queryKey: ConnectQueryKey<O>;
   queryFn:
     | QueryFunction<
         MessageShape<O>,
-        ConnectQueryKey,
+        ConnectQueryKey<O>,
         MessageInitShape<I>[ParamKey]
       >
     | SkipToken;
@@ -187,11 +187,11 @@ export function createInfiniteQueryOptions<
     O,
     ParamKey
   >["getNextPageParam"];
-  queryKey: ConnectQueryKey;
+  queryKey: ConnectQueryKey<O>;
   queryFn:
     | QueryFunction<
         MessageShape<O>,
-        ConnectQueryKey,
+        ConnectQueryKey<O>,
         MessageInitShape<I>[ParamKey]
       >
     | SkipToken;