Allow QueryManager to intercept hook functionality

.changeset/curvy-maps-give.md

@@ -0,0 +1,5 @@
+---
+"@apollo/client": patch
+---
+
+Allow Apollo Client instance to intercept hook functionality

.size-limits.json

@@ -1,4 +1,4 @@
 {
-  "dist/apollo-client.min.cjs": 39075,
+  "dist/apollo-client.min.cjs": 39260,
   "import { ApolloClient, InMemoryCache, HttpLink } from \"dist/index.js\" (production)": 32584
 }

config/rollup.config.js

@@ -8,7 +8,7 @@ import cleanup from "rollup-plugin-cleanup";
 const entryPoints = require("./entryPoints");
 const distDir = "./dist";
 
-const removeComments = cleanup({});
+const removeComments = cleanup({ comments: ["some", /#__PURE__/] });
 
 function isExternal(id, parentId, entryPointsAreExternal = true) {
   let posixId = toPosixPath(id);

src/react/hooks/internal/index.ts

@@ -4,3 +4,4 @@ export { useIsomorphicLayoutEffect } from "./useIsomorphicLayoutEffect.js";
 export { useRenderGuard } from "./useRenderGuard.js";
 export { useLazyRef } from "./useLazyRef.js";
 export { __use } from "./__use.js";
+export { makeHookWrappable } from "./wrapHook.js";

src/react/hooks/internal/wrapHook.ts

@@ -0,0 +1,85 @@
+import type {
+  useQuery,
+  useSuspenseQuery,
+  useBackgroundQuery,
+  useReadQuery,
+  useFragment,
+} from "../index.js";
+import type { QueryManager } from "../../../core/QueryManager.js";
+import type { ApolloClient } from "../../../core/ApolloClient.js";
+import type { ObservableQuery } from "../../../core/ObservableQuery.js";
+
+const wrapperSymbol = Symbol.for("apollo.hook.wrappers");
+
+interface WrappableHooks {
+  useQuery: typeof useQuery;
+  useSuspenseQuery: typeof useSuspenseQuery;
+  useBackgroundQuery: typeof useBackgroundQuery;
+  useReadQuery: typeof useReadQuery;
+  useFragment: typeof useFragment;
+}
+
+/**
+ * @internal
+ * Can be used to correctly type the [Symbol.for("apollo.hook.wrappers")] of
+ * a class that extends `ApolloClient`, to override/wrap hook functionality.
+ */
+export type HookWrappers = {
+  [K in keyof WrappableHooks]?: (
+    originalHook: WrappableHooks[K]
+  ) => WrappableHooks[K];
+};
+
+interface QueryManagerWithWrappers<T> extends QueryManager<T> {
+  [wrapperSymbol]?: HookWrappers;
+}
+
+/**
+ * @internal
+ *
+ * Makes an Apollo Client hook "wrappable".
+ * That means that the Apollo Client instance can expose a "wrapper" that will be
+ * used to wrap the original hook implementation with additional logic.
+ * @example
+ * ```tsx
+ * // this is already done in `@apollo/client` for all wrappable hooks (see `WrappableHooks`)
+ * const wrappedUseQuery = makeHookWrappable('useQuery', useQuery, (_, options) => options.client);
+ *
+ * // this is what a library like `@apollo/client-react-streaming` would do
+ * class ApolloClientWithStreaming extends ApolloClient {
+ *   [Symbol.for("apollo.hook.wrappers")] = {
+ *     useQuery: (original) => (query, options) => {
+ *       console.log("useQuery was called with options", options);
+ *       return original(query, options);
+ *     }
+ *   }
+ * }
+ *
+ * // this will now log the options and then call the original `useQuery`
+ * const client = new ApolloClientWithStreaming({ ... });
+ * wrappedUseQuery(query, { client });
+ * ```
+ */
+export function makeHookWrappable<Name extends keyof WrappableHooks>(
+  hookName: Name,
+  useHook: WrappableHooks[Name],
+  getClientFromOptions: (
+    ...args: Parameters<WrappableHooks[Name]>
+  ) => ObservableQuery<any> | ApolloClient<any>
+): WrappableHooks[Name] {
+  return function (this: any) {
+    const args = arguments as unknown as Parameters<WrappableHooks[Name]>;
+    const queryManager = (
+      getClientFromOptions.apply(this, args) as unknown as {
+        // both `ApolloClient` and `ObservableQuery` have a `queryManager` property
+        // but they're both `private`, so we have to cast around for a bit here.
+        queryManager: QueryManagerWithWrappers<any>;
+      }
+    )["queryManager"];
+    const wrappers = queryManager && queryManager[wrapperSymbol];
+    const wrapper = wrappers && wrappers[hookName];
+    const wrappedHook: WrappableHooks[Name] =
+      wrapper ? wrapper(useHook) : useHook;
+    return (wrappedHook as any).apply(this, args);
+  } as any;
+}

src/react/hooks/useBackgroundQuery.ts

@@ -15,7 +15,7 @@ import {
 } from "../internal/index.js";
 import type { CacheKey, QueryReference } from "../internal/index.js";
 import type { BackgroundQueryHookOptions, NoInfer } from "../types/types.js";
-import { __use } from "./internal/index.js";
+import { __use, makeHookWrappable } from "./internal/index.js";
 import { useWatchQueryOptions } from "./useSuspenseQuery.js";
 import type { FetchMoreFunction, RefetchFunction } from "./useSuspenseQuery.js";
 import { canonicalStringify } from "../../cache/index.js";
@@ -246,3 +246,12 @@ export function useBackgroundQuery<
     { fetchMore, refetch },
   ];
 }
+
+const wrapped = /*#__PURE__*/ makeHookWrappable(
+  "useBackgroundQuery",
+  useBackgroundQuery,
+  (_, options) =>
+    useApolloClient(typeof options === "object" ? options.client : undefined)
+);
+// @ts-expect-error Cannot assign to 'useBackgroundQuery' because it is a function.ts(2630)
+useBackgroundQuery = wrapped;

src/react/hooks/useFragment.ts

@@ -14,7 +14,11 @@ import { useApolloClient } from "./useApolloClient.js";
 import { useSyncExternalStore } from "./useSyncExternalStore.js";
 import type { ApolloClient, OperationVariables } from "../../core/index.js";
 import type { NoInfer } from "../types/types.js";
-import { useDeepMemo, useLazyRef } from "./internal/index.js";
+import {
+  useDeepMemo,
+  useLazyRef,
+  makeHookWrappable,
+} from "./internal/index.js";
 
 export interface UseFragmentOptions<TData, TVars>
   extends Omit<
@@ -112,6 +116,14 @@ export function useFragment<TData = any, TVars = OperationVariables>(
   );
 }
 
+const wrapped = /*#__PURE__*/ makeHookWrappable(
+  "useFragment",
+  useFragment,
+  (options) => useApolloClient(options.client)
+);
+// @ts-expect-error Cannot assign to 'useFragment' because it is a function.ts(2630)
+useFragment = wrapped;
+
 function diffToResult<TData>(
   diff: Cache.DiffResult<TData>
 ): UseFragmentResult<TData> {

src/react/hooks/useQuery.ts

@@ -36,6 +36,7 @@ import {
   isNonEmptyArray,
   maybeDeepFreeze,
 } from "../../utilities/index.js";
+import { makeHookWrappable } from "./internal/index.js";
 
 const {
   prototype: { hasOwnProperty },
@@ -90,6 +91,14 @@ export function useQuery<
   );
 }
 
+const wrapped = /*#__PURE__*/ makeHookWrappable(
+  "useQuery",
+  useQuery,
+  (_, options) => useApolloClient(options && options.client)
+);
+// @ts-expect-error Cannot assign to 'useQuery' because it is a function.ts(2630)
+useQuery = wrapped;
+
 export function useInternalState<TData, TVariables extends OperationVariables>(
   client: ApolloClient<any>,
   query: DocumentNode | TypedDocumentNode<TData, TVariables>

src/react/hooks/useReadQuery.ts

@@ -5,7 +5,7 @@ import {
   updateWrappedQueryRef,
 } from "../internal/index.js";
 import type { QueryReference } from "../internal/index.js";
-import { __use } from "./internal/index.js";
+import { __use, makeHookWrappable } from "./internal/index.js";
 import { toApolloError } from "./useSuspenseQuery.js";
 import { useSyncExternalStore } from "./useSyncExternalStore.js";
 import type { ApolloError } from "../../errors/index.js";
@@ -80,3 +80,10 @@ export function useReadQuery<TData>(
     };
   }, [result]);
 }
+const wrapped = /*#__PURE__*/ makeHookWrappable(
+  "useReadQuery",
+  useReadQuery,
+  (ref) => unwrapQueryRef(ref)["observable"]
+);
+// @ts-expect-error Cannot assign to 'useReadQuery' because it is a function.ts(2630)
+useReadQuery = wrapped;

src/react/hooks/useSuspenseQuery.ts

@@ -20,7 +20,7 @@ import type {
   ObservableQueryFields,
   NoInfer,
 } from "../types/types.js";
-import { __use, useDeepMemo } from "./internal/index.js";
+import { __use, useDeepMemo, makeHookWrappable } from "./internal/index.js";
 import { getSuspenseCache } from "../internal/index.js";
 import { canonicalStringify } from "../../cache/index.js";
 import { skipToken } from "./constants.js";
@@ -281,6 +281,15 @@ export function useSuspenseQuery<
   }, [client, fetchMore, refetch, result, subscribeToMore]);
 }
 
+const wrapped = /*#__PURE__*/ makeHookWrappable(
+  "useSuspenseQuery",
+  useSuspenseQuery,
+  (_, options) =>
+    useApolloClient(typeof options === "object" ? options.client : undefined)
+);
+// @ts-expect-error Cannot assign to 'useSuspenseQuery' because it is a function.ts(2630)
+useSuspenseQuery = wrapped;
+
 function validateOptions(options: WatchQueryOptions) {
   const { query, fetchPolicy, returnPartialData } = options;
 

src/react/internal/index.ts

@@ -9,3 +9,4 @@ export {
   wrapQueryRef,
 } from "./cache/QueryReference.js";
 export type { SuspenseCacheOptions } from "./cache/SuspenseCache.js";
+export type { HookWrappers } from "../hooks/internal/wrapHook.js";