useSubscription.ts

import isEqual from 'lodash/isEqual';
import { useContext, useEffect, useRef } from 'react';
import { ReactRelayContext, requestSubscription } from 'react-relay';
import {
  GraphQLSubscriptionConfig,
  GraphQLTaggedNode,
  OperationType,
} from 'relay-runtime';
import useCommittedRef from '@restart/hooks/useCommittedRef';
import usePrevious from '@restart/hooks/usePrevious';

type RelayOperation<T = unknown> = {
  variables: {
    input: object;
  };
  response: {
    [index: string]: T;
  };
};

interface RelayInput<T extends OperationType> {
  input?: T['variables']['input'];
}

export type RelaySubscriptionConfig<TOperation extends RelayOperation> = Omit<
  GraphQLSubscriptionConfig<TOperation['response']>,
  'variables'
> &
  RelayInput<TOperation>;

export default function useSubscription<
  TOperation extends RelayOperation = never
>(
  subscription: GraphQLTaggedNode,
  config?: Partial<RelaySubscriptionConfig<TOperation>>,
) {
  const currentInput = config && config.input;

  const { environment } = useContext(ReactRelayContext)!;

  const prevValues = usePrevious({ subscription, input: currentInput });
  const configRef = useCommittedRef(config);

  const inputVersionRef = useRef(0);

  if (prevValues && !isEqual(currentInput, prevValues.input)) {
    ++inputVersionRef.current;
  }

  useEffect(() => {
    const { input, ...finalConfig } = configRef.current!;
    const disposable = requestSubscription(environment, {
      subscription,
      variables: { input },
      ...finalConfig,
    });

    return () => {
      disposable.dispose();
    };
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [subscription, inputVersionRef.current, environment]);
}