Add utility helper to update infinite data

README.md

@@ -25,6 +25,7 @@ Connect-Query is an wrapper around [TanStack Query](https://tanstack.com/query)
   - [`createConnectQueryKey`](#createconnectquerykey)
   - [`callUnaryMethod`](#callunarymethod)
   - [`createProtobufSafeUpdater`](#createprotobufsafeupdater)
+  - [`createProtobufSafeInfiniteUpdater`](#createprotobufsafeinfiniteupdater)
   - [`createQueryOptions`](#createqueryoptions)
   - [`createInfiniteQueryOptions`](#createinfinitequeryoptions)
   - [`addStaticKeyToTransport`](#addstatickeytotransport)
@@ -387,6 +388,10 @@ queryClient.setQueryData(
 
 ```
 
+### `createProtobufSafeInfiniteUpdater`
+
+Creates a typesafe updater for infinite queries. Identical to `createProtobufSafeUpdater` except the data received must be of the shape expected for infinite data.
+
 ### `createQueryOptions`
 
 ```ts

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

@@ -18,6 +18,7 @@ import { describe, expect, it } from "vitest";
 
 import {
   assert,
+  createProtobufSafeInfiniteUpdater,
   createProtobufSafeUpdater,
   isAbortController,
 } from "./utils.js";
@@ -184,3 +185,84 @@ describe("createProtobufSafeUpdater", () => {
     });
   });
 });
+
+describe("createProtobufSafeInfiniteUpdater", () => {
+  describe("with update message", () => {
+    const schema = { output: Proto2MessageSchema };
+    const update = {
+      pageParams: [],
+      pages: [
+        {
+          int32Field: 999,
+        },
+      ],
+    };
+    const safeUpdater = createProtobufSafeInfiniteUpdater(schema, update);
+    it("returns update message for previous value undefined", () => {
+      const next = safeUpdater(undefined);
+      expect(next?.pages[0].$typeName).toBe("test.Proto2Message");
+    });
+  });
+
+  describe("with update message init", () => {
+    const schema = { output: Proto2MessageSchema };
+    const update = {
+      pageParams: [],
+      pages: [
+        {
+          int32Field: 999,
+        },
+      ],
+    };
+    const safeUpdater = createProtobufSafeInfiniteUpdater(schema, update);
+    it("returns update message for previous value undefined", () => {
+      const next = safeUpdater(undefined);
+      expect(next?.pages[0].int32Field).toBe(999);
+    });
+    it("returns update message for previous value", () => {
+      const prev = {
+        pageParams: [],
+        pages: [
+          create(Proto2MessageSchema, {
+            int32Field: 123,
+          }),
+        ],
+      };
+      const next = safeUpdater(prev);
+      expect(next?.pages[0].$typeName).toBe(Proto2MessageSchema.typeName);
+      expect(next?.pages[0].int32Field).toBe(999);
+    });
+  });
+
+  describe("with updater function", () => {
+    const schema = { output: Proto2MessageSchema };
+    const safeUpdater = createProtobufSafeInfiniteUpdater(schema, (prev) => {
+      if (prev === undefined) {
+        return undefined;
+      }
+      return {
+        pageParams: [...prev.pageParams, 44],
+        pages: [
+          ...prev.pages,
+          {
+            int32Field: 33,
+            stringField: "whatever",
+          },
+        ],
+      };
+    });
+    it("accepts undefined", () => {
+      const next = safeUpdater(undefined);
+      expect(next).toBeUndefined();
+    });
+    it("can add a new page", () => {
+      const next = safeUpdater({
+        pageParams: [],
+        pages: [],
+      });
+
+      expect(next?.pageParams).toHaveLength(1);
+      expect(next?.pages).toHaveLength(1);
+    });
+  });
+});

packages/connect-query-core/src/utils.ts

@@ -19,6 +19,7 @@ import type {
   MessageShape,
 } from "@bufbuild/protobuf";
 import { create, isMessage } from "@bufbuild/protobuf";
+import type { InfiniteData } from "@tanstack/query-core";
 
 /**
  * Throws an error with the provided message when the condition is `false`
@@ -58,6 +59,16 @@ export type ConnectUpdater<O extends DescMessage> =
   | undefined
   | ((prev?: MessageShape<O>) => MessageShape<O> | undefined);
 
+/**
+ * @see `Updater` from `@tanstack/react-query`
+ */
+export type ConnectInfiniteUpdater<O extends DescMessage> =
+  | InfiniteData<MessageInitShape<O>>
+  | undefined
+  | ((
+      prev?: InfiniteData<MessageShape<O>>,
+    ) => InfiniteData<MessageInitShape<O>> | undefined);
+
 /**
  * This helper makes sure that the type for the original response message is returned.
  */
@@ -78,3 +89,30 @@ export const createProtobufSafeUpdater =
     }
     return updater(prev);
   };
+
+export const createProtobufSafeInfiniteUpdater =
+  <O extends DescMessage>(
+    schema: Pick<DescMethodUnary<never, O>, "output">,
+    updater: ConnectInfiniteUpdater<O>,
+  ) =>
+  (
+    prev?: InfiniteData<MessageShape<O>>,
+  ): InfiniteData<MessageShape<O>> | undefined => {
+    if (typeof updater !== "function") {
+      if (updater === undefined) {
+        return undefined;
+      }
+      return {
+        pageParams: updater.pageParams,
+        pages: updater.pages.map((i) => create(schema.output, i)),
+      };
+    }
+    const result = updater(prev);
+    if (result === undefined) {
+      return undefined;
+    }
+    return {
+      pageParams: result.pageParams,
+      pages: result.pages.map((i) => create(schema.output, i)),
+    };
+  };