Fix onNext callback in requestSubscription

Reviewed By: josephsavona Differential Revision: D27173202 fbshipit-source-id: 456ca308339ea15ff44daf5aedcd83ebcb0927b6
facebook · Mar 23, 2021 · 579b8f2 · 579b8f2
1 parent a041c34
commit 579b8f2
Show file tree

Hide file tree

Showing 5 changed files with 198 additions and 53 deletions.
diff --git a/packages/relay-runtime/store/RelayModernQueryExecutor.js b/packages/relay-runtime/store/RelayModernQueryExecutor.js
@@ -481,6 +481,21 @@ class Executor {
       this._updateOperationTracker(updatedOwners);
       this._processPayloadFollowups(payloadFollowups);
     }
+    if (
+      this._isSubscriptionOperation &&
+      RelayFeatureFlags.ENABLE_UNIQUE_SUBSCRIPTION_ROOT
+    ) {
+      // We attach the id to allow the `requestSubscription` to read from the store using
+      // the current id in its `onNext` callback
+      if (responsesWithData[0].extensions == null) {
+        // $FlowFixMe[cannot-write]
+        responsesWithData[0].extensions = {
+          __relay_subscription_root_id: this._operation.fragment.dataID,
+        };
+      } else {
+        responsesWithData[0].extensions.__relay_subscription_root_id = this._operation.fragment.dataID;
+      }
+    }
     this._sink.next(response);
   }
 

diff --git a/...iption/__tests__/__generated__/requestSubscriptionTestConfigCreateSubscription.graphql.js b/...iption/__tests__/__generated__/requestSubscriptionTestConfigCreateSubscription.graphql.js
diff --git a/...time/subscription/__tests__/__generated__/requestSubscriptionTestExtraFragment.graphql.js b/...time/subscription/__tests__/__generated__/requestSubscriptionTestExtraFragment.graphql.js
diff --git a/packages/relay-runtime/subscription/__tests__/requestSubscription-test.js b/packages/relay-runtime/subscription/__tests__/requestSubscription-test.js
@@ -31,7 +31,6 @@ const {ROOT_ID} = require('../../store/RelayStoreUtils');
 const {createMockEnvironment} = require('relay-test-utils-internal');
 
 describe('requestSubscription-test', () => {
-  RelayFeatureFlags.ENABLE_UNIQUE_SUBSCRIPTION_ROOT = true;
   it('Config: `RANGE_ADD`', () => {
     const environment = createMockEnvironment();
     const store = environment.getStore();
@@ -242,29 +241,37 @@ describe('requestSubscription-test', () => {
         store,
       });
     });
-    it('with cacheConfig', () => {
-      requestSubscription(environment, {
-        subscription: CommentCreateSubscription,
-        variables,
-        cacheConfig: {
-          metadata,
-        },
-      });
 
-      expect(cacheMetadata).toEqual(metadata);
-    });
+    function runTests() {
+      it('with cacheConfig', () => {
+        requestSubscription(environment, {
+          subscription: CommentCreateSubscription,
+          variables,
+          cacheConfig: {
+            metadata,
+          },
+        });
 
-    it('without cacheConfig', () => {
-      requestSubscription(environment, {
-        subscription: CommentCreateSubscription,
-        variables,
+        expect(cacheMetadata).toEqual(metadata);
       });
 
-      expect(cacheMetadata).toEqual(undefined);
-    });
+      it('without cacheConfig', () => {
+        requestSubscription(environment, {
+          subscription: CommentCreateSubscription,
+          variables,
+        });
+
+        expect(cacheMetadata).toEqual(undefined);
+      });
+    }
+    RelayFeatureFlags.ENABLE_UNIQUE_SUBSCRIPTION_ROOT = false;
+    runTests();
+    RelayFeatureFlags.ENABLE_UNIQUE_SUBSCRIPTION_ROOT = true;
+    runTests();
   });
 
   it('does not overwrite existing data', () => {
+    RelayFeatureFlags.ENABLE_UNIQUE_SUBSCRIPTION_ROOT = true;
     const ConfigsQuery = getRequest(graphql`
       query requestSubscriptionTestConfigsQuery {
         viewer {
@@ -279,11 +286,18 @@ describe('requestSubscription-test', () => {
       }
     `);
 
+    graphql`
+      fragment requestSubscriptionTestExtraFragment on Config {
+        isEnabled
+      }
+    `;
+
     const ConfigCreateSubscription = getRequest(graphql`
       subscription requestSubscriptionTestConfigCreateSubscription {
         configCreateSubscribe {
           config {
             name
+            ...requestSubscriptionTestExtraFragment
           }
         }
       }
@@ -364,7 +378,14 @@ describe('requestSubscription-test', () => {
     });
     expect(onNext).toBeCalledTimes(1);
     expect(onNext.mock.calls[0][0]).toEqual({
-      configCreateSubscribe: {config: {name: 'Mark'}},
+      configCreateSubscribe: {
+        config: {
+          name: 'Mark',
+          __id: expect.any(String),
+          __fragments: {requestSubscriptionTestExtraFragment: {}},
+          __fragmentOwner: expect.any(Object),
+        },
+      },
     });
 
     environment.mock.nextValue(ConfigCreateSubscription, {
@@ -396,7 +417,14 @@ describe('requestSubscription-test', () => {
     });
     expect(onNext).toBeCalledTimes(2);
     expect(onNext.mock.calls[1][0]).toEqual({
-      configCreateSubscribe: {config: {name: 'Zuck'}},
+      configCreateSubscribe: {
+        config: {
+          name: 'Zuck',
+          __id: expect.any(String),
+          __fragments: {requestSubscriptionTestExtraFragment: {}},
+          __fragmentOwner: expect.any(Object),
+        },
+      },
     });
   });
 });
diff --git a/packages/relay-runtime/subscription/requestSubscription.js b/packages/relay-runtime/subscription/requestSubscription.js
@@ -22,6 +22,7 @@ const {generateUniqueClientID} = require('../store/ClientID');
 const {
   createOperationDescriptor,
 } = require('../store/RelayModernOperationDescriptor');
+const {createReaderSelector} = require('../store/RelayModernSelector');
 
 import type {DeclarativeMutationConfig} from '../mutations/RelayDeclarativeMutationConfig';
 import type {GraphQLTaggedNode} from '../query/GraphQLTag';
@@ -91,17 +92,24 @@ function requestSubscription<TSubscriptionPayload>(
       updater,
     })
     .map(responses => {
+      let selector = operation.fragment;
       if (RelayFeatureFlags.ENABLE_UNIQUE_SUBSCRIPTION_ROOT) {
+        let nextID;
         if (Array.isArray(responses)) {
-          // $FlowFixMe[incompatible-cast]
-          return (responses.map(
-            response => response.data,
-          ): TSubscriptionPayload);
+          nextID = responses[0]?.extensions?.__relay_subscription_root_id;
+        } else {
+          nextID = responses.extensions?.__relay_subscription_root_id;
+        }
+        if (typeof nextID === 'string') {
+          selector = createReaderSelector(
+            selector.node,
+            nextID,
+            selector.variables,
+            selector.owner,
+          );
         }
-        // $FlowFixMe[incompatible-cast]
-        return (responses.data: TSubscriptionPayload);
       }
-      const data = environment.lookup(operation.fragment).data;
+      const data = environment.lookup(selector).data;
       // $FlowFixMe[incompatible-cast]
       return (data: TSubscriptionPayload);
     })