@pothos/plugin-sub-graph/src/index.ts

A Pothos plugin for creating multiple variants or sub-selections of the same graph

import './global-types';
import SchemaBuilder, {
  BasePlugin,
  type PothosInputFieldConfig,
  type PothosOutputFieldConfig,
  PothosSchemaError,
  type PothosTypeConfig,
  type SchemaTypes,
} from '@pothos/core';
import {
  GraphQLEnumType,
  type GraphQLFieldConfigArgumentMap,
  type GraphQLFieldConfigMap,
  type GraphQLInputFieldConfigMap,
  GraphQLInputObjectType,
  GraphQLInterfaceType,
  type GraphQLNamedType,
  GraphQLObjectType,
  GraphQLScalarType,
  GraphQLSchema,
  GraphQLUnionType,
  getNamedType,
  isInterfaceType,
  isNonNullType,
  isObjectType,
} from 'graphql';
import { replaceType } from './util';

const pluginName = 'subGraph';

export default pluginName;

function matchesSubGraphs(left: string[], right: string[], mode: 'any' | 'all') {
  if (mode === 'all') {
    return right.every((entry) => left.includes(entry));
  }
  for (const entry of left) {
    if (right.includes(entry)) {
      return true;
    }
  }

  return false;
}

export class PothosSubGraphPlugin<Types extends SchemaTypes> extends BasePlugin<Types> {
  static createSubGraph<Types extends SchemaTypes>(
    schema: GraphQLSchema,
    subGraph: string[] | string | { all: string[] },
    builder: PothosSchemaTypes.SchemaBuilder<Types>,
  ) {
    const mode: 'any' | 'all' = Array.isArray(subGraph)
      ? 'any'
      : typeof subGraph === 'string'
        ? 'any'
        : 'all';
    const subGraphs = Array.isArray(subGraph)
      ? subGraph
      : typeof subGraph === 'string'
        ? [subGraph]
        : subGraph.all;

    const config = schema.toConfig();
    const newTypes = PothosSubGraphPlugin.filterTypes(config.types, subGraphs, mode);
    const returnedInterfaces = new Set<string>();

    for (const type of newTypes.values()) {
      if (isObjectType(type) || isInterfaceType(type)) {
        const fields = type.getFields();

        for (const field of Object.values(fields)) {
          const namedType = getNamedType(field.type);

          if (isInterfaceType(namedType)) {
            returnedInterfaces.add(namedType.name);
          }
        }
      }
    }

    function hasReturnedInterface(type: GraphQLInterfaceType | GraphQLObjectType): boolean {
      for (const iface of type.getInterfaces()) {
        if (returnedInterfaces.has(iface.name)) {
          return true;
        }

        if (hasReturnedInterface(iface)) {
          return true;
        }
      }

      return false;
    }

    return new GraphQLSchema({
      directives: config.directives,
      extensions: config.extensions,
      extensionASTNodes: config.extensionASTNodes,
      assumeValid: false,
      query: newTypes.get(schema.getQueryType()?.name ?? 'Query') as GraphQLObjectType,
      mutation: newTypes.get(schema.getMutationType()?.name ?? 'Mutation') as GraphQLObjectType,
      subscription: newTypes.get(
        schema.getSubscriptionType()?.name ?? 'Subscription',
      ) as GraphQLObjectType,
      // Explicitly include types that implement an interface that can be resolved in the subGraph
      types: [...newTypes.values()].filter(
        (type) =>
          builder.options.subGraphs?.explicitlyIncludeType?.(type, subGraphs) ||
          ((isObjectType(type) || isInterfaceType(type)) &&
            hasReturnedInterface(type as GraphQLInterfaceType | GraphQLObjectType)),
      ),
    });
  }

  static filterTypes(types: readonly GraphQLNamedType[], subGraphs: string[], mode: 'any' | 'all') {
    const newTypes = new Map<string, GraphQLNamedType>();

    for (const type of types) {
      if (type.name.startsWith('__')) {
        continue;
      }

      if (
        type.name === 'String' ||
        type.name === 'Int' ||
        type.name === 'Float' ||
        type.name === 'Boolean' ||
        type.name === 'ID'
      ) {
        newTypes.set(type.name, type);
      }

      if (!matchesSubGraphs((type.extensions?.subGraphs as string[]) || [], subGraphs, mode)) {
        continue;
      }

      if (type instanceof GraphQLScalarType || type instanceof GraphQLEnumType) {
        newTypes.set(type.name, type);
      } else if (type instanceof GraphQLObjectType) {
        const typeConfig = type.toConfig();
        newTypes.set(
          type.name,
          new GraphQLObjectType({
            ...typeConfig,
            interfaces: () =>
              typeConfig.interfaces
                .filter((iface) => newTypes.has(iface.name))
                .map((iface) => replaceType(iface, newTypes, typeConfig.name, subGraphs)),
            fields: PothosSubGraphPlugin.filterFields(type, newTypes, subGraphs, mode),
          }),
        );
      } else if (type instanceof GraphQLInterfaceType) {
        const typeConfig = type.toConfig();
        newTypes.set(
          type.name,
          new GraphQLInterfaceType({
            ...typeConfig,
            interfaces: () =>
              typeConfig.interfaces.map((iface) =>
                replaceType(iface, newTypes, typeConfig.name, subGraphs),
              ),
            fields: PothosSubGraphPlugin.filterFields(type, newTypes, subGraphs, mode),
          }),
        );
      } else if (type instanceof GraphQLUnionType) {
        const typeConfig = type.toConfig();
        newTypes.set(
          type.name,
          new GraphQLUnionType({
            ...typeConfig,
            types: () =>
              typeConfig.types.map((member) =>
                replaceType(member, newTypes, typeConfig.name, subGraphs),
              ),
          }),
        );
      } else if (type instanceof GraphQLInputObjectType) {
        const typeConfig = type.toConfig();
        newTypes.set(
          type.name,
          new GraphQLInputObjectType({
            ...typeConfig,
            fields: PothosSubGraphPlugin.mapInputFields(type, newTypes, subGraphs, mode),
          }),
        );
      }
    }

    return newTypes;
  }

  static filterFields(
    type: GraphQLInterfaceType | GraphQLObjectType,
    newTypes: Map<string, GraphQLNamedType>,
    subGraphs: string[],
    mode: 'any' | 'all',
  ) {
    const oldFields = type.getFields();

    return () => {
      const newFields: GraphQLFieldConfigMap<unknown, unknown> = {};

      for (const [fieldName, fieldConfig] of Object.entries(oldFields)) {
        const newArguments: GraphQLFieldConfigArgumentMap = {};

        if (
          !matchesSubGraphs(
            (fieldConfig.extensions?.subGraphs as string[] | undefined) ?? [],
            subGraphs,
            mode,
          ) ||
          !newTypes.has(getNamedType(fieldConfig.type).name)
        ) {
          continue;
        }

        for (const argConfig of fieldConfig.args) {
          const argSubGraphs = argConfig.extensions?.subGraphs as string[] | undefined;

          if (argSubGraphs && !matchesSubGraphs(argSubGraphs, subGraphs, mode)) {
            if (isNonNullType(argConfig.type)) {
              throw new PothosSchemaError(
                `argument ${argConfig.name} of ${type.name}.${fieldName} is NonNull and must be in included in all sub-graphs that include ${type.name}.${fieldName}`,
              );
            }

            continue;
          }

          newArguments[argConfig.name] = {
            description: argConfig.description,
            defaultValue: argConfig.defaultValue,
            extensions: argConfig.extensions,
            astNode: argConfig.astNode,
            deprecationReason: argConfig.deprecationReason,
            type: replaceType(
              argConfig.type,
              newTypes,
              `${argConfig.name} argument of ${type.name}.${fieldConfig.name}`,
              subGraphs,
            ),
          };
        }

        newFields[fieldName] = {
          description: fieldConfig.description,
          resolve: fieldConfig.resolve,
          subscribe: fieldConfig.subscribe,
          deprecationReason: fieldConfig.deprecationReason,
          extensions: fieldConfig.extensions,
          astNode: fieldConfig.astNode,
          type: replaceType(
            fieldConfig.type,
            newTypes,
            `${type.name}.${fieldConfig.name}`,
            subGraphs,
          ),
          args: newArguments,
        };
      }

      return newFields;
    };
  }

  static mapInputFields(
    type: GraphQLInputObjectType,
    newTypes: Map<string, GraphQLNamedType>,
    subGraphs: string[],
    mode: 'any' | 'all',
  ) {
    const oldFields = type.getFields();

    return () => {
      const newFields: GraphQLInputFieldConfigMap = {};

      for (const [fieldName, fieldConfig] of Object.entries(oldFields)) {
        const fieldSubGraphs = fieldConfig.extensions?.subGraphs as string[] | undefined;

        if (fieldSubGraphs && !matchesSubGraphs(fieldSubGraphs, subGraphs, mode)) {
          if (isNonNullType(fieldConfig.type)) {
            throw new PothosSchemaError(
              `${type.name}.${fieldName} is NonNull and must be in included in all sub-graphs that include ${type.name}`,
            );
          }

          continue;
        }

        newFields[fieldName] = {
          description: fieldConfig.description,
          extensions: fieldConfig.extensions,
          astNode: fieldConfig.astNode,
          defaultValue: fieldConfig.defaultValue,
          deprecationReason: fieldConfig.deprecationReason,
          type: replaceType(
            fieldConfig.type,
            newTypes,
            `${type.name}.${fieldConfig.name}`,
            subGraphs,
          ),
        };
      }

      return newFields;
    };
  }

  override afterBuild(schema: GraphQLSchema) {
    if (this.options.subGraph) {
      return PothosSubGraphPlugin.createSubGraph(schema, this.options.subGraph, this.builder);
    }

    return schema;
  }

  override onTypeConfig(typeConfig: PothosTypeConfig) {
    return {
      ...typeConfig,
      extensions: {
        ...typeConfig.extensions,
        subGraphs:
          typeConfig.pothosOptions.subGraphs ??
          this.builder.options.subGraphs?.defaultForTypes ??
          [],
      },
    };
  }

  override onInputFieldConfig(fieldConfig: PothosInputFieldConfig<Types>) {
    if (fieldConfig.pothosOptions.subGraphs) {
      return {
        ...fieldConfig,
        extensions: {
          ...fieldConfig.extensions,
          subGraphs: fieldConfig.pothosOptions.subGraphs,
        },
      };
    }

    return fieldConfig;
  }

  override onOutputFieldConfig(fieldConfig: PothosOutputFieldConfig<Types>) {
    const typeConfig = this.buildCache.getTypeConfig(fieldConfig.parentType);

    if (typeConfig.graphqlKind !== 'Interface' && typeConfig.graphqlKind !== 'Object') {
      return fieldConfig;
    }

    let subGraphs: Types['SubGraphs'][] = [];

    if (fieldConfig.pothosOptions.subGraphs) {
      subGraphs = fieldConfig.pothosOptions.subGraphs;
    } else if (typeConfig.pothosOptions.defaultSubGraphsForFields) {
      subGraphs = typeConfig.pothosOptions.defaultSubGraphsForFields;
    } else if (this.builder.options.subGraphs?.fieldsInheritFromTypes) {
      subGraphs = (typeConfig.extensions?.subGraphs as Types['SubGraphs'][]) || [];
    } else if (this.builder.options.subGraphs?.defaultForFields) {
      subGraphs = this.builder.options.subGraphs?.defaultForFields;
    }

    return {
      ...fieldConfig,
      extensions: {
        ...fieldConfig.extensions,
        subGraphs,
      },
    };
  }
}

SchemaBuilder.registerPlugin(pluginName, PothosSubGraphPlugin);