Use createMessageKey to create query keys and remove default options

Signed-off-by: Gilbert Wong <gilwong3105@gmail.com>

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

@@ -21,7 +21,6 @@ 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 type { SayRequestSchema } from "./gen/eliza_pb.js";
 import { ElizaService } from "./gen/eliza_pb.js";
 import { mockEliza, wrapper } from "./test/test-utils.js";
 
@@ -38,9 +37,7 @@ describe("callUnaryMethod", () => {
               queryFn: async ({
                 queryKey,
                 signal,
-              }: QueryFunctionContext<
-                ConnectQueryKey<typeof SayRequestSchema>
-              >) => {
+              }: QueryFunctionContext<ConnectQueryKey>) => {
                 const transport = mockEliza({
                   sentence: "Response 1",
                 });

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

@@ -12,12 +12,12 @@
 // 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 = {
@@ -33,7 +33,7 @@ describe("makeQueryKey", () => {
     expect(key).toStrictEqual([
       ElizaService.typeName,
       "name",
-      create(SayRequestSchema, { sentence: "someValue" }),
+      createMessageKey(SayRequestSchema, { sentence: "someValue" }),
     ]);
   });
 
@@ -42,7 +42,7 @@ describe("makeQueryKey", () => {
     expect(key).toStrictEqual([
       ElizaService.typeName,
       "name",
-      create(methodDescriptor.input),
+      createMessageKey(methodDescriptor.input, {}),
     ]);
   });
 
@@ -51,7 +51,7 @@ describe("makeQueryKey", () => {
     expect(key).toStrictEqual([
       ElizaService.typeName,
       "name",
-      create(methodDescriptor.input),
+      createMessageKey(methodDescriptor.input, {}),
     ]);
   });
 

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

@@ -12,15 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import type {
-  DescMessage,
-  MessageInitShape,
-  MessageShape,
-} from "@bufbuild/protobuf";
-import { create } from "@bufbuild/protobuf";
+import type { DescMessage, MessageInitShape } from "@bufbuild/protobuf";
 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";
 
 /**
@@ -37,10 +33,10 @@ import type { MethodUnaryDescriptor } from "./method-unary-descriptor.js";
  *   { sentence: "hello there" },
  * ]
  */
-export type ConnectQueryKey<I extends DescMessage> = [
+export type ConnectQueryKey = [
   serviceTypeName: string,
   methodName: string,
-  input: MessageShape<I>,
+  input: Record<string, unknown>,
 ];
 
 /**
@@ -56,21 +52,21 @@ export function createConnectQueryKey<
 >(
   schema: Pick<MethodUnaryDescriptor<I, O>, "input" | "parent" | "name">,
   input?: SkipToken | MessageInitShape<I> | undefined,
-): ConnectQueryKey<I> {
-  return [
-    schema.parent.typeName,
-    schema.name,
-    create(schema.input, input === skipToken || !input ? undefined : input),
-  ];
+): 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<I extends DescMessage> = [
+export type ConnectInfiniteQueryKey = [
   serviceTypeName: string,
   methodName: string,
-  input: MessageShape<I>,
+  input: Record<string, unknown>,
   "infinite",
 ];
 
@@ -83,6 +79,6 @@ export function createConnectInfiniteQueryKey<
 >(
   schema: Pick<MethodUnaryDescriptor<I, O>, "input" | "parent" | "name">,
   input?: SkipToken | MessageInitShape<I> | undefined,
-): ConnectInfiniteQueryKey<I> {
+): ConnectInfiniteQueryKey {
   return [...createConnectQueryKey(schema, input), "infinite"];
 }

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

@@ -21,6 +21,7 @@ import type { Transport } from "@connectrpc/connect";
 import type {
   GetNextPageParamFunction,
   QueryFunction,
+  QueryKey,
   SkipToken,
   UseQueryOptions,
 } from "@tanstack/react-query";
@@ -68,7 +69,7 @@ function createUnaryInfiniteQueryFn<
   },
 ): QueryFunction<
   MessageShape<O>,
-  ConnectInfiniteQueryKey<I>,
+  ConnectInfiniteQueryKey,
   MessageInitShape<I>[ParamKey]
 > {
   return async (context) => {
@@ -107,16 +108,17 @@ export function createInfiniteQueryOptions<
     O,
     ParamKey
   >["getNextPageParam"];
-  queryKey: ConnectInfiniteQueryKey<I>;
+  queryKey: ConnectInfiniteQueryKey;
   queryFn:
     | QueryFunction<
         MessageShape<O>,
-        ConnectInfiniteQueryKey<I>,
+        ConnectInfiniteQueryKey,
         MessageInitShape<I>[ParamKey]
       >
     | SkipToken;
   structuralSharing: Exclude<UseQueryOptions["structuralSharing"], undefined>;
   initialPageParam: MessageInitShape<I>[ParamKey];
+  queryKeyHashFn: (queryKey: QueryKey) => string;
 } {
   const queryKey = createConnectInfiniteQueryKey(
     schema,
@@ -143,5 +145,6 @@ export function createInfiniteQueryOptions<
     queryKey,
     queryFn,
     structuralSharing,
+    queryKeyHashFn: JSON.stringify,
   };
 }

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

@@ -20,6 +20,7 @@ import type {
 import type { Transport } from "@connectrpc/connect";
 import type {
   QueryFunction,
+  QueryKey,
   SkipToken,
   UseQueryOptions as TanStackUseQueryOptions,
 } from "@tanstack/react-query";
@@ -35,7 +36,7 @@ function createUnaryQueryFn<I extends DescMessage, O extends DescMessage>(
   transport: Transport,
   schema: MethodUnaryDescriptor<I, O>,
   input: MessageInitShape<I> | undefined,
-): QueryFunction<MessageShape<O>, ConnectQueryKey<I>> {
+): QueryFunction<MessageShape<O>, ConnectQueryKey> {
   return async (context) => {
     return callUnaryMethod(transport, schema, input, {
       signal: context.signal,
@@ -58,12 +59,13 @@ export function createQueryOptions<
     transport: Transport;
   },
 ): {
-  queryKey: ConnectQueryKey<I>;
-  queryFn: QueryFunction<MessageShape<O>, ConnectQueryKey<I>> | SkipToken;
+  queryKey: ConnectQueryKey;
+  queryFn: QueryFunction<MessageShape<O>, ConnectQueryKey> | SkipToken;
   structuralSharing: Exclude<
     TanStackUseQueryOptions["structuralSharing"],
     undefined
   >;
+  queryKeyHashFn: (queryKey: QueryKey) => string;
 } {
   const queryKey = createConnectQueryKey(schema, input);
   const structuralSharing = createStructuralSharing(schema.output);
@@ -75,5 +77,6 @@ export function createQueryOptions<
     queryKey,
     queryFn,
     structuralSharing,
+    queryKeyHashFn: JSON.stringify,
   };
 }

packages/connect-query/src/default-options.ts

@@ -13,7 +13,6 @@
 // limitations under the License.
 
 import type { DefaultOptions } from "@tanstack/react-query";
-import stableHash from "stable-hash";
 
 /**
  * These default options are required for proper query key hashing.
@@ -25,9 +24,5 @@ import stableHash from "stable-hash";
  * when TanStack Query tries to serialize the value.
  */
 export const defaultOptions = {
-  queries: {
-    // TODO remove once we use createMessageKey()
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any -- stableHash is a function but doesn't seem to satisfy the type. Likely a node16 compat problem.
-    queryKeyHashFn: stableHash as any as (object: any) => string,
-  },
+  queries: {},
 } satisfies DefaultOptions;

packages/connect-query/src/use-infinite-query.ts

@@ -50,7 +50,7 @@ export type UseInfiniteQueryOptions<
     ConnectError,
     InfiniteData<MessageShape<O>>,
     MessageShape<O>,
-    ConnectInfiniteQueryKey<I>,
+    ConnectInfiniteQueryKey,
     MessageInitShape<I>[ParamKey]
   >,
   "getNextPageParam" | "initialPageParam" | "queryFn" | "queryKey"
@@ -104,7 +104,7 @@ export type UseSuspenseInfiniteQueryOptions<
     ConnectError,
     InfiniteData<MessageShape<O>>,
     MessageShape<O>,
-    ConnectInfiniteQueryKey<I>,
+    ConnectInfiniteQueryKey,
     MessageInitShape<I>[ParamKey]
   >,
   "getNextPageParam" | "initialPageParam" | "queryFn" | "queryKey"

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

@@ -39,15 +39,14 @@ import { useTransport } from "./use-transport.js";
  * Options for useQuery
  */
 export type UseQueryOptions<
-  I extends DescMessage,
   O extends DescMessage,
   SelectOutData = MessageShape<O>,
 > = Omit<
   TanStackUseQueryOptions<
     MessageShape<O>,
     ConnectError,
     SelectOutData,
-    ConnectQueryKey<I>
+    ConnectQueryKey
   >,
   "queryFn" | "queryKey"
 > & {
@@ -65,7 +64,7 @@ export function useQuery<
 >(
   schema: MethodUnaryDescriptor<I, O>,
   input?: SkipToken | MessageInitShape<I>,
-  { transport, ...queryOptions }: UseQueryOptions<I, O, SelectOutData> = {},
+  { transport, ...queryOptions }: UseQueryOptions<O, SelectOutData> = {},
 ): UseQueryResult<SelectOutData, ConnectError> {
   const transportFromCtx = useTransport();
   const baseOptions = createQueryOptions(schema, input, {
@@ -81,15 +80,14 @@ export function useQuery<
  * Options for useSuspenseQuery
  */
 export type UseSuspenseQueryOptions<
-  I extends DescMessage,
   O extends DescMessage,
   SelectOutData = 0,
 > = Omit<
   TanStackUseSuspenseQueryOptions<
     MessageShape<O>,
     ConnectError,
     SelectOutData,
-    ConnectQueryKey<I>
+    ConnectQueryKey
   >,
   "queryFn" | "queryKey"
 > & {
@@ -110,7 +108,7 @@ export function useSuspenseQuery<
   {
     transport,
     ...queryOptions
-  }: UseSuspenseQueryOptions<I, O, SelectOutData> = {},
+  }: UseSuspenseQueryOptions<O, SelectOutData> = {},
 ): UseSuspenseQueryResult<SelectOutData, ConnectError> {
   const transportFromCtx = useTransport();
   const baseOptions = createQueryOptions(schema, input, {