Relay hook to handle authenticantion

useAuth.js

import { ROOT_ID } from 'relay-runtime';
import { useRelayEnvironment } from 'react-relay/hooks';
import { useLocation, useHistory } from 'react-router-dom';
import { commitLocalUpdate } from 'react-relay'
import { useMutation } from 'relay-hooks/lib';

import { AuthUserMutation } from 'mutations/AuthUserMutation';

export const TOKEN_KEY = 'KEY';

export const useAuthUser = () => {
  const environment = useRelayEnvironment();

  const loginUpdater = store => {
    const me = store.getRootField('AuthUser').getLinkedRecord('me');
    store.get(ROOT_ID).setLinkedRecord(me, 'me');
  };

  const onCompletedLogin = data => {
    const { AuthUser } = data;
    const { token } = AuthUser;
    localStorage.setItem(TOKEN_KEY, token);
  };
  
  const [mutate] = useMutation(AuthUserMutation, {
    onCompleted,
    updater: loginUpdater,
  });
  
 const authUser = ({ email, password }) => 
 	mutate({ variables: { input: { email, password } } });

  
  const logoffUpdater = store => {
    const me = store.get(ROOT_ID).getLinkedRecord('me');
    return store.delete(me.getDataID());
  };
  
  const logoffUser = () => {
    localStorage.removeItem(TOKEN_KEY);
    commitLocalUpdate(environment, logoffUpdater);
  }
  
  return { authUser, logoffUser };
};