AddInlineFragmentToType.js

// @flow
/* eslint-disable no-use-before-define, consistent-return, no-prototype-builtins, no-underscore-dangle */
// This was mostly ripped from:
// https://github.com/apollographql/graphql-tools/blob/master/src/transforms/ReplaceFieldWithFragment.ts
// It was easier to modify an existing transform to work than it was to
// write a new transform from scratch.
import {
  DocumentNode,
  GraphQLSchema,
  GraphQLType,
  InlineFragmentNode,
  Kind,
  SelectionSetNode,
  TypeInfo,
  OperationDefinitionNode,
  parse,
  visit,
  visitWithTypeInfo,
  SelectionNode,
} from 'graphql';
import { Request, Transform } from 'graphql-tools';

export default class AddInlineFragmentToType extends Transform {
  targetSchema: GraphQLSchema;
  mapping: FieldToFragmentMapping;

  constructor(
    targetSchema: GraphQLSchema,
    fragments: Array<{
      typeName: string;
      fragment: string;
    }>,
  ) {
    super();
    this.targetSchema = targetSchema;
    this.mapping = {};
    fragments.forEach(({ fragment }) => {
      const parsedFragment = parseFragmentToInlineFragment(fragment);
      const actualTypeName = parsedFragment.typeCondition.name.value;

      if (this.mapping[actualTypeName]) {
        this.mapping[actualTypeName].push(parsedFragment);
      } else {
        this.mapping[actualTypeName] = [parsedFragment];
      }
    })
  }

  transformRequest(originalRequest: Request): Request {
    const document = addFragmentsToTypes(
      this.targetSchema,
      originalRequest.document,
      this.mapping,
    );
    return {
      ...originalRequest,
      document,
    };
  }
}

type FieldToFragmentMapping = {
  [typeName: string]: InlineFragmentNode[]
};

function addFragmentsToTypes(
  targetSchema: GraphQLSchema,
  document: DocumentNode,
  mapping: FieldToFragmentMapping,
): DocumentNode {
  const typeInfo = new TypeInfo(targetSchema);
  return visit(
    document,
    visitWithTypeInfo(typeInfo, {
      [Kind.SELECTION_SET](
        node: SelectionSetNode,
      ): SelectionSetNode | null | undefined {
        const parentType: GraphQLType = typeInfo.getParentType();
        if (parentType) {
          const parentTypeName = parentType.name;
          let selections = node.selections;

          if (mapping[parentTypeName]) {
            const fragments = mapping[parentTypeName];
            if (fragments && fragments.length > 0) {
              const fragment = concatInlineFragments(
                parentTypeName,
                fragments,
              );
              selections = selections.concat(fragment);
            }
          }

          if (selections !== node.selections) {
            return {
              ...node,
              selections,
            };
          }
        }
      },
    }),
  );
}

function parseFragmentToInlineFragment(
  definitions: string,
): InlineFragmentNode {
  if (definitions.trim().startsWith('fragment')) {
    const document = parse(definitions);
    for (let i = 0; i < document.definitions.length; i++) {
      const definition = document.definitions[i];
      if (definition.kind === Kind.FRAGMENT_DEFINITION) {
        return {
          kind: Kind.INLINE_FRAGMENT,
          typeCondition: definition.typeCondition,
          selectionSet: definition.selectionSet,
        };
      }
    }
  }

  const query: OperationDefinitionNode = parse(`{${definitions}}`)
    .definitions[0];
  for (let i = 0; i < query.selectionSet.selections.length; i++) {
    const selection = query.selectionSet.selections[i];
    if (selection.kind === Kind.INLINE_FRAGMENT) {
      return selection;
    }
  }

  throw new Error('Could not parse fragment');
}

function concatInlineFragments(
  type: string,
  fragments: InlineFragmentNode[],
): InlineFragmentNode {
  const fragmentSelections: SelectionNode[] = fragments.reduce(
    (selections, fragment) =>
      selections.concat(fragment.selectionSet.selections),
    [],
  );

  const deduplicatedFragmentSelection: SelectionNode[] = deduplicateSelection(
    fragmentSelections,
  );

  return {
    kind: Kind.INLINE_FRAGMENT,
    typeCondition: {
      kind: Kind.NAMED_TYPE,
      name: {
        kind: Kind.NAME,
        value: type,
      },
    },
    selectionSet: {
      kind: Kind.SELECTION_SET,
      selections: deduplicatedFragmentSelection,
    },
  };
}

function deduplicateSelection(nodes: SelectionNode[]): SelectionNode[] {
  const selectionMap = nodes.reduce(
    (map: { [key: string]: SelectionNode }, node: SelectionNode) => {
      switch (node.kind) {
        case 'Field': {
          if (node.alias) {
            if (map.hasOwnProperty(node.alias.value)) {
              return map;
            }
            return {
              ...map,
              [node.alias.value]: node,
            };
          } else if (map.hasOwnProperty(node.name.value)) {
            return map;
          }

          return {
            ...map,
            [node.name.value]: node,
          };
        }
        case 'FragmentSpread': {
          if (map.hasOwnProperty(node.name.value)) {
            return map;
          }
          return {
            ...map,
            [node.name.value]: node,
          };
        }
        case 'InlineFragment': {
          if (map.__fragment) {
            const fragment: InlineFragmentNode = map.__fragment;

            return {
              ...map,
              __fragment: concatInlineFragments(
                fragment.typeCondition.name.value,
                [fragment, node],
              ),
            };
          }
          return {
            ...map,
            __fragment: node,
          };
        }
        default: {
          return map;
        }
      }
    },
    {},
  );

  const selection = Object.keys(selectionMap).reduce(
    (selectionList, node) => selectionList.concat(selectionMap[node]),
    [],
  );

  return selection;
}