GraphQL: Merge Extensions Into AST

mergeExtensionsIntoAST.js

const invariant = require('invariant');
const { Kind } = require('graphql');

const byKindGetInfo = {
  // SchemaDefinition
  [Kind.SCHEMA_DEFINITION]: def => ({
    isExtension: false,
    type: 'schema',
    typeName: 'schema',
  }),
  // ScalarTypeDefinition
  [Kind.SCALAR_TYPE_DEFINITION]: def => ({
    isExtension: false,
    type: 'scalar',
    typeName: `scalar ${def.name.value}`,
  }),
  // ObjectTypeDefinition
  [Kind.OBJECT_TYPE_DEFINITION]: def => ({
    isExtension: false,
    type: 'type',
    typeName: `type ${def.name.value}`,
  }),
  // InterfaceTypeDefinition
  [Kind.INTERFACE_TYPE_DEFINITION]: def => ({
    isExtension: false,
    type: 'interface',
    typeName: `interface ${def.name.value}`,
  }),
  // UnionTypeDefinition
  [Kind.UNION_TYPE_DEFINITION]: def => ({
    isExtension: false,
    type: 'union',
    typeName: `union ${def.name.value}`,
  }),
  // EnumTypeDefinition
  [Kind.ENUM_TYPE_DEFINITION]: def => ({
    isExtension: false,
    type: 'enum',
    typeName: `enum ${def.name.value}`,
  }),
  // InputObjectTypeDefinition
  [Kind.INPUT_OBJECT_TYPE_DEFINITION]: def => ({
    isExtension: false,
    type: 'input',
    typeName: `input ${def.name.value}`,
  }),
  // DirectiveDefinition
  [Kind.DIRECTIVE_DEFINITION]: def => ({
    isExtension: false,
    type: 'directive',
    typeName: `directive ${def.name.value}`,
  }),

  // SchemaExtension
  [Kind.SCHEMA_EXTENSION]: def => ({
    isExtension: true,
    type: 'schema',
    typeName: 'schema',
  }),
  // ScalarTypeExtension
  [Kind.SCALAR_TYPE_EXTENSION]: def => ({
    isExtension: true,
    type: 'scalar',
    typeName: `scalar ${def.name.value}`,
  }),
  // ObjectTypeExtension
  [Kind.OBJECT_TYPE_EXTENSION]: def => ({
    isExtension: true,
    type: 'type',
    typeName: `type ${def.name.value}`,
  }),
  // InterfaceTypeExtension
  [Kind.INTERFACE_TYPE_EXTENSION]: def => ({
    isExtension: true,
    type: 'interface',
    typeName: `interface ${def.name.value}`,
  }),
  // UnionTypeExtension
  [Kind.UNION_TYPE_EXTENSION]: def => ({
    isExtension: true,
    type: 'union',
    typeName: `union ${def.name.value}`,
  }),
  // EnumTypeExtension
  [Kind.ENUM_TYPE_EXTENSION]: def => ({
    isExtension: true,
    type: 'enum',
    typeName: `enum ${def.name.value}`,
  }),
  // InputObjectTypeExtension
  [Kind.INPUT_OBJECT_TYPE_EXTENSION]: def => ({
    isExtension: true,
    type: 'input',
    typeName: `input ${def.name.value}`,
  }),
};

function extendDefinition(def, ext) {
  const defInfo = byKindGetInfo[def.kind](def);
  const extInfo = byKindGetInfo[ext.kind](ext);
  invariant(
    defInfo.type === extInfo.type,
    'Types must be same: %s != %s',
    defInfo.type,
    extInfo.type,
  );
  invariant(!defInfo.isExtension, 'Extended type must be definition type');
  invariant(extInfo.isExtension, 'Extending type must be extension type');

  const extendLocation = (loc, loc2) => ({
    ...loc,
    ext: loc.ext ? [...loc.ext, loc2] : [loc2],
  });

  switch (defInfo.type) {
    case 'schema': {
      return {
        ...def,
        directives: [...def.directives, ...ext.directives],
        operationTypes: [...def.operationTypes, ...ext.operationTypes],
        loc: extendLocation(def.loc, ext.loc),
      };
    }
    case 'scalar': {
      return {
        ...def,
        directives: [...def.directives, ...ext.directives],
        loc: extendLocation(def.loc, ext.loc),
      };
    }
    case 'type': {
      return {
        ...def,
        interfaces: [...def.interfaces, ...ext.interfaces],
        directives: [...def.directives, ...ext.directives],
        fields: [...def.fields, ...ext.fields],
        loc: extendLocation(def.loc, ext.loc),
      };
    }
    case 'interface': {
      return {
        ...def,
        directives: [...def.directives, ...ext.directives],
        fields: [...def.fields, ...ext.fields],
        loc: extendLocation(def.loc, ext.loc),
      };
    }
    case 'union': {
      return {
        ...def,
        directives: [...def.directives, ...ext.directives],
        types: [...def.types, ...ext.types],
        loc: extendLocation(def.loc, ext.loc),
      };
    }
    case 'enum': {
      return {
        ...def,
        directives: [...def.directives, ...ext.directives],
        values: [...def.values, ...ext.values],
        loc: extendLocation(def.loc, ext.loc),
      };
    }
    case 'input': {
      return {
        ...def,
        directives: [...def.directives, ...ext.directives],
        fields: [...def.fields, ...ext.fields],
        loc: extendLocation(def.loc, ext.loc),
      };
    }
    default: {
      invariant(false, 'Unhandled type for merge: %s', defInfo.type);
      return def;
    }
  }
}

function mergeExtensionsIntoAST(inAst) {
  invariant(inAst.kind === 'Document', 'Document node required');
  const definitions = new Map();
  const extensions = new Map();

  // collect definitions and extensions
  inAst.definitions.forEach(def => {
    invariant(def, 'Definition expected');

    const getKey = byKindGetInfo[def.kind];
    invariant(getKey, 'Cannot retrieve key for %s', def.kind);
    const { isExtension, typeName } = getKey(def);
    if (isExtension) {
      if (extensions.has(typeName)) {
        extensions.get(typeName).push(def);
      } else {
        extensions.set(typeName, [def]);
      }
    } else {
      invariant(
        !definitions.has(typeName),
        'Schema cannot contain multiple definitions: "%s"',
        typeName,
      );
      definitions.set(typeName, def);
    }
  });

  for (const [key, extDefs] of extensions) {
    const def = definitions.get(key);
    definitions.set(key, extDefs.reduce(extendDefinition, def));
  }

  return {
    ...inAst,
    definitions: [...definitions.values()],
  };
}

module.exports = {
  mergeExtensionsIntoAST,
};

utils.js

const util = require('util');
const path = require('path');
const fs = require('fs');
const glob = require('glob');
const invariant = require('invariant');
const {
  Source,
  parse,
  print,
  buildASTSchema,
  TokenKind,
  graphql,
  getIntrospectionQuery,
} = require('graphql');
const { mergeExtensionsIntoAST } = require('./mergeExtensionsIntoAST');
const { format } = require('prettier');
const supportsColor = require('supports-color');
const Chalk = require('chalk').constructor;
const chalk = new Chalk({ level: supportsColor.stderr.level });

const readFile = util.promisify(fs.readFile);
const asyncGlob = util.promisify(glob);

const relativePath = fullPath =>
  path.relative(path.join(__dirname, '../../schema/'), fullPath);

async function loadSources(
  globPattern = path.join(__dirname, '../../schema/**/*.gql'),
) {
  const fullPaths = await asyncGlob(globPattern, {});
  const sources = await Promise.all(
    fullPaths.map(
      async fullPath =>
        new Source(await readFile(fullPath, 'utf-8'), relativePath(fullPath)),
    ),
  );
  return sources;
}

function parseSources(sources) {
  return sources.map(source => {
    try {
      return { source, ...parse(source) };
    } catch (error) {
      const errorMessage = error.locations
        ? `${error.message}${error.locations
            .map(
              ({ line, column }) => `\n  at ${source.name}:${line}:${column}`,
            )
            .join('')}`
        : `${error.message} at ${source.name}`;

      const coloredErrorMessage = error.locations
        ? `${chalk.redBright(error.message)}${error.locations
            .map(
              ({ line, column }) =>
                `\n  at ${chalk.cyanBright(
                  `${source.name}:${line}:${column}`,
                )}`,
            )
            .join('')}`
        : `${error.message} at ${source.name}`;

      console.error(
        `${chalk.whiteBright.bold.bgRed('ERROR:')} ${coloredErrorMessage}`,
      );
      return { source, error, errorMessage };
    }
  });
}

// At first, this will sort by kind.
// Second criteria is ASCII name
// Third is order of keys in this POJO map ;-)
const sortKeyByKind = {
  DirectiveDefinition: 1,

  EnumTypeDefinition: 2,
  EnumTypeExtension: 2,

  ScalarTypeDefinition: 3,
  ScalarTypeExtension: 3,

  InterfaceTypeDefinition: 4,
  InterfaceTypeExtension: 4,

  ObjectTypeDefinition: 5,
  ObjectTypeExtension: 5,

  UnionTypeDefinition: 6,
  UnionTypeExtension: 6,

  InputObjectTypeDefinition: 7,
  InputObjectTypeExtension: 7,

  SchemaDefinition: 8,
  SchemaExtension: 8,
};

const formatKind = kind =>
  ({
    DirectiveDefinition: 'directive',
    EnumTypeDefinition: 'enum',
    EnumTypeExtension: 'extend enum',
    ScalarTypeDefinition: 'scalar',
    ScalarTypeExtension: 'extend scalar',
    InterfaceTypeDefinition: 'interface',
    InterfaceTypeExtension: 'extend interface',
    ObjectTypeDefinition: 'type',
    ObjectTypeExtension: 'extend type',
    UnionTypeDefinition: 'union',
    UnionTypeExtension: 'extend union',
    InputObjectTypeDefinition: 'input',
    InputObjectTypeExtension: 'extend input',
    SchemaDefinition: 'schema',
    SchemaExtension: 'extend schema',
  }[kind] || kind);

const typeKeys = Object.keys(sortKeyByKind);

const makeComparator = (...selectors) => (a, b) => {
  for (const selector of selectors) {
    const valA = selector(a);
    const valB = selector(b);
    if (valA < valB) return -1;
    if (valA > valB) return 1;
  }
  return 0;
};

function sortDefinitions(document) {
  document.definitions.sort(
    makeComparator(
      def => sortKeyByKind[def.kind],
      def => (def.name && def.name.value) || '',
      def => typeKeys.indexOf(def.kind),
      def => (def.loc && def.loc.source && def.loc.source.name) || '',
      def => (def.loc && def.loc.start) || 0,
    ),
  );
  return document;
}

const ignoreTokens = new Set([
  TokenKind.BLOCK_STRING,
  TokenKind.STRING,
  TokenKind.COMMENT,
]);

function getRealDefinitionFirstToken(first) {
  let token = first;
  while (token && ignoreTokens.has(token.kind)) token = token.next;
  return token;
}

function formatDefinitionPosition(def, { color } = {}) {
  const formattedKind = formatKind(def.kind);
  const loc = getRealDefinitionFirstToken(def.loc.startToken);
  const location = `${def.loc.source.name}:${loc.line}:${loc.column}`;
  return [
    color ? chalk.yellowBright.bold(formattedKind) : formattedKind,
    (def.name && def.name.value) || null,
    color ? chalk.blueBright(location) : location,
  ]
    .filter(x => x !== null)
    .join(' ');
}

function mergeParsedIntoDocument(parsedDocuments) {
  const result = parsedDocuments.reduce(
    (result, document) => {
      invariant(
        document.kind === 'Document',
        'mergeParsedIntoDocument can only accept list of GraphQL Documents',
      );
      const tmp = {
        sources: document.source
          ? [...result.sources, document.source]
          : result.sources,
        kind: 'Document',
        definitions: [...result.definitions, ...document.definitions],
      };
      if (document.error) {
        tmp.errors = tmp.errors
          ? [...tmp.errors, document.error]
          : [document.error];
      }
      return tmp;
    },
    {
      sources: [],
      kind: 'Document',
      definitions: [],
    },
  );

  return result;
}

function loadAndParse(glob) {
  return loadSources(glob)
    .then(parseSources)
    .then(mergeParsedIntoDocument)
    .then(sortDefinitions);
}

function printTopDefinitions(document) {
  document.definitions.forEach(def => {
    console.error(formatDefinitionPosition(def, { color: true }));
  });
  return document;
}

function runIntrospectionQuery(schema) {
  return graphql(schema, getIntrospectionQuery({ descriptions: true }));
}

module.exports = {
  loadSources,
  parseSources,
  mergeParsedIntoDocument,
  sortDefinitions,
  loadAndParse,
  printTopDefinitions,
  buildASTSchema,
  runIntrospectionQuery,
  mergeExtensionsIntoAST,
};

🎉 👍

Thanks!