V2: Replace disableQuery with skipToken

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

@@ -13,11 +13,11 @@
 // 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 { disableQuery } from "./utils.js";
 
 describe("makeQueryKey", () => {
   const methodDescriptor = {
@@ -46,8 +46,8 @@ describe("makeQueryKey", () => {
     ]);
   });
 
-  it("makes a query key with a disabled input", () => {
-    const key = createConnectQueryKey(methodDescriptor, disableQuery);
+  it("makes a query key with a skipToken", () => {
+    const key = createConnectQueryKey(methodDescriptor, skipToken);
     expect(key).toStrictEqual([
       ElizaService.typeName,
       "name",

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

@@ -18,10 +18,10 @@ import type {
   MessageShape,
 } from "@bufbuild/protobuf";
 import { create } from "@bufbuild/protobuf";
+import type { SkipToken } from "@tanstack/react-query";
+import { skipToken } from "@tanstack/react-query";
 
 import type { MethodUnaryDescriptor } from "./method-unary-descriptor.js";
-import type { DisableQuery } from "./utils.js";
-import { disableQuery } from "./utils.js";
 
 /**
  * TanStack Query requires query keys in order to decide when the query should automatically update.
@@ -55,12 +55,12 @@ export function createConnectQueryKey<
   O extends DescMessage,
 >(
   schema: Pick<MethodUnaryDescriptor<I, O>, "input" | "parent" | "name">,
-  input?: DisableQuery | MessageInitShape<I> | undefined,
+  input?: SkipToken | MessageInitShape<I> | undefined,
 ): ConnectQueryKey<I> {
   return [
     schema.parent.typeName,
     schema.name,
-    create(schema.input, input === disableQuery || !input ? undefined : input),
+    create(schema.input, input === skipToken || !input ? undefined : input),
   ];
 }
 
@@ -82,7 +82,7 @@ export function createConnectInfiniteQueryKey<
   O extends DescMessage,
 >(
   schema: Pick<MethodUnaryDescriptor<I, O>, "input" | "parent" | "name">,
-  input?: DisableQuery | MessageInitShape<I> | undefined,
+  input?: SkipToken | MessageInitShape<I> | undefined,
 ): ConnectInfiniteQueryKey<I> {
   return [...createConnectQueryKey(schema, input), "infinite"];
 }

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

@@ -21,11 +21,14 @@ import type { CallOptions, ConnectError, Transport } from "@connectrpc/connect";
 import type {
   GetNextPageParamFunction,
   InfiniteData,
+  Optional,
   QueryFunction,
+  SkipToken,
   UseInfiniteQueryOptions,
   UseQueryOptions,
   UseSuspenseInfiniteQueryOptions,
 } from "@tanstack/react-query";
+import { skipToken } from "@tanstack/react-query";
 
 import { callUnaryMethod } from "./call-unary-method.js";
 import {
@@ -34,7 +37,7 @@ import {
 } from "./connect-query-key.js";
 import type { MethodUnaryDescriptor } from "./method-unary-descriptor.js";
 import { createStructuralSharing } from "./structural-sharing.js";
-import { assert, type DisableQuery, disableQuery } from "./utils.js";
+import { assert } from "./utils.js";
 
 /**
  * Options specific to connect-query
@@ -65,7 +68,7 @@ export type CreateInfiniteQueryOptions<
   O extends DescMessage,
   ParamKey extends keyof MessageInitShape<I>,
 > = ConnectInfiniteQueryOptions<I, O, ParamKey> &
-  Omit<
+  Optional<
     UseInfiniteQueryOptions<
       MessageShape<O>,
       ConnectError,
@@ -74,7 +77,7 @@ export type CreateInfiniteQueryOptions<
       ConnectInfiniteQueryKey<I>,
       MessageInitShape<I>[ParamKey]
     >,
-    "getNextPageParam" | "initialPageParam" | "queryFn" | "queryKey"
+    "queryKey"
   >;
 
 /**
@@ -85,7 +88,7 @@ export type CreateSuspenseInfiniteQueryOptions<
   O extends DescMessage,
   ParamKey extends keyof MessageInitShape<I>,
 > = ConnectInfiniteQueryOptions<I, O, ParamKey> &
-  Omit<
+  Optional<
     UseSuspenseInfiniteQueryOptions<
       MessageShape<O>,
       ConnectError,
@@ -94,7 +97,7 @@ export type CreateSuspenseInfiniteQueryOptions<
       ConnectInfiniteQueryKey<I>,
       MessageInitShape<I>[ParamKey]
     >,
-    "getNextPageParam" | "initialPageParam" | "queryFn" | "queryKey"
+    "queryKey"
   >;
 
 function createUnaryInfiniteQueryFn<
@@ -103,7 +106,7 @@ function createUnaryInfiniteQueryFn<
   ParamKey extends keyof MessageInitShape<I>,
 >(
   schema: MethodUnaryDescriptor<I, O>,
-  input: DisableQuery | MessageInitShape<I>,
+  input: MessageInitShape<I>,
   {
     callOptions,
     transport,
@@ -119,7 +122,6 @@ function createUnaryInfiniteQueryFn<
   MessageInitShape<I>[ParamKey]
 > {
   return async (context) => {
-    assert(input !== disableQuery, "Disabled query cannot be fetched");
     assert("pageParam" in context, "pageParam must be part of context");
 
     const inputCombinedWithPageParam = {
@@ -138,9 +140,6 @@ function createUnaryInfiniteQueryFn<
 
 /**
  * Query the method provided. Maps to useInfiniteQuery on tanstack/react-query
- *
- * @param schema
- * @returns
  */
 export function createUseInfiniteQueryOptions<
   I extends DescMessage,
@@ -149,7 +148,7 @@ export function createUseInfiniteQueryOptions<
 >(
   schema: MethodUnaryDescriptor<I, O>,
   input:
-    | DisableQuery
+    | SkipToken
     | (MessageInitShape<I> & Required<Pick<MessageInitShape<I>, ParamKey>>),
   {
     transport,
@@ -164,38 +163,42 @@ export function createUseInfiniteQueryOptions<
     ParamKey
   >["getNextPageParam"];
   queryKey: ConnectInfiniteQueryKey<I>;
-  queryFn: QueryFunction<
-    MessageShape<O>,
-    ConnectInfiniteQueryKey<I>,
-    MessageInitShape<I>[ParamKey]
-  >;
-  structuralSharing?: Exclude<UseQueryOptions["structuralSharing"], undefined>;
+  queryFn:
+    | QueryFunction<
+        MessageShape<O>,
+        ConnectInfiniteQueryKey<I>,
+        MessageInitShape<I>[ParamKey]
+      >
+    | SkipToken;
+  structuralSharing: Exclude<UseQueryOptions["structuralSharing"], undefined>;
   initialPageParam: MessageInitShape<I>[ParamKey];
-  enabled?: false;
 } {
   const queryKey = createConnectInfiniteQueryKey(
     schema,
-    input === disableQuery
+    input === skipToken
       ? undefined
       : {
           ...input,
           [pageParamKey]: undefined,
         },
   );
   const structuralSharing = createStructuralSharing(schema.output);
+  const queryFn =
+    input === skipToken
+      ? skipToken
+      : createUnaryInfiniteQueryFn(schema, input, {
+          transport,
+          callOptions,
+          pageParamKey,
+        });
   return {
     getNextPageParam,
     initialPageParam:
-      input === disableQuery
+      input === skipToken
         ? (undefined as MessageInitShape<I>[ParamKey])
         : (input[pageParamKey] as MessageInitShape<I>[ParamKey]),
     queryKey,
-    queryFn: createUnaryInfiniteQueryFn(schema, input, {
-      transport,
-      callOptions,
-      pageParamKey,
-    }),
+    queryFn,
     structuralSharing,
-    ...(input === disableQuery ? { enabled: false } : undefined),
   };
 }

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

@@ -0,0 +1,46 @@
+// Copyright 2021-2023 The Connect Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import { skipToken } from "@tanstack/react-query";
+import { describe, expect, it } from "vitest";
+
+import { createConnectQueryKey } from "./connect-query-key.js";
+import { createUseQueryOptions } from "./create-use-query-options.js";
+import { ElizaService } from "./gen/eliza_pb.js";
+import { mockEliza } from "./test/test-utils.js";
+
+// TODO: maybe create a helper to take a service and method and generate this.
+const sayMethodDescriptor = ElizaService.method.say;
+
+const mockedElizaTransport = mockEliza();
+
+describe("createUseQueryOptions", () => {
+  it("honors skipToken", () => {
+    const opt = createUseQueryOptions(sayMethodDescriptor, skipToken, {
+      transport: mockedElizaTransport,
+    });
+    expect(opt.queryFn).toBe(skipToken);
+  });
+  it("sets queryKey", () => {
+    const want = createConnectQueryKey(sayMethodDescriptor, { sentence: "hi" });
+    const opt = createUseQueryOptions(
+      sayMethodDescriptor,
+      { sentence: "hi" },
+      {
+        transport: mockedElizaTransport,
+      },
+    );
+    expect(opt.queryKey).toStrictEqual(want);
+  });
+});

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

@@ -19,17 +19,19 @@ import type {
 } from "@bufbuild/protobuf";
 import type { CallOptions, ConnectError, Transport } from "@connectrpc/connect";
 import type {
+  Optional,
   QueryFunction,
+  SkipToken,
   UseQueryOptions,
   UseSuspenseQueryOptions,
 } from "@tanstack/react-query";
+import { skipToken } from "@tanstack/react-query";
 
 import { callUnaryMethod } from "./call-unary-method.js";
 import type { ConnectQueryKey } from "./connect-query-key.js";
 import { createConnectQueryKey } from "./connect-query-key.js";
 import type { MethodUnaryDescriptor } from "./method-unary-descriptor.js";
 import { createStructuralSharing } from "./structural-sharing.js";
-import { assert, type DisableQuery, disableQuery } from "./utils.js";
 
 export interface ConnectQueryOptions {
   /** The transport to be used for the fetching. */
@@ -46,14 +48,14 @@ export type CreateQueryOptions<
   O extends DescMessage,
   SelectOutData = MessageShape<O>,
 > = ConnectQueryOptions &
-  Omit<
+  Optional<
     UseQueryOptions<
       MessageShape<O>,
       ConnectError,
       SelectOutData,
       ConnectQueryKey<I>
     >,
-    "queryFn" | "queryKey"
+    "queryKey"
   >;
 
 /**
@@ -64,19 +66,19 @@ export type CreateSuspenseQueryOptions<
   O extends DescMessage,
   SelectOutData = 0,
 > = ConnectQueryOptions &
-  Omit<
+  Optional<
     UseSuspenseQueryOptions<
       MessageShape<O>,
       ConnectError,
       SelectOutData,
       ConnectQueryKey<I>
     >,
-    "queryFn" | "queryKey"
+    "queryKey"
   >;
 
 function createUnaryQueryFn<I extends DescMessage, O extends DescMessage>(
   schema: MethodUnaryDescriptor<I, O>,
-  input: DisableQuery | MessageInitShape<I> | undefined,
+  input: MessageInitShape<I> | undefined,
   {
     callOptions,
     transport,
@@ -86,7 +88,6 @@ function createUnaryQueryFn<I extends DescMessage, O extends DescMessage>(
   },
 ): QueryFunction<MessageShape<O>, ConnectQueryKey<I>> {
   return async (context) => {
-    assert(input !== disableQuery, "Disabled query cannot be fetched");
     return callUnaryMethod(schema, input, {
       callOptions: {
         ...callOptions,
@@ -105,7 +106,7 @@ export function createUseQueryOptions<
   O extends DescMessage,
 >(
   schema: MethodUnaryDescriptor<I, O>,
-  input: DisableQuery | MessageInitShape<I> | undefined,
+  input: SkipToken | MessageInitShape<I> | undefined,
   {
     transport,
     callOptions,
@@ -114,19 +115,18 @@ export function createUseQueryOptions<
   },
 ): {
   queryKey: ConnectQueryKey<I>;
-  queryFn: QueryFunction<MessageShape<O>, ConnectQueryKey<I>>;
-  structuralSharing?: Exclude<UseQueryOptions["structuralSharing"], undefined>;
-  enabled?: false;
+  queryFn: QueryFunction<MessageShape<O>, ConnectQueryKey<I>> | SkipToken;
+  structuralSharing: Exclude<UseQueryOptions["structuralSharing"], undefined>;
 } {
   const queryKey = createConnectQueryKey(schema, input);
   const structuralSharing = createStructuralSharing(schema.output);
+  const queryFn =
+    input === skipToken
+      ? skipToken
+      : createUnaryQueryFn(schema, input, { transport, callOptions });
   return {
     queryKey,
-    queryFn: createUnaryQueryFn(schema, input, {
-      transport,
-      callOptions,
-    }),
+    queryFn,
     structuralSharing,
-    ...(input === disableQuery ? { enabled: false } : undefined),
   };
 }

packages/connect-query/src/index.ts

@@ -20,7 +20,7 @@ export {
   createConnectQueryKey,
   createConnectInfiniteQueryKey,
 } from "./connect-query-key.js";
-export { disableQuery, createProtobufSafeUpdater } from "./utils.js";
+export { createProtobufSafeUpdater } from "./utils.js";
 export { useTransport, TransportProvider } from "./use-transport.js";
 export type { CreateInfiniteQueryOptions as UseInfiniteQueryOptions } from "./create-use-infinite-query-options.js";
 export {
@@ -32,7 +32,7 @@ export { useQuery, useSuspenseQuery } from "./use-query.js";
 export type { UseMutationOptions } from "./use-mutation.js";
 export { useMutation } from "./use-mutation.js";
 export { defaultOptions } from "./default-options.js";
-export type { DisableQuery, ConnectUpdater } from "./utils.js";
+export type { ConnectUpdater } from "./utils.js";
 export { callUnaryMethod } from "./call-unary-method.js";
 export type { MethodUnaryDescriptor } from "./method-unary-descriptor.js";
 export { createUseInfiniteQueryOptions as createInfiniteQueryOptions } from "./create-use-infinite-query-options.js";