Use message and transport keys for query keys

Signed-off-by: Timo Stamm <[email protected]>

packages/connect-query/src/call-unary-method.test.ts

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import { create } from "@bufbuild/protobuf";
 import type { QueryFunctionContext } from "@tanstack/react-query";
 import { useQueries } from "@tanstack/react-query";
 import { renderHook, waitFor } from "@testing-library/react";
@@ -21,30 +22,35 @@ import { callUnaryMethod } from "./call-unary-method.js";
 import type { ConnectQueryKey } from "./connect-query-key.js";
 import { createConnectQueryKey } from "./connect-query-key.js";
 import { defaultOptions } from "./default-options.js";
-import { ElizaService } from "./gen/eliza_pb.js";
+import type { SayRequest } from "./gen/eliza_pb.js";
+import { ElizaService, SayRequestSchema } from "./gen/eliza_pb.js";
 import { mockEliza, wrapper } from "./test/test-utils.js";
 
 describe("callUnaryMethod", () => {
   it("can be used with useQueries", async () => {
+    const transport = mockEliza({
+      sentence: "Response 1",
+    });
     const { result } = renderHook(
       () => {
+        const input: SayRequest = create(SayRequestSchema, {
+          sentence: "query 1",
+        });
         const [query1] = useQueries({
           queries: [
             {
-              queryKey: createConnectQueryKey(ElizaService.method.say, {
-                sentence: "query 1",
+              queryKey: createConnectQueryKey({
+                method: ElizaService.method.say,
+                input,
+                transport,
               }),
               queryFn: async ({
-                queryKey,
                 signal,
               }: QueryFunctionContext<ConnectQueryKey>) => {
-                const transport = mockEliza({
-                  sentence: "Response 1",
-                });
                 const res = await callUnaryMethod(
                   transport,
                   ElizaService.method.say,
-                  queryKey[2],
+                  input,
                   {
                     signal,
                   },

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

@@ -12,52 +12,65 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import { create } from "@bufbuild/protobuf";
 import { skipToken } from "@tanstack/react-query";
 import { describe, expect, it } from "vitest";
 
 import { createConnectQueryKey } from "./connect-query-key.js";
 import { ElizaService, SayRequestSchema } from "./gen/eliza_pb.js";
 import { createMessageKey } from "./message-key.js";
 
-describe("makeQueryKey", () => {
-  const methodDescriptor = {
-    input: SayRequestSchema,
-    name: "name",
-    parent: ElizaService,
-  };
-
+describe("createConnectQueryKey", () => {
   it("makes a query key with input", () => {
-    const key = createConnectQueryKey(methodDescriptor, {
-      sentence: "someValue",
+    const key = createConnectQueryKey({
+      method: ElizaService.method.say,
+      input: create(SayRequestSchema, { sentence: "hi" }),
     });
     expect(key).toStrictEqual([
-      ElizaService.typeName,
-      "name",
-      createMessageKey(SayRequestSchema, { sentence: "someValue" }),
+      "connect-query",
+      {
+        serviceName: ElizaService.typeName,
+        methodName: ElizaService.method.say.name,
+        cardinality: "finite",
+        input: createMessageKey(SayRequestSchema, { sentence: "hi" }),
+      },
     ]);
   });
 
-  it("allows empty inputs", () => {
-    const key = createConnectQueryKey(methodDescriptor);
+  it("allows input: undefined", () => {
+    const key = createConnectQueryKey({
+      method: ElizaService.method.say,
+      input: undefined,
+    });
     expect(key).toStrictEqual([
-      ElizaService.typeName,
-      "name",
-      createMessageKey(methodDescriptor.input, {}),
+      "connect-query",
+      {
+        serviceName: ElizaService.typeName,
+        methodName: ElizaService.method.say.name,
+        cardinality: "finite",
+      },
     ]);
   });
 
-  it("makes a query key with a skipToken", () => {
-    const key = createConnectQueryKey(methodDescriptor, skipToken);
+  it("allows to omit input", () => {
+    const key = createConnectQueryKey({
+      method: ElizaService.method.say,
+    });
     expect(key).toStrictEqual([
-      ElizaService.typeName,
-      "name",
-      createMessageKey(methodDescriptor.input, {}),
+      "connect-query",
+      {
+        serviceName: ElizaService.typeName,
+        methodName: ElizaService.method.say.name,
+        cardinality: "finite",
+      },
     ]);
   });
 
-  it("generates identical keys when input is empty or the default is explicitly sent", () => {
-    const key1 = createConnectQueryKey(methodDescriptor, {});
-    const key2 = createConnectQueryKey(methodDescriptor, { sentence: "" });
-    expect(key1).toStrictEqual(key2);
+  it("skipToken sets input: 'skipped'", () => {
+    const key = createConnectQueryKey({
+      method: ElizaService.method.say,
+      input: skipToken,
+    });
+    expect(key[1].input).toBe("skipped");
   });
 });

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

@@ -12,16 +12,23 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import type { DescMessage, MessageInitShape } from "@bufbuild/protobuf";
+import type {
+  DescMethod,
+  DescService,
+  MessageInitShape,
+} from "@bufbuild/protobuf";
+import type { Transport } from "@connectrpc/connect";
 import type { SkipToken } from "@tanstack/react-query";
-import { skipToken } from "@tanstack/react-query";
 
 import { createMessageKey } from "./message-key.js";
-import type { MethodUnaryDescriptor } from "./method-unary-descriptor.js";
+import { createTransportKey } from "./transport-key.js";
 
 /**
  * TanStack Query requires query keys in order to decide when the query should automatically update.
  *
+ *
+ * TODO
+ *
  * `QueryKey`s in TanStack Query are usually arbitrary, but Connect-Query uses the approach of creating a query key that begins with the least specific information: the service's `typeName`, followed by the method name, and ending with the most specific information to identify a particular request: the input message itself.
  *
  * For example, for a query key might look like this:
@@ -34,51 +41,120 @@ import type { MethodUnaryDescriptor } from "./method-unary-descriptor.js";
  * ]
  */
 export type ConnectQueryKey = [
-  serviceTypeName: string,
-  methodName: string,
-  input: Record<string, unknown>,
+  "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;
+    /**
+     * Whether this is an infinite query, or a regular one.
+     */
+    cardinality?: "infinite" | "finite";
+    /**
+     * A key for the request message, created with createMessageKey(),
+     * or "skipped".
+     */
+    input?: Record<string, unknown> | "skipped";
+  },
 ];
 
+type KeyParams<Desc extends DescMethod | DescService> = Desc extends DescMethod
+  ? {
+      /**
+       * Set `serviceName` and `methodName` in the key.
+       */
+      method: Desc;
+      /**
+       * Set `input` in the key:
+       * - If a SkipToken is provided, `input` is "skipped".
+       * - If an init shape is provided, `input` is set to a message key.
+       * - If omitted or undefined, `input` is not set in the key.
+       */
+      input?: MessageInitShape<Desc["input"]> | SkipToken | undefined;
+      /**
+       * Set `transport` in the key.
+       */
+      transport?: Transport;
+      /**
+       * Set `cardinality` in the key - "finite" by default.
+       */
+      cardinality?: "finite" | "infinite" | "any";
+      /**
+       * If omit the field with this name from the key for infinite queries.
+       */
+      pageParamKey?: keyof MessageInitShape<Desc["input"]>;
+    }
+  : {
+      /**
+       * Set `serviceName` in the key, and omit `methodName`.
+       */
+      service: Desc;
+      /**
+       * Set `transport` in the key.
+       */
+      transport?: Transport;
+      /**
+       * Set `cardinality` in the key - "finite" by default.
+       */
+      cardinality?: "finite" | "infinite" | "any";
+    };
+
 /**
  * TanStack Query requires query keys in order to decide when the query should automatically update.
  *
  * In Connect-Query, much of this is handled automatically by this function.
  *
  * @see ConnectQueryKey for information on the components of Connect-Query's keys.
  */
-export function createConnectQueryKey<
-  I extends DescMessage,
-  O extends DescMessage,
->(
-  schema: Pick<MethodUnaryDescriptor<I, O>, "input" | "parent" | "name">,
-  input?: SkipToken | MessageInitShape<I> | undefined,
+export function createConnectQueryKey<Desc extends DescMethod | DescService>(
+  params: KeyParams<Desc>,
 ): ConnectQueryKey {
-  const key =
-    input === skipToken || input === undefined
-      ? createMessageKey(schema.input, {} as MessageInitShape<DescMessage & I>)
-      : createMessageKey(schema.input, input);
-  return [schema.parent.typeName, schema.name, key];
-}
-
-/**
- * Similar to @see ConnectQueryKey, but for infinite queries.
- */
-export type ConnectInfiniteQueryKey = [
-  serviceTypeName: string,
-  methodName: string,
-  input: Record<string, unknown>,
-  "infinite",
-];
-
-/**
- * Similar to @see createConnectQueryKey, but for infinite queries.
- */
-export function createConnectInfiniteQueryKey<
-  I extends DescMessage,
-  O extends DescMessage,
->(
-  schema: Pick<MethodUnaryDescriptor<I, O>, "input" | "parent" | "name">,
-  input?: SkipToken | MessageInitShape<I> | undefined,
-): ConnectInfiniteQueryKey {
-  return [...createConnectQueryKey(schema, input), "infinite"];
+  const props: ConnectQueryKey[1] =
+    "method" in params
+      ? {
+          serviceName: params.method.parent.typeName,
+          methodName: params.method.name,
+        }
+      : {
+          serviceName: params.service.typeName,
+        };
+  if (params.transport !== undefined) {
+    props.transport = createTransportKey(params.transport);
+  }
+  // eslint-disable-next-line @typescript-eslint/switch-exhaustiveness-check -- "Cases not matched: undefined" 🤷
+  switch (params.cardinality) {
+    case undefined:
+    case "finite":
+      props.cardinality = "finite";
+      break;
+    case "infinite":
+      props.cardinality = "infinite";
+      break;
+    case "any":
+      break;
+  }
+  if ("method" in params && typeof params.input == "symbol") {
+    props.input = "skipped";
+  }
+  if ("method" in params) {
+    if (typeof params.input == "symbol") {
+      props.input = "skipped";
+    } else if (params.input !== undefined) {
+      props.input = createMessageKey(
+        params.method.input,
+        params.input,
+        params.pageParamKey,
+      );
+    }
+  }
+  return ["connect-query", props];
 }

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

@@ -29,8 +29,8 @@ import { skipToken } from "@tanstack/react-query";
 
 import { callUnaryMethod } from "./call-unary-method.js";
 import {
-  type ConnectInfiniteQueryKey,
-  createConnectInfiniteQueryKey,
+  type ConnectQueryKey,
+  createConnectQueryKey,
 } from "./connect-query-key.js";
 import type { MethodUnaryDescriptor } from "./method-unary-descriptor.js";
 import { createStructuralSharing } from "./structural-sharing.js";
@@ -69,7 +69,7 @@ function createUnaryInfiniteQueryFn<
   },
 ): QueryFunction<
   MessageShape<O>,
-  ConnectInfiniteQueryKey,
+  ConnectQueryKey,
   MessageInitShape<I>[ParamKey]
 > {
   return async (context) => {
@@ -108,27 +108,24 @@ export function createInfiniteQueryOptions<
     O,
     ParamKey
   >["getNextPageParam"];
-  queryKey: ConnectInfiniteQueryKey;
+  queryKey: ConnectQueryKey;
   queryFn:
     | QueryFunction<
         MessageShape<O>,
-        ConnectInfiniteQueryKey,
+        ConnectQueryKey,
         MessageInitShape<I>[ParamKey]
       >
     | SkipToken;
   structuralSharing: Exclude<UseQueryOptions["structuralSharing"], undefined>;
   initialPageParam: MessageInitShape<I>[ParamKey];
   queryKeyHashFn: (queryKey: QueryKey) => string;
 } {
-  const queryKey = createConnectInfiniteQueryKey(
-    schema,
-    input === skipToken
-      ? undefined
-      : {
-          ...input,
-          [pageParamKey]: undefined,
-        },
-  );
+  const queryKey = createConnectQueryKey({
+    cardinality: "infinite",
+    method: schema,
+    transport,
+    input,
+  });
   const structuralSharing = createStructuralSharing(schema.output);
   const queryFn =
     input === skipToken