Expose the internal state of usePaginatedQuery through an internal hook (#41552)

GitOrigin-RevId: 68406e0ae95cd472a76bad1690b2358be420ca23

Author: Nicolapps

src/react/use_paginated_query.test.tsx

@@ -1,7 +1,7 @@
 /**
  * @vitest-environment custom-vitest-environment.ts
  */
-import { expect, vi, test, describe, beforeEach } from "vitest";
+import { expect, vi, test, describe, beforeEach, it } from "vitest";
 import { act, renderHook } from "@testing-library/react";
 import React from "react";
 
@@ -26,6 +26,9 @@ import {
   insertAtTop,
   insertAtPosition,
   PaginatedQueryReference,
+  usePaginatedQueryInternal,
+  includePage,
+  page,
 } from "./use_paginated_query.js";
 import { OptimisticLocalStore } from "../browser/index.js";
 
@@ -400,6 +403,54 @@ describe("usePaginatedQuery pages", () => {
       "item4S",
     ]);
   });
+
+  describe("usePaginatedQueryInternal", () => {
+    it("doesn’t expose the page keys by default", () => {
+      const { result } = renderHook(
+        () => usePaginatedQueryInternal(query, {}, { initialNumItems: 1 }),
+        { wrapper },
+      );
+      mockPage(
+        {
+          numItems: 1,
+          cursor: null,
+        },
+        {
+          page: [{ name: "item1" }],
+          continueCursor: "abc",
+          isDone: true,
+        },
+      );
+      const item = result.current.user.results[0];
+      expect(page in item).to.be.false;
+      expect("0" in result.current.internal.state.queries).to.be.true;
+    });
+
+    it("exposes the page keys when asked", () => {
+      const { result } = renderHook(
+        () =>
+          usePaginatedQueryInternal(
+            query,
+            {},
+            { initialNumItems: 1, [includePage]: true },
+          ),
+        { wrapper },
+      );
+      mockPage(
+        {
+          numItems: 1,
+          cursor: null,
+        },
+        {
+          page: [{ name: "item1" }],
+          continueCursor: "abc",
+          isDone: true,
+        },
+      );
+      const item = result.current.user.results[0];
+      expect(item[page]).to.equal("0");
+    });
+  });
 });
 
 describe("PaginatedQueryArgs", () => {

src/react/use_paginated_query.ts

@@ -164,6 +164,32 @@ export function usePaginatedQuery<Query extends PaginatedQueryReference>(
   args: PaginatedQueryArgs<Query> | "skip",
   options: { initialNumItems: number },
 ): UsePaginatedQueryReturnType<Query> {
+  const { user } = usePaginatedQueryInternal(query, args, options);
+  return user;
+}
+
+/** @internal */
+export const includePage = Symbol("includePageKeys");
+
+/** @internal */
+export const page = Symbol("page");
+
+/**
+ * @internal
+ */
+export function usePaginatedQueryInternal<
+  Query extends PaginatedQueryReference,
+>(
+  query: Query,
+  args: PaginatedQueryArgs<Query> | "skip",
+  options: {
+    initialNumItems: number;
+    [includePage]?: boolean;
+  },
+): {
+  user: UsePaginatedQueryReturnType<Query>;
+  internal: { state: UsePaginatedQueryState };
+} {
   if (
     typeof options?.initialNumItems !== "number" ||
     options.initialNumItems < 0
@@ -234,6 +260,7 @@ export function usePaginatedQuery<Query extends PaginatedQueryReference>(
 
   const resultsObject = useQueries(currState.queries);
 
+  const isIncludingPageKeys = options[includePage] ?? false;
   const [results, maybeLastResult]: [
     Value[],
     undefined | PaginationResult<Value>,
@@ -303,7 +330,14 @@ export function usePaginatedQuery<Query extends PaginatedQueryReference>(
         // page and return 'LoadingMore' while the page is splitting.
         return [allItems, undefined];
       }
-      allItems.push(...currResult.page);
+      allItems.push(
+        ...(isIncludingPageKeys
+          ? currResult.page.map((i: any) => ({
+              ...i,
+              [page]: pageKey.toString(),
+            }))
+          : currResult.page),
+      );
     }
     return [allItems, currResult];
   }, [
@@ -313,6 +347,7 @@ export function usePaginatedQuery<Query extends PaginatedQueryReference>(
     options.initialNumItems,
     createInitialState,
     logger,
+    isIncludingPageKeys,
   ]);
 
   const statusObject = useMemo(() => {
@@ -379,8 +414,11 @@ export function usePaginatedQuery<Query extends PaginatedQueryReference>(
   }, [maybeLastResult, currState.nextPageKey]);
 
   return {
-    results,
-    ...statusObject,
+    user: {
+      results,
+      ...statusObject,
+    },
+    internal: { state: currState },
   };
 }