Fix: removeQuery not removing a fetch's reject fn, mistakenly triggering "store reset when query in flight" error when store resets

packages/apollo-client/src/core/QueryManager.ts

@@ -73,10 +73,12 @@ export class QueryManager<TStore> {
   // subscriptions as well
   private queries: Map<string, QueryInfo> = new Map();
 
-  // A set of Promise reject functions for fetchQuery promises that have not
+  // A map of Promise reject functions for fetchQuery promises that have not
   // yet been resolved, used to keep track of in-flight queries so that we can
   // reject them in case a destabilizing event occurs (e.g. Apollo store reset).
-  private fetchQueryRejectFns = new Set<Function>();
+  // The key is in the format of `query:${queryId}` or `fetchRequest:${queryId}`,
+  // depending on where the promise's rejection function was created from.
+  private fetchQueryRejectFns = new Map<string, Function>();
 
   // A map going from the name of a query to an observer issued for it by watchQuery. This is
   // generally used to refetches for refetchQueries and to update mutation results through
@@ -700,8 +702,9 @@ export class QueryManager<TStore> {
     }
 
     return new Promise<ApolloQueryResult<T>>((resolve, reject) => {
-      this.fetchQueryRejectFns.add(reject);
-      this.watchQuery<T>(options, false)
+      const watchedQuery = this.watchQuery<T>(options, false);
+      this.fetchQueryRejectFns.set(`query:${watchedQuery.queryId}`, reject);
+      watchedQuery
         .result()
         .then(resolve, reject)
         // Since neither resolve nor reject throw or return a value, this .then
@@ -710,7 +713,9 @@ export class QueryManager<TStore> {
         // since resolve and reject are mutually idempotent. In fact, it would
         // not be incorrect to let reject functions accumulate over time; it's
         // just a waste of memory.
-        .then(() => this.fetchQueryRejectFns.delete(reject));
+        .then(() =>
+          this.fetchQueryRejectFns.delete(`query:${watchedQuery.queryId}`),
+        );
     });
   }
 
@@ -960,6 +965,12 @@ export class QueryManager<TStore> {
   public removeQuery(queryId: string) {
     const { subscriptions } = this.getQuery(queryId);
     // teardown all links
+    // Both `QueryManager.fetchRequest` and `QueryManager.query` create separate promises
+    // that each add their reject functions to fetchQueryRejectFns.
+    // A query created with `QueryManager.query()` could trigger a `QueryManager.fetchRequest`.
+    // The same queryId could have two rejection fns for two promises
+    this.fetchQueryRejectFns.delete(`query:${queryId}`);
+    this.fetchQueryRejectFns.delete(`fetchRequest:${queryId}`);
     subscriptions.forEach(x => x.unsubscribe());
     this.queries.delete(queryId);
   }
@@ -1092,12 +1103,8 @@ export class QueryManager<TStore> {
     let resultFromStore: any;
     let errorsFromStore: any;
 
-    let rejectFetchPromise: (reason?: any) => void;
-
     return new Promise<ApolloQueryResult<T>>((resolve, reject) => {
-      // Need to assign the reject function to the rejectFetchPromise variable
-      // in the outer scope so that we can refer to it in the .catch handler.
-      this.fetchQueryRejectFns.add((rejectFetchPromise = reject));
+      this.fetchQueryRejectFns.set(`fetchRequest:${queryId}`, reject);
 
       const subscription = execute(this.deduplicator, operation).subscribe({
         next: (result: ExecutionResult) => {
@@ -1164,7 +1171,7 @@ export class QueryManager<TStore> {
           }
         },
         error: (error: ApolloError) => {
-          this.fetchQueryRejectFns.delete(reject);
+          this.fetchQueryRejectFns.delete(`fetchRequest:${queryId}`);
 
           this.setQuery(queryId, ({ subscriptions }) => ({
             subscriptions: subscriptions.filter(x => x !== subscription),
@@ -1173,7 +1180,7 @@ export class QueryManager<TStore> {
           reject(error);
         },
         complete: () => {
-          this.fetchQueryRejectFns.delete(reject);
+          this.fetchQueryRejectFns.delete(`fetchRequest:${queryId}`);
 
           this.setQuery(queryId, ({ subscriptions }) => ({
             subscriptions: subscriptions.filter(x => x !== subscription),
@@ -1193,7 +1200,7 @@ export class QueryManager<TStore> {
         subscriptions: subscriptions.concat([subscription]),
       }));
     }).catch(error => {
-      this.fetchQueryRejectFns.delete(rejectFetchPromise);
+      this.fetchQueryRejectFns.delete(`fetchRequest:${queryId}`);
       throw error;
     });
   }

packages/apollo-client/src/core/__tests__/QueryManager/index.ts

@@ -3229,6 +3229,42 @@ describe('QueryManager', () => {
       );
     });
 
+    it('should not error on a stopped query()', done => {
+      let queryManager: QueryManager<NormalizedCacheObject>;
+      const query = gql`
+        query {
+          author {
+            firstName
+            lastName
+          }
+        }
+      `;
+
+      const data = {
+        author: {
+          firstName: 'John',
+          lastName: 'Smith',
+        },
+      };
+
+      const link = new ApolloLink(
+        () =>
+          new Observable(observer => {
+            observer.next({ data });
+          }),
+      );
+
+      queryManager = createQueryManager({ link });
+
+      const queryId = '1';
+      queryManager
+        .fetchQuery(queryId, { query })
+        .catch(e => done.fail('Exception thrown for stopped query'));
+
+      queryManager.removeQuery(queryId);
+      queryManager.resetStore().then(() => done());
+    });
+
     it('should throw an error on an inflight fetch query if the store is reset', done => {
       const query = gql`
         query {