Allow QueryManager to intercept hook functionality (#11617)

* Allow Apollo Client instance to intercept hook functionality

* update api extractor

* changeset

* keep PURE comments when building cjs

* shave a few bytes

* Workaround for `useReadableQuery` without wrapping `Provider`.

* update size-limits

* Update src/react/hooks/internal/wrapHook.ts

* put wrappers on `QueryManager` instead

* add `__NO_SIDE_EFFECTS__` annotation

* swap call order

* better tree-shaking approach

* adjust comment

* simplify implementation by just calling `wrapHook`

* adjust comments

* slight type adjustment

.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": 39209,
   "import { ApolloClient, InMemoryCache, HttpLink } from \"dist/index.js\" (production)": 32584
 }

config/rollup.config.js

@@ -8,7 +8,9 @@ import cleanup from "rollup-plugin-cleanup";
 const entryPoints = require("./entryPoints");
 const distDir = "./dist";
 
-const removeComments = cleanup({});
+const removeComments = cleanup({
+  comments: ["some", /#__PURE__/, /#__NO_SIDE_EFFECTS__/],
+});
 
 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 { wrapHook } from "./wrapHook.js";

src/react/hooks/internal/wrapHook.ts

@@ -0,0 +1,88 @@
+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")] property of
+ * `QueryManager`, 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`)
+ * // following this pattern
+ * function useQuery() {
+ *   return wrapHook('useQuery', _useQuery, options.client)(query, options);
+ * }
+ * function _useQuery(query, options) {
+ *   // original implementation
+ * }
+ *
+ * // this is what a library like `@apollo/client-react-streaming` would do
+ * class ApolloClientWithStreaming extends ApolloClient {
+ *   constructor(options) {
+ *     super(options);
+ *     this.queryManager[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({ ... });
+ * useQuery(query, { client });
+ * ```
+ */
+export function wrapHook<Hook extends (...args: any[]) => any>(
+  hookName: keyof WrappableHooks,
+  useHook: Hook,
+  clientOrObsQuery: ObservableQuery<any> | ApolloClient<any>
+): Hook {
+  const queryManager = (
+    clientOrObsQuery 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: undefined | ((wrap: Hook) => Hook) =
+    wrappers && (wrappers[hookName] as any);
+  return wrapper ? wrapper(useHook) : useHook;
+}

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, wrapHook } from "./internal/index.js";
 import { useWatchQueryOptions } from "./useSuspenseQuery.js";
 import type { FetchMoreFunction, RefetchFunction } from "./useSuspenseQuery.js";
 import { canonicalStringify } from "../../cache/index.js";
@@ -182,6 +182,26 @@ export function useBackgroundQuery<
 ): [
   QueryReference<TData, TVariables> | undefined,
   UseBackgroundQueryResult<TData, TVariables>,
+] {
+  return wrapHook(
+    "useBackgroundQuery",
+    _useBackgroundQuery,
+    useApolloClient(typeof options === "object" ? options.client : undefined)
+  )(query, options);
+}
+
+function _useBackgroundQuery<
+  TData = unknown,
+  TVariables extends OperationVariables = OperationVariables,
+>(
+  query: DocumentNode | TypedDocumentNode<TData, TVariables>,
+  options:
+    | (SkipToken &
+        Partial<BackgroundQueryHookOptionsNoInfer<TData, TVariables>>)
+    | BackgroundQueryHookOptionsNoInfer<TData, TVariables>
+): [
+  QueryReference<TData, TVariables> | undefined,
+  UseBackgroundQueryResult<TData, TVariables>,
 ] {
   const client = useApolloClient(options.client);
   const suspenseCache = getSuspenseCache(client);

src/react/hooks/useFragment.ts

@@ -14,7 +14,7 @@ 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, wrapHook } from "./internal/index.js";
 
 export interface UseFragmentOptions<TData, TVars>
   extends Omit<
@@ -53,6 +53,16 @@ export type UseFragmentResult<TData> =
 
 export function useFragment<TData = any, TVars = OperationVariables>(
   options: UseFragmentOptions<TData, TVars>
+): UseFragmentResult<TData> {
+  return wrapHook(
+    "useFragment",
+    _useFragment,
+    useApolloClient(options.client)
+  )(options);
+}
+
+function _useFragment<TData = any, TVars = OperationVariables>(
+  options: UseFragmentOptions<TData, TVars>
 ): UseFragmentResult<TData> {
   const { cache } = useApolloClient(options.client);
 

src/react/hooks/useQuery.ts

@@ -36,6 +36,7 @@ import {
   isNonEmptyArray,
   maybeDeepFreeze,
 } from "../../utilities/index.js";
+import { wrapHook } from "./internal/index.js";
 
 const {
   prototype: { hasOwnProperty },
@@ -85,6 +86,20 @@ export function useQuery<
     NoInfer<TVariables>
   > = Object.create(null)
 ): QueryResult<TData, TVariables> {
+  return wrapHook(
+    "useQuery",
+    _useQuery,
+    useApolloClient(options && options.client)
+  )(query, options);
+}
+
+function _useQuery<
+  TData = any,
+  TVariables extends OperationVariables = OperationVariables,
+>(
+  query: DocumentNode | TypedDocumentNode<TData, TVariables>,
+  options: QueryHookOptions<NoInfer<TData>, NoInfer<TVariables>>
+) {
   return useInternalState(useApolloClient(options.client), query).useQuery(
     options
   );

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, wrapHook } from "./internal/index.js";
 import { toApolloError } from "./useSuspenseQuery.js";
 import { useSyncExternalStore } from "./useSyncExternalStore.js";
 import type { ApolloError } from "../../errors/index.js";
@@ -38,6 +38,16 @@ export interface UseReadQueryResult<TData = unknown> {
 
 export function useReadQuery<TData>(
   queryRef: QueryReference<TData>
+): UseReadQueryResult<TData> {
+  return wrapHook(
+    "useReadQuery",
+    _useReadQuery,
+    unwrapQueryRef(queryRef)["observable"]
+  )(queryRef);
+}
+
+function _useReadQuery<TData>(
+  queryRef: QueryReference<TData>
 ): UseReadQueryResult<TData> {
   const internalQueryRef = React.useMemo(
     () => unwrapQueryRef(queryRef),

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, wrapHook } from "./internal/index.js";
 import { getSuspenseCache } from "../internal/index.js";
 import { canonicalStringify } from "../../cache/index.js";
 import { skipToken } from "./constants.js";
@@ -174,6 +174,22 @@ export function useSuspenseQuery<
   options:
     | (SkipToken & Partial<SuspenseQueryHookOptions<TData, TVariables>>)
     | SuspenseQueryHookOptions<TData, TVariables> = Object.create(null)
+): UseSuspenseQueryResult<TData | undefined, TVariables> {
+  return wrapHook(
+    "useSuspenseQuery",
+    _useSuspenseQuery,
+    useApolloClient(typeof options === "object" ? options.client : undefined)
+  )(query, options);
+}
+
+function _useSuspenseQuery<
+  TData = unknown,
+  TVariables extends OperationVariables = OperationVariables,
+>(
+  query: DocumentNode | TypedDocumentNode<TData, TVariables>,
+  options:
+    | (SkipToken & Partial<SuspenseQueryHookOptions<TData, TVariables>>)
+    | SuspenseQueryHookOptions<TData, TVariables>
 ): UseSuspenseQueryResult<TData | undefined, TVariables> {
   const client = useApolloClient(options.client);
   const suspenseCache = getSuspenseCache(client);

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";