FragmentWatcherLinks.ts. Part of Benchling's double-writes approach to adopting GraphQL

FragmentWatcherLink.ts

/**
 * Apollo Link that notifies about queried fragments when data comes back
 *
 * We're intentionally not notifying about inline fragments since they generally serve a different purpose
 * (assisting in polymorphism, rather than reusing query chunks).
 */
import {ApolloLink, FetchResult, NextLink, Observable, Operation} from 'apollo-link';
import {
  ASTKindToNode,
  ASTNode,
  DocumentNode,
  FieldNode,
  FragmentDefinitionNode,
  InlineFragmentNode,
  SelectionNode,
  SelectionSetNode,
} from 'graphql/language/ast';
import _, {noop} from 'underscore';

import {assert, assertAndContinue} from 'Asserts';
import {possibleTypeNamesForTypeCondition} from 'graphql-client/possibleTypeNamesForFragment';
import {assertNever} from 'util/Types';
import {OpSyncState} from './OpSyncState';

function getASTNodesOfKind<K extends keyof ASTKindToNode>(
  nodes: ReadonlyArray<ASTNode>,
  kind: K
): ReadonlyArray<ASTKindToNode[K]> {
  return nodes.filter((node) => node.kind === kind) as Array<ASTKindToNode[K]>;
}

function getSingleASTNodeOfKind<K extends keyof ASTKindToNode>(
  nodes: ReadonlyArray<ASTNode>,
  kind: K
): ASTKindToNode[K] {
  const matching = getASTNodesOfKind(nodes, kind);
  if (matching.length !== 1) {
    throw new Error(`Expected single node of kind=${kind}`);
  }
  return matching[0];
}


function fieldResultKey(fieldNode: FieldNode): string {
  return (fieldNode.alias && fieldNode.alias.value) || fieldNode.name.value;
}

interface Data {
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  [key: string]: any;
}

function shouldSearchFragmentNode(
  fragmentNode: InlineFragmentNode | FragmentDefinitionNode,
  data: Data
): boolean {
  // We don't return "fragment matches" for inline fragments since they typically serve a different purpose
  // from normal fragment spreads (e.g. union/polymorphic field selection or grouping fields under a shared
  // directive).
  // This function determines whether we should search into the node - whether the result actually matches
  // what the fragment (inline or typed spread) switched on, if anything.
  return (
    data &&
    (!fragmentNode.typeCondition ||
      possibleTypeNamesForTypeCondition(fragmentNode.typeCondition).includes(data.__typename))
  );
}

/**
 * Given a SelectionSet and overall data object, recursively extract the selections from the data.
 *
 * A SelectionSet is a set of selections, where a selection could be a field (which may have a nested
 * selection set), a fragment spread, or an inline fragment spread (see
 * https://graphql.org/learn/queries/#inline-fragments).
 *
 *   query {
 *     foo
 *     ...bar
 *     baz {
 *       bang
 *     }
 *     ... on florp {
 *       <...>
 *     }
 *   }
 *
 * In this query, the top-level SelectionSet has children Selection(foo), FragmentSpread(bar), and
 * Selection(baz) with a nested SelectionSet consisting of bang.
 *
 * This is important because multiple fragments could request overlapping subsets of data (to an arbitrary
 * depth), and we don't want to return more data than the fragment is supposed to have. For example, given a
 * query:
 *
 *   query {
 *     ...bar
 *     baz {
 *       bang
 *     }
 *   }
 *
 * and fragment bar:
 *
 *   fragment bar on <...> {
 *     baz {
 *       boop
 *     }
 *   }
 *
 * We wouldn't want the bar fragment to just return data.baz, because that'd include the additionally
 * requested bang field.
 *
 * Finally, arrays are a little tricky here - since we're not reading from the schema (primarily for
 * simplicity's sake, though it's certainly an option), we don't know which fields are supposed to be arrays,
 * and thus need to carefully check the type of data.
 */
function extractSelectionSetData(
  fragDefsByName: {[key: string]: FragmentDefinitionNode},
  selectionSet: SelectionSetNode,
  data: Data | null
): Data | null {
  if (data == null) {
    return null;
  }
  if (Array.isArray(data)) {
    return data.map((item) => extractSelectionSetData(fragDefsByName, selectionSet, item));
  }

  let result: Data = {};

  for (const selection of selectionSet.selections) {
    if (selection.kind === 'Field') {
      const key = fieldResultKey(selection);
      if (selection.selectionSet === undefined) {
        // e.g. foo in the above example - set result.foo = data.foo
        // result[key] may be null
        result[key] = data[key];
      } else {
        // e.g. baz in the above example - set result.baz to be the recursively computed selection data
        // result[key] may be null
        result[key] = extractSelectionSetData(fragDefsByName, selection.selectionSet, data[key]);
      }
    } else if (
      selection.kind === 'FragmentSpread' &&
      shouldSearchFragmentNode(fragDefsByName[selection.name.value], data)
    ) {
      // e.g. bar in the above example - since a fragment spread is splatting fields, splat the corresponding
      // data into result.
      const fragmentResult = extractFragmentData(fragDefsByName, fragDefsByName[selection.name.value], data);
      assert(fragmentResult != null, 'Fragment must only be null if data is null');
      result = {...result, ...fragmentResult};
    } else if (selection.kind === 'InlineFragment' && shouldSearchFragmentNode(selection, data)) {
      // e.g. "... on florp" in the above example - similarly splat the fields here
      const inlineResult = extractSelectionSetData(fragDefsByName, selection.selectionSet, data);
      assert(inlineResult != null, 'InlineFragment must only be null if data is null');
      result = {...result, ...inlineResult};
    } else if (selection.kind !== 'InlineFragment' && selection.kind !== 'FragmentSpread') {
      assertAndContinue(false, `Unexpected selection kind=${assertNever(selection).kind}`);
    }
  }

  return result;
}

function extractFragmentData(
  fragDefsByName: {[key: string]: FragmentDefinitionNode},
  fragDef: FragmentDefinitionNode,
  data: Data
): Data | null {
  return extractSelectionSetData(fragDefsByName, fragDef.selectionSet, data);
}

// Extracted data we'll eventually notify about
export interface FragmentResult {
  fragmentName: string;
  data: Data | null;
}

function assertNoConditionalFieldSelection(selection: SelectionNode): void {
  // The GraphQL standard defines two directives - include and skip - which tell the server to conditionally
  // include/exclude certain fields from the query. This is useful for network-sensitive applications that
  // want to conserve bytes over the wire, but that's not a major goal of ours. It's tricky to provide a
  // consistent behavior when these are present without building first-class support, so disallow them.
  assert(
    !(selection.directives || []).some((d) => d.name.value === 'skip' || d.name.value === 'include'),
    'Conditional field selection is not supported'
  );
}

function recursivelyExtractResultsFromSelectionSet(
  fragDefsByName: {[key: string]: FragmentDefinitionNode},
  selectionSet: SelectionSetNode,
  data: Data
): Array<FragmentResult> {
  if (Array.isArray(data)) {
    return data.flatMap((item) =>
      recursivelyExtractResultsFromSelectionSet(fragDefsByName, selectionSet, item)
    );
  }

  // For each selection in the selectionSet:
  //   - If the kind is "FragmentSpread" and the typename matches, include its name and data as a result, then
  //     recurse to look for other fragments.
  //   - If the kind is "Field" and its selectionSet is not undefined (nested selection), recurse to look for
  //     nested fragments.
  //   - If the kind is "InlineFragment" and the typename matches, recurse to look for nested fragments.

  let results: Array<FragmentResult> = [];
  for (const selection of selectionSet.selections) {
    assertNoConditionalFieldSelection(selection);

    if (
      selection.kind === 'FragmentSpread' &&
      shouldSearchFragmentNode(fragDefsByName[selection.name.value], data)
    ) {
      const fragDefinition = fragDefsByName[selection.name.value];
      results = [
        ...results,
        ...recursivelyExtractResultsFromSelectionSet(fragDefsByName, fragDefinition.selectionSet, data),
        // Add the current fragment at the end since:
        //
        //   - Some legacy data stitching logic may expect that once a top-level object is available, all of
        //     its children are available as well (e.g. when a SampleGroup is in the store, all associated
        //     StageRunSamples are in the store)
        //   - Fragments are synced to the store in the order they're returned here
        {
          fragmentName: selection.name.value,
          data: extractFragmentData(fragDefsByName, fragDefinition, data),
        },
      ];
    } else if (
      selection.kind === 'Field' &&
      selection.selectionSet !== undefined &&
      data != null &&
      data[fieldResultKey(selection)] != null
    ) {
      results = [
        ...results,
        ...recursivelyExtractResultsFromSelectionSet(
          fragDefsByName,
          selection.selectionSet,
          data[fieldResultKey(selection)]
        ),
      ];
    } else if (selection.kind === 'InlineFragment' && shouldSearchFragmentNode(selection, data)) {
      results = [
        ...results,
        ...recursivelyExtractResultsFromSelectionSet(fragDefsByName, selection.selectionSet, data),
      ];
    }
  }

  return results;
}

// Exported for tests, but the main export is FragmentWatcherLink.
export function computeFragmentResults(query: DocumentNode, data: Data): Array<FragmentResult> {
  const opDef = getSingleASTNodeOfKind(query.definitions, 'OperationDefinition');
  if (opDef.operation !== 'query') {
    return [];
  }

  const fragDefs = getASTNodesOfKind(query.definitions, 'FragmentDefinition');
  const fragDefsByName = _.indexBy(fragDefs, (def) => def.name.value);

  return recursivelyExtractResultsFromSelectionSet(fragDefsByName, opDef.selectionSet, data);
}

type PendingFragmentResult = FragmentResult & {onSynced: () => void};

class FragmentWatcherLink extends ApolloLink {
  onFragmentsFound: (fragmentResults: Array<FragmentResult>) => void;
  pendingFragmentResults: Array<PendingFragmentResult>;
  debouncedNotifyFragments: () => void;

  constructor(onFragmentsFound: (fragmentResults: Array<FragmentResult>) => void) {
    super();
    this.onFragmentsFound = onFragmentsFound;

    // Debounce actually posting the update in case this is a batch request - we'd be processing a lot of
    // (likely similar) operations in a row, and don't want to repeatedly sync them all to Flux/Redux.
    this.pendingFragmentResults = [];
    this.debouncedNotifyFragments = _.debounce(() => {
      this.onFragmentsFound(this.pendingFragmentResults);
      this.pendingFragmentResults.forEach((result) => {
        result.onSynced();
      });
      this.pendingFragmentResults = [];
    }, 1 /* ms */);
  }

  request(operation: Operation, forward: NextLink | undefined): Observable<FetchResult> {
    // Store the context as a separate variable since we debounce calling setSyncState(), and Apollo may try
    // to dispose the operation after the query completes.
    // setSyncState is only passed by useSkippableBenchlingQuery (which is in the call tree of everything in
    // BenchlingQuery.tsx), so direct queries (i.e. getApolloClient().query()) don't pass that context. For
    // now, assume that direct queries don't rely on data being synced.
    const {setSyncState = noop} = operation.getContext();

    setSyncState(OpSyncState.WILL_SYNC);
    const onQuerySynced = () => {
      setSyncState(OpSyncState.SYNC_COMPLETE);
    };

    // forward is assumed to be non-null since this isn't a terminating link
    // https://www.apollographql.com/docs/link/overview#terminating
    return forward!(operation).map((data) => {
      if (!data.data) {
        // Request errored; ignore
        return data;
      }

      // TODO: We may need to dedupe the fragment results. That should help performance, but it may also save
      // us from errors if we turn on server-side response deduplication (roughly, that works by only
      // returning id and __typename for objects that have already been serialized somewhere in the response)
      const newResults = computeFragmentResults(operation.query, data.data);

      if (newResults.length) {
        const onFragmentSynced = _.after(newResults.length, onQuerySynced) as () => void;
        this.pendingFragmentResults = [
          ...this.pendingFragmentResults,
          ...newResults.map((result) => ({...result, onSynced: onFragmentSynced})),
        ];
        this.debouncedNotifyFragments();
      } else {
        onQuerySynced();
      }

      return data;
    });
  }
}

export default FragmentWatcherLink;