GraphQL DataLoaderライブラリ dataloden の使い方

こんにちは、spotチームの名嘉眞です。spotチームはキャンプ場検索サービス(hinata spot)を開発しております。私はspotチームのバックエンド担当として日々Goを書いてます。

spotチームでは、GraphQLのDataLoaderライブラリとしてdataloaden を使用しています。今回は、dataloadenの使い方について記載しました。

GraphQLでのデータ取得について

GraphQLは、1回のクエリで必要なフィールドだけ取得できる仕組みになっていて便利です。ただ、よく考えずにバックエンドの実装を行うと、クエリによっては使わないフィールドの値も必ず取得するような実装になったり、DBからデータを取得する場合は、必ずJOINするような実装になる可能性があります。

動的にJOINしたり、バックエンド側でリクエストに応じて取得したいフィールドに合ったデータを生成することは出来ると思いますが、バックエンドの実装が少し複雑になるかと思います。

# schema.graphql

type Contract {
  """
  id: ID!
  """
  name: String!
  """
  payment: Payment!  # paymentフィールド(支払い方法)はcontractとは別テーブルのデータとする
}

query {
  getContract(contract: { id: "1" }) {
    id
    name
    payment # このpaymentフィールドはクライアント側からのリクエストで必ず指定されるわけでは無い
}

このような問題を解決する方法として、対象フィールドのGraphQL resolver(GraphQLバックエンド関数)を別で作成し、GraphQLサーバがレスポンスを組み立てるという方法があります。

この方法を使うと各resolver自体の実装は対象のフィールドのデータ取得だけでシンプルになり、バックエンド側で複雑な実装はせずに済みます。また、クエリに応じて動的にフィールドのデータを取得することが出来ます。

resolverを別で作成する方法に関して詳しくは記載しませんが、spotチームだと gqlgen というライブラリを使用しており、以下のようにgqlgenの設定ファイル(gqlgen.yml)を変更することで作成出来ます。

# gqlgen.yml

models:
  Contract:
    fields:
      payment:
        resolver: true

次のようなインターフェースが作成されます。contract resolver というインターフェースにPaymentというメソッドが追加されました。このインターフェースを満たす実装を行うことでデータの取得ができます。

// generated.go
type ContractResolver interface {
    Payment(ctx context.Context, obj *models.Contract) (*models.Payment, error)
}

ただ、注意しないといけない点があり、resolverを別にしただけだと、N＋1問題が発生します。

query {
  getContracts {
    id
    name
    payment # paymentフィールドは別resolverで取得することになっている
}

上記のクエリは、契約情報の一覧を取得するクエリだとします。このクエリは契約名や、契約の支払い情報を複数取得できます。

paymentフィールドは別resolverで取得するようにした場合、バックエンド側でSQLは下記のように発行されます。

# 契約一覧を取得 
SELECT id, name FROM contracts;
# 上記のSQLで取得された契約情報の数だけ、下記のSQLも発行される
SELECT * FROM payments WHERE contract_id = ?
SELECT * FROM payments WHERE contract_id = ?
SELECT * FROM payments WHERE contract_id = ?

上記のようなN＋1問題を解決する方法が、DataLoaderを使ったデータ取得です。 DataLoaderを使うと上記のSQLを下記のように変更することが出来ます。

# 契約一覧を取得 
SELECT id, name FROM contracts;
# paymentはまとめて取得
SELECT * FROM payments WHERE contract_id IN (?, ?, ?)

DataLoaderとは

DataLoaderとは、データ取得をバッチ化するためのライブラリです。
具体的には、複数のデータを取得するリクエストが発生した時に、定義した待ち時間(Wait)のあいだ、取得したいデータのキー(オブジェクトのIDなど)を蓄積し、まとめてデータ取得処理を実行する仕組みです。データ取得処理の回数を減らすことができます。
取得したいデータのキーがバッチ化する数という感じになります。バッチ化する数(MaxBatch)と、待ち時間(Wait) は任意に指定することができます。
待ち時間(Wait) の値が大きくなるとバッチ化できる範囲が広がりますが、その分レスポンスタイムが遅くなるおそれがあります。

dataloadenの使い方

dataloadenは、GoのDataLoader生成のためのライブラリです。gqlgenの作者が作成しています。また、gqlgenのDataLoaderの説明にも、dataloadenを使ったDataLoaderの設定方法が記載されています。

使い方の例として、これまでの説明でもあったpayment(支払い情報)を取得するDataLoaderを、dataloadenを使って作成します。

まず、DataLoaderファイルを生成したいディレクトリで下記のコマンドを実行します。DataLoaderの名前、 int、stringなどのデータを取得するためのキーの型、取得するデータの型を指定してコマンドを実行します。

# dataloaden `DataLoaderの名前` `keytype(int, stringなど)` `取得するデータの型`
dataloaden PaymentLoader string '*github.com/vivit/spot/model.Payment'

上記のコマンドを実行すると、paymentloader_gen.go というファイルが自動生成されます。ファイルの中には、下記のように PaymentLoader という構造体が定義されています。 PaymentLoader のfetchフィールドは関数型で、引数に keys []string 、戻り値に []*model.Payment が定義されています。先ほど実行したdataloadenのコマンドで指定した型です。

// PaymentLoader batches and caches requests
type PaymentLoader struct {
    // this method provides the data for the loader
    fetch func(keys []string) ([]*model.Payment, []error)

    // how long to done before sending a batch
    wait time.Duration

    // this will limit the maximum number of keys to send in one batch, 0 = no limit
    maxBatch int

    // INTERNAL

    // lazily created cache
    cache map[string]*model.Payment

    // the current batch. keys will continue to be collected until timeout is hit,
    // then everything will be sent to the fetch method and out to the listeners
    batch *paymentLoaderBatch

    // mutex to prevent races
    mu sync.Mutex
}

次にdataloadenで生成した仕組みを使うプログラムを作成します。(名前はなんでも大丈夫です。) GraphQLサーバーのミドルウェアとして使うための関数(PaymentLoaderMiddleware) と、resolverからDataLoaderを使うための関数(Payment) を定義します。

// paymentloader.go
package loaders

import (
    // 省略
)
const paymentLoaderKey = "paymentLoader"

func PaymentLoaderMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // dataloadenで生成されたファイルに定義されているNewPaymentLoaderを使います。
        // NewPaymentLoaderの引数にはこれまた生成されたファイルに定義されているPaymentLoaderConfigを設定します。
        // PaymentLoaderConfigにMaxBatch, Wait, データをfeatchする関数を定義します。
        loader := NewPaymentLoader(PaymentLoaderConfig{
            MaxBatch: 30,　                  // 最大MaxBatchの数だけ、keysに値が設定される(contractID)
            Wait:     10 * time.Millisecond, // 10ms待機する。待機した時間の間で設定されたkeysをもとに後述の処理を行う
            Fetch: func(keys []string) ([]*model.Payment, []error) {
                ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
                defer cancel()

                // DB からpaymentを取得する、この関数は、keysの値で下記のようなSQLを実行する
                // SELECT * FROM payments WHERE contract_id IN (keys[0], keys[1], keys[2])
                reply, err := repogitory.GetPayments(ctx, keys)
                if err != nil {
                    return nil, []error{err}
                }

                payments := map[string]*model.Payment{}
                for _, p := range reply {
                    payments[p.contractID] = p
                }

                // keyの値にあったpaymentを、keysの並び順にしたpaymentのスライスを作る
                result := make([]*model.Payment, len(keys))
                for i, key := range keys {
                    result[i] = payments[key]
                }

                return result, nil
            },
        })
        ctx := context.WithValue(r.Context(), paymentLoaderKey, loader)
        r = r.WithContext(ctx)
        next.ServeHTTP(w, r)
    })
}

func Payment(ctx context.Context, id string) (*model.Payment, error) {
    v := ctx.Value(paymentLoaderKey)
    loader, ok := v.(*PaymentLoader)

    if !ok {
        return nil, xerrors.New("failed to get loader from current context")
    }

    return loader.Load(id)
}

そして上記で作成した関数を使う部分です。作成した関数、PaymentMiddlewareをミドルウェアとして、http.Handleで定義して使います。 ※このあたりはいろいろ定義の仕方があると思うので簡略した書き方にしています。

// handler.go
http.Handle("/graphql", loader.PaymentMiddleware(http.Handler型が入る))

もう一つ作成した関数、PaymentはcontractResolverからの呼び出して使います。記事の最初の方で記載した、gqlgenで自動生成されたインターフェースを満たすように実装します。

// contract_resolver.go
func (*contractResolver) Payment(ctx context.Context, s *models.Contract) (*models.Payment, error) {
    p, err := loaders.Payment(ctx, s.GetId())
    
    return &model.Payment{Payment: p}, err
}

以上で、設定は完了です。以下のようなクエリを実行したログを確認して、contractを検索するSQLとcontractに紐づくpaymentを検索するSQLが確認できれば成功です。

query {
  getContracts {
    id
    name
    payment
  }
}

paymentフィールドを指定した場合だけpaymentを取得するSQLが実行されるのと、その実行回数も親のオブジェクトであるcontractの件数以下になっていると思います。

まとめ

GraphQLの良さを最大限に活かすには、DataLoaderは必要不可欠かなと思いました。DataLoaderを使うことで、動的に取得したいデータだけを取得できますし、バックエンド側の実装もシンプルになりました。

まだ、datalodenで生成されたファイルの内容を全て理解できたわけでは無いので、これからもいろいろ使ってみるのと、spotサービスでDataLoaderを使えていないフィールドもあるので、その改善も進めていこうと思います。