move executeQuery to useLazyQuery

src/react/hooks/useLazyQuery.ts

@@ -2,15 +2,25 @@ import type { DocumentNode } from "graphql";
 import type { TypedDocumentNode } from "@graphql-typed-document-node/core";
 import * as React from "rehackt";
 
-import type { OperationVariables } from "../../core/index.js";
+import type {
+  ApolloQueryResult,
+  OperationVariables,
+} from "../../core/index.js";
 import { mergeOptions } from "../../utilities/index.js";
 import type {
   LazyQueryHookExecOptions,
   LazyQueryHookOptions,
   LazyQueryResultTuple,
   NoInfer,
+  QueryHookOptions,
+  QueryResult,
 } from "../types/types.js";
-import { useInternalState, useQueryWithInternalState } from "./useQuery.js";
+import type { InternalState } from "./useQuery.js";
+import {
+  createWatchQueryOptions,
+  useInternalState,
+  useQueryWithInternalState,
+} from "./useQuery.js";
 import { useApolloClient } from "./useApolloClient.js";
 
 // The following methods, when called will execute the query, regardless of
@@ -94,7 +104,8 @@ export function useLazyQuery<
     useQueryResult.observable.options.initialFetchPolicy ||
     internalState.getDefaultFetchPolicy();
 
-  const { forceUpdateState, obsQueryFields } = internalState;
+  const { obsQueryFields } = internalState;
+  const forceUpdateState = React.useReducer((tick) => tick + 1, 0)[1];
   // We use useMemo here to make sure the eager methods have a stable identity.
   const eagerMethods = React.useMemo(() => {
     const eagerMethods: Record<string, any> = {};
@@ -141,18 +152,76 @@ export function useLazyQuery<
         ...execOptionsRef.current,
       });
 
-      const promise = internalState
-        .executeQuery({ ...options, skip: false }, false)
-        .then((queryResult) => Object.assign(queryResult, eagerMethods));
+      const promise = executeQuery(
+        { ...options, skip: false },
+        false,
+        internalState,
+        forceUpdateState
+      ).then((queryResult) => Object.assign(queryResult, eagerMethods));
 
       // Because the return value of `useLazyQuery` is usually floated, we need
       // to catch the promise to prevent unhandled rejections.
       promise.catch(() => {});
 
       return promise;
     },
-    [eagerMethods, initialFetchPolicy, internalState]
+    [eagerMethods, forceUpdateState, initialFetchPolicy, internalState]
   );
 
   return [execute, result];
 }
+
+function executeQuery<TData, TVariables extends OperationVariables>(
+  options: QueryHookOptions<TData, TVariables> & {
+    query?: DocumentNode;
+  },
+  hasRenderPromises: boolean,
+  internalState: InternalState<TData, TVariables>,
+  forceUpdate: () => void
+) {
+  if (options.query) {
+    internalState.query = options.query;
+  }
+
+  internalState.watchQueryOptions = createWatchQueryOptions(
+    (internalState.queryHookOptions = options),
+    internalState,
+    hasRenderPromises
+  );
+
+  const concast = internalState.observable.reobserveAsConcast(
+    internalState.getObsQueryOptions()
+  );
+
+  // Make sure getCurrentResult returns a fresh ApolloQueryResult<TData>,
+  // but save the current data as this.previousData, just like setResult
+  // usually does.
+  internalState.previousData =
+    internalState.result?.data || internalState.previousData;
+  internalState.result = void 0;
+  forceUpdate();
+
+  return new Promise<QueryResult<TData, TVariables>>((resolve) => {
+    let result: ApolloQueryResult<TData>;
+
+    // Subscribe to the concast independently of the ObservableQuery in case
+    // the component gets unmounted before the promise resolves. This prevents
+    // the concast from terminating early and resolving with `undefined` when
+    // there are no more subscribers for the concast.
+    concast.subscribe({
+      next: (value) => {
+        result = value;
+      },
+      error: () => {
+        resolve(
+          internalState.toQueryResult(
+            internalState.observable.getCurrentResult()
+          )
+        );
+      },
+      complete: () => {
+        resolve(internalState.toQueryResult(result));
+      },
+    });
+  });
+}

src/react/hooks/useQuery.ts

@@ -238,8 +238,6 @@ export function useQueryWithInternalState<
           return () => {};
         }
 
-        internalState.forceUpdate = handleStoreChange;
-
         const onNext = () => {
           const previousResult = internalState.result;
           // We use `getCurrentResult()` instead of the onNext argument because
@@ -256,7 +254,7 @@ export function useQueryWithInternalState<
             return;
           }
 
-          internalState.setResult(result);
+          internalState.setResult(result, handleStoreChange);
         };
 
         const onError = (error: Error) => {
@@ -274,12 +272,15 @@ export function useQueryWithInternalState<
             (previousResult && previousResult.loading) ||
             !equal(error, previousResult.error)
           ) {
-            internalState.setResult({
-              data: (previousResult && previousResult.data) as TData,
-              error: error as ApolloError,
-              loading: false,
-              networkStatus: NetworkStatus.error,
-            });
+            internalState.setResult(
+              {
+                data: (previousResult && previousResult.data) as TData,
+                error: error as ApolloError,
+                loading: false,
+                networkStatus: NetworkStatus.error,
+              },
+              handleStoreChange
+            );
           }
         };
 
@@ -291,7 +292,6 @@ export function useQueryWithInternalState<
         // happen in very fast succession.
         return () => {
           setTimeout(() => subscription.unsubscribe());
-          internalState.forceUpdate = () => internalState.forceUpdateState();
         };
       },
       // eslint-disable-next-line react-compiler/react-compiler
@@ -323,17 +323,8 @@ export function useInternalState<TData, TVariables extends OperationVariables>(
   client: ApolloClient<any>,
   query: DocumentNode | TypedDocumentNode<TData, TVariables>
 ): InternalState<TData, TVariables> {
-  // By default, InternalState.prototype.forceUpdate is an empty function, but
-  // we replace it here (before anyone has had a chance to see this state yet)
-  // with a function that unconditionally forces an update, using the latest
-  // setTick function. Updating this state by calling state.forceUpdate or the
-  // uSES notification callback are the only way we trigger React component updates.
-  const forceUpdateState = React.useReducer((tick) => tick + 1, 0)[1];
-
   function createInternalState(previous?: InternalState<TData, TVariables>) {
-    return Object.assign(new InternalState(client, query, previous), {
-      forceUpdateState,
-    });
+    return new InternalState(client, query, previous);
   }
 
   let [state, updateState] = React.useState(createInternalState);
@@ -351,7 +342,7 @@ export function useInternalState<TData, TVariables extends OperationVariables>(
   return state;
 }
 // A function to massage options before passing them to ObservableQuery.
-function createWatchQueryOptions<
+export function createWatchQueryOptions<
   TData = any,
   TVariables extends OperationVariables = OperationVariables,
 >(
@@ -412,10 +403,11 @@ function createWatchQueryOptions<
   return watchQueryOptions;
 }
 
+export { type InternalState };
 class InternalState<TData, TVariables extends OperationVariables> {
   constructor(
     public readonly client: ReturnType<typeof useApolloClient>,
-    public readonly query: DocumentNode | TypedDocumentNode<TData, TVariables>,
+    public query: DocumentNode | TypedDocumentNode<TData, TVariables>,
     previous?: InternalState<TData, TVariables>
   ) {
     verifyDocumentType(query, DocumentType.Query);
@@ -429,74 +421,6 @@ class InternalState<TData, TVariables extends OperationVariables> {
     }
   }
 
-  /**
-   * Forces an update using local component state.
-   * As this is not batched with `useSyncExternalStore` updates,
-   * this is only used as a fallback if the `useSyncExternalStore` "force update"
-   * method is not registered at the moment.
-   * See https://github.com/facebook/react/issues/25191
-   *  */
-  forceUpdateState() {
-    // Replaced (in useInternalState) with a method that triggers an update.
-    invariant.warn(
-      "Calling default no-op implementation of InternalState#forceUpdate"
-    );
-  }
-
-  /**
-   * Will be overwritten by the `useSyncExternalStore` "force update" method
-   * whenever it is available and reset to `forceUpdateState` when it isn't.
-   */
-  forceUpdate = () => this.forceUpdateState();
-
-  executeQuery(
-    options: QueryHookOptions<TData, TVariables> & {
-      query?: DocumentNode;
-    },
-    hasRenderPromises: boolean
-  ) {
-    if (options.query) {
-      Object.assign(this, { query: options.query });
-    }
-
-    this.watchQueryOptions = createWatchQueryOptions(
-      (this.queryHookOptions = options),
-      this,
-      hasRenderPromises
-    );
-
-    const concast = this.observable.reobserveAsConcast(
-      this.getObsQueryOptions()
-    );
-
-    // Make sure getCurrentResult returns a fresh ApolloQueryResult<TData>,
-    // but save the current data as this.previousData, just like setResult
-    // usually does.
-    this.previousData = this.result?.data || this.previousData;
-    this.result = void 0;
-    this.forceUpdate();
-
-    return new Promise<QueryResult<TData, TVariables>>((resolve) => {
-      let result: ApolloQueryResult<TData>;
-
-      // Subscribe to the concast independently of the ObservableQuery in case
-      // the component gets unmounted before the promise resolves. This prevents
-      // the concast from terminating early and resolving with `undefined` when
-      // there are no more subscribers for the concast.
-      concast.subscribe({
-        next: (value) => {
-          result = value;
-        },
-        error: () => {
-          resolve(this.toQueryResult(this.observable.getCurrentResult()));
-        },
-        complete: () => {
-          resolve(this.toQueryResult(result));
-        },
-      });
-    });
-  }
-
   public queryHookOptions!: QueryHookOptions<TData, TVariables>;
   public watchQueryOptions!: WatchQueryOptions<TVariables, TData>;
 
@@ -569,15 +493,18 @@ class InternalState<TData, TVariables extends OperationVariables> {
   public result: undefined | ApolloQueryResult<TData>;
   public previousData: undefined | TData;
 
-  public setResult(nextResult: ApolloQueryResult<TData>) {
+  public setResult(
+    nextResult: ApolloQueryResult<TData>,
+    forceUpdate: () => void
+  ) {
     const previousResult = this.result;
     if (previousResult && previousResult.data) {
       this.previousData = previousResult.data;
     }
     this.result = nextResult;
     // Calling state.setResult always triggers an update, though some call sites
     // perform additional equality checks before committing to an update.
-    this.forceUpdate();
+    forceUpdate();
     this.handleErrorOrCompleted(nextResult, previousResult);
   }