経験は何よりも饒舌

10年後に真価を発揮するかもしれないブログ 

graphql-validation-complexity から学ぶGraphQLのAST走査

オライリー『初めてのGraphQL』の「7.3.4 クエリの複雑さ制限」で、graphql-validation-complexity というライブラリが紹介されていて、面白そうだったのでのぞいてみた。

github.com

graphql-validation-complexityは、クエリに対して複雑度を計算し、その複雑度に制限を設けることができる。
デフォルトでそれぞれのフィールドに値を設定し、リストが入れ子になるたびに値を10倍して、複雑度を計算している。
テストコードで例を示すと、以下のように複雑度が計算される。

graphql-validation-complexity/ComplexityVisitor.test.js at master · 4Catalyzer/graphql-validation-complexity · GitHub

query {
  item {
    name # 複雑度1
    item {
      name # 複雑度10
    }
  }
  list {
    item {
      name # 複雑度10
    }
    list {
      name # 複雑度100
    }
  }
  nonNullItem {
    name # 複雑度1
  }
  nonNullList {
    name # 複雑度1
  }
} # 総複雑度123

この複雑度の計算の解釈が正しいことは、以下のテストが通ったことからも保証できる。

describe('simple queries', () => {
    it('should calculate the correct cost', () => {
      const cost = checkCost(`
        query {
          list {
            item {
              name
            }
            list {
              name
            }
          }
        }
      `);

      expect(cost).toBe(110);
    });
});

実際にどう使うかというと、例えばApolloを使う場合は、以下のようにvalidationRulesフィールドに設定する。

const apolloServer = new ApolloServer({
  schema,
  validationRules: [createComplexityLimitRule(1000)],
});

では、このcreateComplexityLimitRuleの中身を中身をみていく。

createComplexityLimitRule関数は、src/index.js##L14-L17でexportされており、引数に設定値をmaxCostとしてとっている。
graphql-validation-complexity/index.js at 1be812179e5da7d850b76441b135ed178ce769f5 · 4Catalyzer/graphql-validation-complexity · GitHub

export function createComplexityLimitRule(
  maxCost,
  { onCost, createError, formatErrorMessage, ...options } = {},
) {

costが設定値のmaxCostを超えていたら、エラーが吐かれるようになっている。

graphql-validation-complexity/index.js at 1be812179e5da7d850b76441b135ed178ce769f5 · 4Catalyzer/graphql-validation-complexity · GitHub

if (cost > maxCost) {
        context.reportError(
         createError
            ? createError(cost, node)
            : new GraphQLError(formatErrorMessage(cost), [node]),
        );
}

では、costはどのように求められているのかを調べると、以下のコードが目についた。

Document: {
   enter(node) {
      visit(node, visitWithTypeInfo(typeInfo, visitor));
    },
    leave(node) {
      const cost = visitor.getCost();

ここのvisit関数はfrom 'graphql';でインポートされているため、GraphQLのドキュメントを見てみると、以下のような説明があった。

graphql.org

visit() will walk through an AST using a depth first traversal, calling the visitor's enter function at each node in the traversal, and calling the leave function after visiting that node and all of its child nodes.

By returning different values from the enter and leave functions, the behavior of the visitor can be altered, including skipping over a sub-tree of the AST (by returning false), editing the AST by returning a value or null to remove the value, or to stop the whole traversal by returning BREAK.

When using visit() to edit an AST, the original AST will not be modified, and a new version of the AST with the changes applied will be returned from the visit function.

visit()は、深さ優先のトラバーサルを用いてASTを走査し、トラバーサル内の各ノードでビジターのenter関数を呼び出し、そのノードとその子ノードすべてを訪れた後にleave関数を呼び出します。

enter関数とleave関数から異なる値を返すことで、ビジターの動作を変更することができます。例えば、falseを返すことでASTのサブツリーをスキップしたり、値を返すことでASTを編集したり、nullを返すことで値を削除したり、BREAKを返すことでトラバーサル全体を停止したりすることができます。

visit()を使ってASTを編集する場合、元のASTは変更されず、変更が適用された新しいバージョンのASTがvisit関数から返されます。

このenterにあるvisit()によりAST走査がはじまるらしい。
テストコードを見てみると、const ast = parse(query);で明示的にqueryがパースされていたので、enter(node)nodeはASTであることがわかる。

graphql-validation-complexity/ComplexityVisitor.test.js at 1be812179e5da7d850b76441b135ed178ce769f5 · 4Catalyzer/graphql-validation-complexity · GitHub

describe('ComplexityVisitor', () => {
  const typeInfo = new TypeInfo(schema);

  function checkCost(query) {
    const ast = parse(query);
    const context = new ValidationContext(schema, ast, typeInfo);
    const visitor = new ComplexityVisitor(context, {});

    deepFreeze(ast); // ensure we aren't mutating accidentally

    visit(ast, visitWithTypeInfo(typeInfo, visitor));

    return visitor.getCost();
  }

  describe('simple queries', () => {
    it('should calculate the correct cost', () => {
      const cost = checkCost(`
        query {
          item {
            name
            item {
              name
            }
           ...
      `);

      expect(cost).toBe(123);
    });
  });

astvisitWithTypeInfo(typeInfo, visitor)デバッグすると以下のようであった。

console.log
    {
      kind: 'Document',
      definitions: [
        {
          kind: 'OperationDefinition',
          operation: 'query',
          name: undefined,
          variableDefinitions: [],
          directives: [],
          selectionSet: [Object],
          loc: [Object]
        }
      ],
      loc: { start: 0, end: 370 }
    }

console.log
    { enter: [Function: enter], leave: [Function: leave] }

astをさらに詳しくデバッグすると、queryのフィールドはSelectionSetで分割されていることがわかった。
console.log(ast.definitions[0].selectionSet);

console.log
    {
      kind: 'SelectionSet',
      selections: [
        {
          kind: 'Field',
          alias: undefined,
          name: [Object],
          arguments: [],
          directives: [],
          selectionSet: [Object],
          loc: [Object]
        },
        {
          kind: 'Field',
          alias: undefined,
          name: [Object],
          arguments: [],
          directives: [],
          selectionSet: [Object],
          loc: [Object]
        },
        {
          kind: 'Field',
          alias: undefined,
          name: [Object],
          arguments: [],
          directives: [],
          selectionSet: [Object],
          loc: [Object]
        },
        {
          kind: 'Field',
          alias: undefined,
          name: [Object],
          arguments: [],
          directives: [],
          selectionSet: [Object],
          loc: [Object]
        }
      ],
      loc: { start: 15, end: 363 }
    }

では、実際に複雑度を計算している箇所をみていく。

const cost = visitor.getCost();getCost関数はsrc/ComplexityVisitor.js#L221-L223で定義されていて、

graphql-validation-complexity/ComplexityVisitor.js at 1be812179e5da7d850b76441b135ed178ce769f5 · 4Catalyzer/graphql-validation-complexity · GitHub

  getCost() {
    return this.cost;
  }

119行目から127行目あたりで実際にcostの計算がされていた。

graphql-validation-complexity/ComplexityVisitor.js at 1be812179e5da7d850b76441b135ed178ce769f5 · 4Catalyzer/graphql-validation-complexity · GitHub

 enterField() {
    this.costFactor *= this.getFieldCostFactor();

    this.cost += this.costFactor * this.getFieldCost();
  }

  leaveField() {
    this.costFactor /= this.getFieldCostFactor();
  }

astselectionSetがいじられている箇所があったことから、selectionSetを使って複雑度に重みを与えていることが想像できる。

graphql-validation-complexity/ComplexityVisitor.js at 1be812179e5da7d850b76441b135ed178ce769f5 · 4Catalyzer/graphql-validation-complexity · GitHub

    this.SelectionSet = this.flattenFragmentSpreads;
  }

  flattenFragmentSpreads(selectionSet) {
    const nextSelections = selectionSet.selections.flatMap((node) => {
  ....

複雑度が計算されて値に応じてエラーが返る流れはなんとなく掴めた。
が、graphql-validation-complexityとGraphqlのvisit()がどのように噛み合ってるのかがいまいちよくわからない。
ComplexityVisitorクラスのコンストラクタにあるthis.Fieldenter/leaveが発火しているっぽいが、そもそもthis.Fieldが何を指しているのかわかっていない。

    this.Field = {
      enter: this.enterField,
      leave: this.leaveField,
    };

さっきみたように、visitWithTypeInfoによってenter/leaveが作られているから、visitWithTypeInfoを調べればわかるかもしれないが、ドキュメントを調べても出てこないから諦める。

const visitor = new ComplexityVisitor(context, options);
const typeInfo = context._typeInfo || new TypeInfo(context.getSchema());

console.log(visitWithTypeInfo(typeInfo, visitor))
// { enter: [Function: enter], leave: [Function: leave] }

にしてもいろんなところでASTが使われている。