selection set decorator

selection-set.decorator.ts

/**
 * EXPERIMENTAL auto populator for mongo using graphql selections set
 */

import { createParamDecorator, ExecutionContext } from "@nestjs/common";
import { Model, DocumentQuery, Document, QueryPopulateOptions } from "mongoose";

import {
  FieldNode,
  ArgumentNode,
  GraphQLResolveInfo,
  FragmentSpreadNode,
  InlineFragmentNode,
  FragmentDefinitionNode,
} from "graphql";

interface Selections {
  [field: string]: {
    arguments: readonly ArgumentNode[] | undefined;
    selections: Selections;
    typeCondition?: string;
  };
}

function getPopulateOptions<TResult, T extends Document>(
  selectionSet: Selections,
  model: Model<T>,
) {
  const select = new Set<string>();
  const populate = new Array<QueryPopulateOptions>();
  // ensure we select all keys of the selection set
  if (selectionSet)
    Object.keys(selectionSet).forEach(field => select.add(field));

  const visitPopulate = (
    subSelections: Selections,
    field: string,
    modelName: string,
  ): QueryPopulateOptions => {
    if (!subSelections[field].selections) {
      return { path: field };
    } else {
      const refFields = Object.keys(subSelections[field].selections)
        .map(selection => {
          const subSchema = model.model(
            subSelections[field].selections[selection].typeCondition ??
              modelName,
          ).schema;
          return {
            field: selection,
            schema: subSchema,
            pathOptions:
              subSchema.path(selection)?.options ||
              subSchema.virtualpath(selection)?.options,
          };
        })
        .filter(s => s.pathOptions?.ref);

      const options: {
        limit: number | undefined;
      } = { limit: undefined };

      // handle limit argument for subfield
      // eg ... { members(limit: 10) }
      const limitArg = subSelections[field].arguments?.find(
        arg => arg.name.value === "limit",
      );

      if (limitArg?.value.kind === "IntValue") {
        options.limit = Number(limitArg.value.value);
      }

      return {
        path: field,
        populate: refFields.map(sub =>
          visitPopulate(
            subSelections[field].selections,
            sub.field,
            sub.pathOptions.ref,
          ),
        ),
        options,
        select: [
          ...refFields.map(r => r.pathOptions.localField),
          ...Object.keys(subSelections[field].selections),
        ].filter(sel => model.model(modelName).schema.paths[sel]),
      };
    }
  };

  Object.keys(selectionSet).forEach(field => {
    const pathOptions =
      model.schema.path(field)?.options ||
      model.schema.virtualpath(field)?.options;
    // ref fields are "virtual"
    if (pathOptions?.ref) {
      // visit each sub selection set to build a tree of fields we need to populate
      const populateOptions = visitPopulate(
        selectionSet,
        field,
        pathOptions.ref,
      );

      // ensure we select the localField, otherwise population won't work
      if (pathOptions.localField) {
        select.add(pathOptions.localField);
      }
      populate.push(populateOptions);
    }
  });

  return { select, populate };
}

function applySelectionSet<TResult, T extends Document>(
  selectionSet: Selections,
  query: DocumentQuery<TResult, T>,
  model: Model<T>,
): void {
  const options = getPopulateOptions(selectionSet, model);

  options.select.forEach(field => query.select(field));
  options.populate.forEach(options => query.populate(options));
}

type FieldNodeType = FieldNode | FragmentSpreadNode | InlineFragmentNode;

function getSelectionSet(
  node: FieldNode,
  fragments: { [key: string]: FragmentDefinitionNode },
) {
  const fieldTreeFromAST = (
    asts: readonly FieldNodeType[],
    fragments: { [key: string]: FragmentDefinitionNode },
    init: Selections,
    typeCondition?: string,
  ) => {
    return asts.reduce(function(tree, val) {
      if (val.kind === "Field") {
        tree[val.name.value] = tree[val.name.value] || {
          selections: {},
          typeCondition,
          arguments: [],
        };
        if (val.selectionSet) {
          fieldTreeFromAST(
            val.selectionSet.selections,
            fragments,
            tree[val.name.value].selections,
          );
        } else {
          tree[val.name.value].arguments = val.arguments;
        }
      } else if (val.kind === "FragmentSpread") {
        fieldTreeFromAST(
          fragments[val.name.value].selectionSet.selections,
          fragments,
          tree,
          typeCondition,
        );
      } else if (val.kind === "InlineFragment") {
        fieldTreeFromAST(
          val.selectionSet.selections,
          fragments,
          tree,
          val.typeCondition?.name.value,
        );
      } // else ignore
      return tree;
    }, init);
  };

  const selections = node.selectionSet?.selections
    ? fieldTreeFromAST(
        node.selectionSet?.selections as FieldNodeType[],
        fragments,
        {},
      )
    : {};

  return selections;
}

export function fromGqlInfo(
  info: GraphQLResolveInfo,
  filter?: (selections: Selections) => Selections,
): SelectionSet {
  let selections = getSelectionSet(info.fieldNodes[0], info.fragments);
  if (filter) {
    selections = filter(selections);
  }

  return {
    selections,
    populatedQuery: (query, schema) => {
      applySelectionSet(selections, query, schema);
      return query;
    },
    getPopulateOptions: schema => {
      return getPopulateOptions(selections, schema);
    },
  };
}

/**
 * Returns an interface that can be used to populate a mongodb query
 * using the selection set of the currently executed GraphQL query
 */
export const SelectionSet = createParamDecorator(
  (
    filter: (selections: Selections) => Selections,
    context: ExecutionContext,
  ): SelectionSet => {
    const [, , , info] = context.getArgs();
    return fromGqlInfo(info, filter);
  },
);

export interface SelectionSet {
  populatedQuery<TResult, T extends Document>(
    query: DocumentQuery<TResult, T>,
    model: Model<any>,
  ): DocumentQuery<TResult, T>;
  getPopulateOptions<T extends Document>(
    model: Model<T>,
  ): { select: Set<string>; populate: QueryPopulateOptions[] };
  selections: Selections;
}

Usage:

  @Query(() => Foo)
  async getFoo(
    @SelectionSet() selections: SelectionSet,
  ) {
    this.someService.findById(someId, selections);
  }

  async findById(id: string, selectionSet?: SelectionSet) {
    const query = this.model.findById(id);
    selectionSet?.populatedQuery(query, this.model);

    return query;
  }