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": 39283,
   "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 { makeHookWrappable } from "./wrapHook.js";

src/react/hooks/internal/wrapHook.ts

@@ -0,0 +1,98 @@
+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);
+ *
+ * // although for tree-shaking purposes, in reality it looks more like
+ * function useQuery() {
+ *   useQuery = makeHookWrappable('useQuery', (_, options) => options.client, _useQuery);
+ *   return useQuery.apply(null, arguments as any);
+ * }
+ * function _useQuery() {
+ *   // 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({ ... });
+ * wrappedUseQuery(query, { client });
+ * ```
+ */
+/*#__NO_SIDE_EFFECTS__*/
+export function makeHookWrappable<Name extends keyof WrappableHooks>(
+  hookName: Name,
+  getClientFromOptions: (
+    ...args: Parameters<WrappableHooks[Name]>
+  ) => ObservableQuery<any> | ApolloClient<any>,
+  useHook: WrappableHooks[Name]
+): 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";
@@ -170,7 +170,19 @@ export function useBackgroundQuery<
   UseBackgroundQueryResult<TData, TVariables>,
 ];
 
-export function useBackgroundQuery<
+export function useBackgroundQuery() {
+  // @ts-expect-error Cannot assign to 'useBackgroundQuery' because it is a function.ts(2630)
+  useBackgroundQuery = makeHookWrappable(
+    "useBackgroundQuery",
+    (_, options) =>
+      useApolloClient(typeof options === "object" ? options.client : undefined),
+    _useBackgroundQuery as any
+  );
+
+  return useBackgroundQuery.apply(null, arguments as any);
+}
+
+function _useBackgroundQuery<
   TData = unknown,
   TVariables extends OperationVariables = OperationVariables,
 >(

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<
@@ -53,6 +57,21 @@ export type UseFragmentResult<TData> =
 
 export function useFragment<TData = any, TVars = OperationVariables>(
   options: UseFragmentOptions<TData, TVars>
+): UseFragmentResult<TData>;
+
+export function useFragment<TData = any, TVars = OperationVariables>() {
+  // @ts-expect-error Cannot assign to 'useFragment' because it is a function.ts(2630)
+  useFragment = makeHookWrappable(
+    "useFragment",
+    (options) => useApolloClient(options.client),
+    _useFragment
+  );
+
+  return useFragment.apply(null, arguments as any);
+}
+
+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 { makeHookWrappable } from "./internal/index.js";
 
 const {
   prototype: { hasOwnProperty },
@@ -78,13 +79,31 @@ const {
 export function useQuery<
   TData = any,
   TVariables extends OperationVariables = OperationVariables,
+>(
+  query: DocumentNode | TypedDocumentNode<TData, TVariables>,
+  options?: QueryHookOptions<NoInfer<TData>, NoInfer<TVariables>>
+): QueryResult<TData, TVariables>;
+
+export function useQuery() {
+  // @ts-expect-error Cannot assign to 'useQuery' because it is a function. ts(2630)
+  useQuery = makeHookWrappable(
+    "useQuery",
+    (_, options) => useApolloClient(options && options.client),
+    _useQuery
+  );
+  return useQuery.apply(null, arguments as any);
+}
+
+function _useQuery<
+  TData = any,
+  TVariables extends OperationVariables = OperationVariables,
 >(
   query: DocumentNode | TypedDocumentNode<TData, TVariables>,
   options: QueryHookOptions<
     NoInfer<TData>,
     NoInfer<TVariables>
   > = Object.create(null)
-): QueryResult<TData, 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, makeHookWrappable } from "./internal/index.js";
 import { toApolloError } from "./useSuspenseQuery.js";
 import { useSyncExternalStore } from "./useSyncExternalStore.js";
 import type { ApolloError } from "../../errors/index.js";
@@ -38,6 +38,20 @@ export interface UseReadQueryResult<TData = unknown> {
 
 export function useReadQuery<TData>(
   queryRef: QueryReference<TData>
+): UseReadQueryResult<TData>;
+
+export function useReadQuery() {
+  // @ts-expect-error Cannot assign to 'useReadQuery' because it is a function.ts(2630)
+  useReadQuery = makeHookWrappable(
+    "useReadQuery",
+    (ref) => unwrapQueryRef(ref)["observable"],
+    _useReadQuery
+  );
+  return useReadQuery.apply(null, arguments as any);
+}
+
+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, makeHookWrappable } from "./internal/index.js";
 import { getSuspenseCache } from "../internal/index.js";
 import { canonicalStringify } from "../../cache/index.js";
 import { skipToken } from "./constants.js";
@@ -166,7 +166,18 @@ export function useSuspenseQuery<
     | SuspenseQueryHookOptions<NoInfer<TData>, NoInfer<TVariables>>
 ): UseSuspenseQueryResult<TData | undefined, TVariables>;
 
-export function useSuspenseQuery<
+export function useSuspenseQuery() {
+  // @ts-expect-error Cannot assign to 'useSuspenseQuery' because it is a function. ts(2630)
+  useSuspenseQuery = makeHookWrappable(
+    "useSuspenseQuery",
+    (_, options) =>
+      useApolloClient(typeof options === "object" ? options.client : undefined),
+    _useSuspenseQuery
+  );
+  return useSuspenseQuery.apply(null, arguments as any);
+}
+
+function _useSuspenseQuery<
   TData = unknown,
   TVariables extends OperationVariables = OperationVariables,
 >(

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