Repository: FormidableLabs/urql (ts)

Issue:
Completed subscription not reflected in hook data
There are two cases where a subscription query completed. Either it unmounts or the server completes the source. i was expecting that on source completion the "fetching" value of useSubscription would turn to false but it still stays the same.

I went through the code of the subscription exchange and found out that the complete event from forwardSubscription is sent to the wonka stream via the complete function but for an unknown reason it does not propagate up to the parent. 


Pull request description:
Implement active teardowns and add useQuery features to useSubscription
Resolve #407 

I think the idea of adding "active teardowns" is pretty valuable. It may be required in the future for someone to actively cancel ongoing queries from exchanges. This is already necessary to correctly implement ending subscriptions.

This hence also adds missing features to the `useSubscription` hook, like `executeSubscription` and `args.pause`. It also ensures that the `handler` arg is a constant ref, since reexecuting subscriptions when it changes wouldn't make much sense.

Wonka had a slight bug in `onEnd`. The operator wasn't propagating `Close` signals correctly, so the use of `onEnd` would prevent unsubscriptions from propagating upwards.

Patch:
diff --git a/package.json b/package.json
index f30e03d1..4317b340 100644
--- a/package.json
+++ b/package.json
@@ -136,6 +136,6 @@
   },
   "dependencies": {
     "fast-json-stable-stringify": "^2.0.0",
-    "wonka": "^3.2.0"
+    "wonka": "^3.2.1"
   }
 }
diff --git a/src/client.ts b/src/client.ts
index 970d3d74..56e55bd0 100755
--- a/src/client.ts
+++ b/src/client.ts
@@ -8,6 +8,7 @@ import {
   share,
   Source,
   take,
+  takeUntil,
   merge,
   interval,
   fromValue,
@@ -179,8 +180,14 @@ export class Client {
       );
     }
 
+    const teardown$ = pipe(
+      this.operations$,
+      filter(op => op.operationName === 'teardown' && op.key === key)
+    );
+
     const result$ = pipe(
       operationResults$,
+      takeUntil(teardown$),
       onStart<OperationResult>(() => this.onOperationStart(operation)),
       onEnd<OperationResult>(() => this.onOperationEnd(operation))
     );
diff --git a/src/components/Subscription.ts b/src/components/Subscription.ts
index cfc86dcb..885f8c5e 100644
--- a/src/components/Subscription.ts
+++ b/src/components/Subscription.ts
@@ -1,4 +1,5 @@
 import { ReactElement } from 'react';
+import { OperationContext } from '../types';
 
 import {
   useSubscription,
@@ -9,12 +10,19 @@ import {
 
 export interface SubscriptionProps<T, R, V> extends UseSubscriptionArgs<V> {
   handler?: SubscriptionHandler<T, R>;
-  children: (arg: UseSubscriptionState<R>) => ReactElement<any>;
+  children: (arg: SubscriptionState<R>) => ReactElement<any>;
+}
+
+export interface SubscriptionState<T> extends UseSubscriptionState<T> {
+  executeSubscription: (opts?: Partial<OperationContext>) => void;
 }
 
 export function Subscription<T = any, R = T, V = any>(
   props: SubscriptionProps<T, R, V>
 ): ReactElement<any> {
-  const [state] = useSubscription<T, R, V>(props, props.handler);
-  return props.children(state);
+  const [state, executeSubscription] = useSubscription<T, R, V>(
+    props,
+    props.handler
+  );
+  return props.children({ ...state, executeSubscription });
 }
diff --git a/src/exchanges/subscription.ts b/src/exchanges/subscription.ts
index 3c7eb2e7..70c37418 100644
--- a/src/exchanges/subscription.ts
+++ b/src/exchanges/subscription.ts
@@ -59,7 +59,7 @@ const isSubscriptionOperation = (operation: Operation) =>
 
 export const subscriptionExchange = ({
   forwardSubscription,
-}: SubscriptionExchangeOpts): Exchange => ({ forward }) => {
+}: SubscriptionExchangeOpts): Exchange => ({ client, forward }) => {
   const createSubscriptionSource = (
     operation: Operation
   ): Source<OperationResult> => {
@@ -72,13 +72,27 @@ export const subscriptionExchange = ({
     });
 
     return make<OperationResult>(([next, complete]) => {
+      let isComplete = false;
+
       const sub = observableish.subscribe({
         next: result => next(makeResult(operation, result)),
         error: err => next(makeErrorResult(operation, err)),
-        complete,
+        complete: () => {
+          if (!isComplete) {
+            client.reexecuteOperation({
+              ...operation,
+              operationName: 'teardown',
+            });
+          }
+
+          complete();
+        },
       });
 
-      return () => sub.unsubscribe();
+      return () => {
+        isComplete = true;
+        sub.unsubscribe();
+      };
     });
   };
 
diff --git a/src/hooks/useQuery.ts b/src/hooks/useQuery.ts
index 03bb45dd..df3c094b 100644
--- a/src/hooks/useQuery.ts
+++ b/src/hooks/useQuery.ts
@@ -1,6 +1,6 @@
 import { DocumentNode } from 'graphql';
 import { useCallback, useContext, useRef } from 'react';
-import { pipe, subscribe } from 'wonka';
+import { pipe, onEnd, subscribe } from 'wonka';
 import { Context } from '../context';
 import { OperationContext, RequestPolicy } from '../types';
 import { CombinedError, noop } from '../utils';
@@ -61,6 +61,7 @@ export const useQuery = <T = any, V = object>(
           ...args.context,
           ...opts,
         }),
+        onEnd(() => setState(s => ({ ...s, fetching: false }))),
         subscribe(({ data, error, extensions }) => {
           setState({ fetching: false, data, error, extensions });
         })
diff --git a/src/hooks/useSubscription.ts b/src/hooks/useSubscription.ts
index 8562c414..a441e1b6 100644
--- a/src/hooks/useSubscription.ts
+++ b/src/hooks/useSubscription.ts
@@ -1,15 +1,17 @@
 import { DocumentNode } from 'graphql';
-import { useCallback, useContext, useEffect, useRef } from 'react';
-import { pipe, subscribe } from 'wonka';
+import { useCallback, useContext, useRef } from 'react';
+import { pipe, onEnd, subscribe } from 'wonka';
 import { Context } from '../context';
 import { CombinedError, noop } from '../utils';
 import { useRequest } from './useRequest';
+import { useImmediateEffect } from './useImmediateEffect';
 import { useImmediateState } from './useImmediateState';
 import { OperationContext } from '../types';
 
 export interface UseSubscriptionArgs<V> {
   query: DocumentNode | string;
   variables?: V;
+  pause?: boolean;
   context?: Partial<OperationContext>;
 }
 
@@ -22,49 +24,71 @@ export interface UseSubscriptionState<T> {
   extensions?: Record<string, any>;
 }
 
-export type UseSubscriptionResponse<T> = [UseSubscriptionState<T>];
+export type UseSubscriptionResponse<T> = [
+  UseSubscriptionState<T>,
+  (opts?: Partial<OperationContext>) => void
+];
 
 export const useSubscription = <T = any, R = T, V = object>(
   args: UseSubscriptionArgs<V>,
   handler?: SubscriptionHandler<T, R>
 ): UseSubscriptionResponse<R> => {
   const unsubscribe = useRef(noop);
+  const handlerRef = useRef(handler);
   const client = useContext(Context);
 
   const [state, setState] = useImmediateState<UseSubscriptionState<R>>({
-    fetching: true,
+    fetching: false,
     error: undefined,
     data: undefined,
     extensions: undefined,
   });
 
+  // Update handler on constant ref, since handler changes shouldn't
+  // trigger a new subscription run
+  handlerRef.current = handler;
+
   // This creates a request which will keep a stable reference
   // if request.key doesn't change
   const request = useRequest(args.query, args.variables);
 
-  const executeSubscription = useCallback(() => {
-    unsubscribe.current();
+  const executeSubscription = useCallback(
+    (opts?: Partial<OperationContext>) => {
+      unsubscribe.current();
+
+      setState(s => ({ ...s, fetching: true }));
+
+      [unsubscribe.current] = pipe(
+        client.executeSubscription(request, {
+          ...args.context,
+          ...opts,
+        }),
+        onEnd(() => setState(s => ({ ...s, fetching: false }))),
+        subscribe(({ data, error, extensions }) => {
+          const { current: handler } = handlerRef;
+
+          setState(s => ({
+            fetching: true,
+            data: typeof handler === 'function' ? handler(s.data, data) : data,
+            error,
+            extensions,
+          }));
+        })
+      );
+    },
+    [client, request, setState, args.context]
+  );
 
-    [unsubscribe.current] = pipe(
-      client.executeSubscription(request, args.context || {}),
-      subscribe(({ data, error, extensions }) => {
-        setState(s => ({
-          fetching: true,
-          data: handler !== undefined ? handler(s.data, data) : data,
-          error,
-          extensions,
-        }));
-      })
-    );
-  }, [client, handler, request, setState, args.context]);
+  useImmediateEffect(() => {
+    if (args.pause) {
+      unsubscribe.current();
+      setState(s => ({ ...s, fetching: false }));
+      return noop;
+    }
 
-  // Trigger subscription on query change
-  // We don't use useImmediateEffect here as we have no way of
-  // unsubscribing from subscriptions during SSR
-  useEffect(() => {
     executeSubscription();
     return () => unsubscribe.current(); // eslint-disable-line
-  }, [executeSubscription]);
+  }, [executeSubscription, args.pause, setState]);
 
-  return [state];
+  return [state, executeSubscription];
 };
diff --git a/yarn.lock b/yarn.lock
index a4e398f4..61af5ae0 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -6255,10 +6255,10 @@ wide-align@^1.1.0:
   dependencies:
     string-width "^1.0.2 || 2"
 
-wonka@^3.2.0:
-  version "3.2.0"
-  resolved "https://registry.yarnpkg.com/wonka/-/wonka-3.2.0.tgz#40b9ce8e0c375127437661290e98e5d9e5278440"
-  integrity sha512-sHCLouGzTclUlp1sdXAhuTqGMIoAQfpBit9puepaINPVqDBRtC8DPWtTUHC0DJIdFSO+DsO86AFBFCDor7Sn0g==
+wonka@^3.2.1:
+  version "3.2.1"
+  resolved "https://registry.yarnpkg.com/wonka/-/wonka-3.2.1.tgz#f66123fabbb260f23a4e89f1526317169203e258"
+  integrity sha512-LpquZ9PJtYJ31AlMPp5tEba1y+S+wUmY7OyFLGZCmGSNO1ZOTTZGp/HbieFzP5fVJ9fWXCus3ecgDjwAssHSGg==
 
 wordwrap@~0.0.2:
   version "0.0.3"


Test patch:
diff --git a/src/exchanges/subscription.test.ts b/src/exchanges/subscription.test.ts
index 93cddd7a..759b4233 100644
--- a/src/exchanges/subscription.test.ts
+++ b/src/exchanges/subscription.test.ts
@@ -1,16 +1,24 @@
 import { print } from 'graphql';
-import { empty, fromValue, pipe, Source, take, toPromise } from 'wonka';
+import {
+  empty,
+  publish,
+  fromValue,
+  pipe,
+  Source,
+  take,
+  toPromise,
+} from 'wonka';
 import { Client } from '../client';
 import { subscriptionOperation, subscriptionResult } from '../test-utils';
 import { OperationResult } from '../types';
 import { subscriptionExchange, SubscriptionForwarder } from './subscription';
 
-const exchangeArgs = {
-  forward: () => empty as Source<OperationResult>,
-  client: {} as Client,
-};
-
 it('should return response data from forwardSubscription observable', async () => {
+  const exchangeArgs = {
+    forward: () => empty as Source<OperationResult>,
+    client: {} as Client,
+  };
+
   const unsubscribe = jest.fn();
   const forwardSubscription: SubscriptionForwarder = operation => {
     expect(operation.query).toBe(print(subscriptionOperation.query));
@@ -38,3 +46,29 @@ it('should return response data from forwardSubscription observable', async () =
   expect(data).toMatchSnapshot();
   expect(unsubscribe).toHaveBeenCalled();
 });
+
+it('should tear down the operation if the source subscription ends', () => {
+  const reexecuteOperation = jest.fn();
+  const unsubscribe = jest.fn();
+
+  const exchangeArgs = {
+    forward: () => empty as Source<OperationResult>,
+    client: { reexecuteOperation: reexecuteOperation as any } as Client,
+  };
+
+  const forwardSubscription: SubscriptionForwarder = () => ({
+    subscribe(observer) {
+      observer.complete();
+      return { unsubscribe };
+    },
+  });
+
+  pipe(
+    fromValue(subscriptionOperation),
+    subscriptionExchange({ forwardSubscription })(exchangeArgs),
+    publish
+  );
+
+  expect(unsubscribe).not.toHaveBeenCalled();
+  expect(reexecuteOperation).toHaveBeenCalled();
+});
diff --git a/src/hooks/__snapshots__/useSubscription.test.tsx.snap b/src/hooks/__snapshots__/useSubscription.test.tsx.snap
index b4db0876..3202ca54 100644
--- a/src/hooks/__snapshots__/useSubscription.test.tsx.snap
+++ b/src/hooks/__snapshots__/useSubscription.test.tsx.snap
@@ -2,8 +2,8 @@
 
 exports[`on initial useEffect initialises default state 1`] = `
 Object {
-  "data": undefined,
-  "error": undefined,
+  "data": 1234,
+  "error": 5678,
   "extensions": undefined,
   "fetching": true,
 }
diff --git a/src/hooks/useQuery.test.tsx b/src/hooks/useQuery.test.tsx
index 2740be19..b1269b0b 100644
--- a/src/hooks/useQuery.test.tsx
+++ b/src/hooks/useQuery.test.tsx
@@ -20,7 +20,7 @@ jest.mock('../client', () => {
 
 import React, { FC } from 'react';
 import renderer, { act } from 'react-test-renderer';
-import { pipe, onStart, onEnd, interval } from 'wonka';
+import { pipe, onStart, onEnd, empty, never } from 'wonka';
 import { createClient } from '../client';
 import { OperationContext } from '../types';
 import { useQuery, UseQueryArgs, UseQueryState } from './useQuery';
@@ -185,7 +185,7 @@ describe('on unmount', () => {
   beforeEach(() => {
     client.executeQuery.mockReturnValue(
       pipe(
-        interval(400),
+        never,
         onStart(start),
         onEnd(unsubscribe)
       )
@@ -194,17 +194,25 @@ describe('on unmount', () => {
 
   it('unsubscribe is called', () => {
     const wrapper = renderer.create(<QueryUser {...props} />);
-    wrapper.unmount();
+    act(() => wrapper.unmount());
     expect(start).toBeCalledTimes(1);
     expect(unsubscribe).toBeCalledTimes(1);
   });
 });
 
+describe('active teardown', () => {
+  it('sets fetching to false when the source ends', () => {
+    client.executeQuery.mockReturnValueOnce(empty);
+    renderer.create(<QueryUser {...props} />);
+    expect(client.executeQuery).toHaveBeenCalled();
+    expect(state).toMatchObject({ fetching: false });
+  });
+});
+
 describe('execute query', () => {
   it('triggers query execution', () => {
     renderer.create(<QueryUser {...props} />);
-    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-    act(() => execute!());
+    act(() => execute && execute());
     expect(client.executeQuery).toBeCalledTimes(2);
   });
 });
diff --git a/src/hooks/useSubscription.test.tsx b/src/hooks/useSubscription.test.tsx
index 2a1db21f..87360ca0 100644
--- a/src/hooks/useSubscription.test.tsx
+++ b/src/hooks/useSubscription.test.tsx
@@ -2,9 +2,9 @@
 jest.mock('../client', () => {
   const d = { data: 1234, error: 5678 };
   // eslint-disable-next-line @typescript-eslint/no-var-requires
-  const { fromArray } = require('wonka');
+  const { merge, fromValue, never } = require('wonka');
   const mock = {
-    executeSubscription: jest.fn(() => fromArray([d])),
+    executeSubscription: jest.fn(() => merge([fromValue(d), never])),
   };
 
   return {
@@ -15,24 +15,27 @@ jest.mock('../client', () => {
 
 import React, { FC } from 'react';
 import renderer, { act } from 'react-test-renderer';
+import { empty } from 'wonka';
 // @ts-ignore - data is imported from mock only
 import { createClient, data } from '../client';
-import { useSubscription } from './useSubscription';
+import { useSubscription, UseSubscriptionState } from './useSubscription';
 import { OperationContext } from '../types';
 
 // @ts-ignore
 const client = createClient() as { executeSubscription: jest.Mock };
 const query = 'subscription Example { example }';
-let state: any;
+
+let state: UseSubscriptionState<any> | undefined;
+let execute: ((opts?: Partial<OperationContext>) => void) | undefined;
 
 const SubscriptionUser: FC<{
   q: string;
   handler?: (prev: any, data: any) => any;
   context?: Partial<OperationContext>;
-}> = ({ q, handler, context }) => {
-  const [s] = useSubscription({ query: q, context }, handler);
-  state = s;
-  return <p>{s.data}</p>;
+  pause?: boolean;
+}> = ({ q, handler, context, pause = false }) => {
+  [state, execute] = useSubscription({ query: q, context, pause }, handler);
+  return <p>{state.data}</p>;
 };
 
 beforeEach(() => {
@@ -121,3 +124,38 @@ it('calls handler', () => {
   expect(handler).toBeCalledTimes(1);
   expect(handler).toBeCalledWith(undefined, 1234);
 });
+
+describe('active teardown', () => {
+  it('sets fetching to false when the source ends', () => {
+    client.executeSubscription.mockReturnValueOnce(empty);
+    renderer.create(<SubscriptionUser q={query} />);
+    expect(client.executeSubscription).toHaveBeenCalled();
+    expect(state).toMatchObject({ fetching: false });
+  });
+});
+
+describe('execute subscription', () => {
+  it('triggers subscription execution', () => {
+    renderer.create(<SubscriptionUser q={query} />);
+    act(() => execute && execute());
+    expect(client.executeSubscription).toBeCalledTimes(2);
+  });
+});
+
+describe('pause', () => {
+  const props = { q: query };
+
+  it('skips executing the query if pause is true', () => {
+    renderer.create(<SubscriptionUser {...props} pause={true} />);
+    expect(client.executeSubscription).not.toBeCalled();
+  });
+
+  it('skips executing queries if pause updates to true', () => {
+    const wrapper = renderer.create(<SubscriptionUser {...props} />);
+
+    wrapper.update(<SubscriptionUser {...props} pause={true} />);
+    wrapper.update(<SubscriptionUser {...props} pause={true} />);
+    expect(client.executeSubscription).toBeCalledTimes(1);
+    expect(state).toMatchObject({ fetching: false });
+  });
+});