GraphQL RubyのDataloaderを使ってみる

GraphQL Rubyを使っているときにSQLのN+1問題を解決するためにバッチロードを利用すると思います。
有名なgemとしては [GraphQL::Batch](https://github.com/Shopify/graphql-batch) や [BatchLoader](https://github.com/exAspArk/batch-loader)、 [Dataloader](https://github.com/sheerun/dataloader)などがあります。
個人的には使ったことがあるのは[Dataloader](https://github.com/sheerun/dataloader)なのですが、2018年から更新されていないようなのと、去年GraphQL Ruby自体にDataloaderが入ったので今回はこちらを触ってみました。

## GraphQL::Dataloader の概要

ドキュメントはこちらです。

https://graphql-ruby.org/dataloader/overview.html

GraphQL::Dataloaderはデータベースアクセスを効率的に行うためのツールで、RubyのFiberを使っており、Ruby 3のノンブロッキングFiberもサポートしているようです。影響を受けたものとしては以下2つが挙げられています。

- https://github.com/bessey/graphql-fiber-test/tree/no-gem-changes
- https://github.com/shopify/graphql-batch

[Dataloader vs. GraphQL-Batch](https://graphql-ruby.org/dataloader/adopting.html)でも、GraphQL-Batchなどの他のローダーはPromiseを使っている一方で、このDataloaderはFiberを使っているということが特徴として挙げられています。committerのRobert Mosolgoによると、Promiseを使うと複雑になりやすいところを、Rubyにもともと備わっている機能であるFiberを使うことで、他に何も使わずとも並列I/Oが使えるというところで選んだようです。

https://rmosolgo.github.io/ruby/graphql/2021/02/13/graphql-dataloader.html

## 前提

- graphql v1.13.2
- rails v7.0.0

## 前準備

サンプルのデータをもとに検証してみます。
シンプルにUser, Aritcle, Likeの3つのテーブルを作ります。

```rb
class User < ApplicationRecord
  has_many :articles, foreign_key: 'author_id'
end

class Article < ApplicationRecord
  belongs_to :author, class_name: 'User'
  has_many :likes
end

class Like < ApplicationRecord
  belongs_to :article, validate: true
  belongs_to :user, validate: true
end
```

フィールドはこんな感じです。

```rb
module Types
  class UserType < Types::BaseObject
    field :id, ID, null: false
    field :name, String, null: false
    field :email, String, null: false
    field :articles, Types::ArticleType.connection_type, null: false
  end
end

module Types
  class ArticleType < Types::BaseObject
    field :id, ID, null: false
    field :title, String, null: false
    field :body, String, null: false
    field :author, Types::UserType, null: false
    field :liked, Types::LikeConnection, null: false, method: :likes
  end
end

module Types
  class LikeType < Types::BaseObject
    field :id, ID, null: false
    field :user, Types::UserType, null: false
    field :article, Types::ArticleType, null: false
    field :liked_at, Int, null: false, method: :created_at
  end
end
```

今回欲しいクエリはこんな感じで、現在のユーザーが投稿した記事一覧とその記事にLikeしたユーザーの一覧を取得します。一部ページネーションでConnectionを使っています。

```graphql
query {
  currentUser {
    name
    articles {
      edges {
        node {
          title
          body
          liked {
            count
            edges {
              node {
                user {
                  name
                }
              }
            }
          }
        }
      }
    }
  }
}
```

<details><summary>ダミーデータも生成しておく</summary><div>

取得するデータを生成しておきます。
クエリの取得結果はこのようになりました。

```json
{
  "data": {
    "currentUser": {
      "name": "test_user1",
      "articles": {
        "edges": [
          {
            "node": {
              "title": "first article",
              "body": "first article body",
              "liked": {
                "count": 4,
                "edges": [
                  {
                    "node": {
                      "user": {
                        "name": "test_user2"
                      }
                    }
                  },
                  {
                    "node": {
                      "user": {
                        "name": "test_user3"
                      }
                    }
                  },
                  {
                    "node": {
                      "user": {
                        "name": "test_user4"
                      }
                    }
                  },
                  {
                    "node": {
                      "user": {
                        "name": "test_user5"
                      }
                    }
                  }
                ]
              }
            }
          },
          {
            "node": {
              "title": "second article",
              "body": "second article body",
              "liked": {
                "count": 1,
                "edges": [
                  {
                    "node": {
                      "user": {
                        "name": "test_user2"
                      }
                    }
                  }
                ]
              }
            }
          },
          {
            "node": {
              "title": "third article",
              "body": "third article body",
              "liked": {
                "count": 2,
                "edges": [
                  {
                    "node": {
                      "user": {
                        "name": "test_user3"
                      }
                    }
                  },
                  {
                    "node": {
                      "user": {
                        "name": "test_user4"
                      }
                    }
                  }
                ]
              }
            }
          },
          {
            "node": {
              "title": "forth article",
              "body": "forth article body",
              "liked": {
                "count": 0,
                "edges": []
              }
            }
          },
          {
            "node": {
              "title": "fifth article",
              "body": "fifth article body",
              "liked": {
                "count": 1,
                "edges": [
                  {
                    "node": {
                      "user": {
                        "name": "test_user5"
                      }
                    }
                  }
                ]
              }
            }
          }
        ]
      }
    }
  }
}
```

</div></details>

## まずクエリを投げてみる

このまま素直にクエリを投げると、素直にN+1が発生してくれるのでこれを解消していきます。

```sql
  User Load (0.1ms)  SELECT "users".* FROM "users" WHERE "users"."id" = ? LIMIT ?  [["id", 1], ["LIMIT", 1]]

  Article Load (0.2ms)  SELECT "articles".* FROM "articles" WHERE "articles"."author_id" = ? LIMIT ?  [["author_id", 1], ["LIMIT", 20]]
  Like Load (0.1ms)  SELECT "likes".* FROM "likes" WHERE "likes"."article_id" = ? LIMIT ?  [["article_id", 1], ["LIMIT", 20]]
  Like Load (0.1ms)  SELECT "likes".* FROM "likes" WHERE "likes"."article_id" = ? LIMIT ?  [["article_id", 2], ["LIMIT", 20]]
  Like Load (0.2ms)  SELECT "likes".* FROM "likes" WHERE "likes"."article_id" = ? LIMIT ?  [["article_id", 3], ["LIMIT", 20]]
  Like Load (0.2ms)  SELECT "likes".* FROM "likes" WHERE "likes"."article_id" = ? LIMIT ?  [["article_id", 4], ["LIMIT", 20]]
  Like Load (0.5ms)  SELECT "likes".* FROM "likes" WHERE "likes"."article_id" = ? LIMIT ?  [["article_id", 5], ["LIMIT", 20]]

  User Load (0.1ms)  SELECT "users".* FROM "users" WHERE "users"."id" = ? LIMIT ?  [["id", 2], ["LIMIT", 1]]
  User Load (0.2ms)  SELECT "users".* FROM "users" WHERE "users"."id" = ? LIMIT ?  [["id", 3], ["LIMIT", 1]]
  User Load (0.3ms)  SELECT "users".* FROM "users" WHERE "users"."id" = ? LIMIT ?  [["id", 4], ["LIMIT", 1]]
  User Load (0.2ms)  SELECT "users".* FROM "users" WHERE "users"."id" = ? LIMIT ?  [["id", 5], ["LIMIT", 1]]
  User Load (0.2ms)  SELECT "users".* FROM "users" WHERE "users"."id" = ? LIMIT ?  [["id", 2], ["LIMIT", 1]]
  User Load (0.2ms)  SELECT "users".* FROM "users" WHERE "users"."id" = ? LIMIT ?  [["id", 3], ["LIMIT", 1]]
  User Load (0.7ms)  SELECT "users".* FROM "users" WHERE "users"."id" = ? LIMIT ?  [["id", 4], ["LIMIT", 1]]
  User Load (0.2ms)  SELECT "users".* FROM "users" WHERE "users"."id" = ? LIMIT ?  [["id", 5], ["LIMIT", 1]]
```

## Dataloaderの導入

新規プロジェクトで`rails generate graphql:install`した場合はすでに有効化されていますが、schemaに以下の行を追加するだけです。

```diff_ruby
class MySchema < GraphQL::Schema
  # ...
+  use GraphQL::Dataloader
end
```

### シングルレコードを取得

まずはbelongs_toの関係になっている、`like.user`部分のシングルレコードの取得部分からから準備します。

実装するものとしては、

- `Source`
- resolverのメソッド

の２つです。

說明する前に先にコードを載せます。
こちらはほぼドキュメントに乗っているサンプルそのままです。

```rb:app/graphql/sources/user_by_id.rb
 module Sources
  class UserById < GraphQL::Dataloader::Source
    def initialize
      @model_class = ::User
    end

    def fetch(ids)
      records = @model_class.where(id: ids)
      ids.map { |id| records.find { |r| r.id == id.to_i } }
    end
  end
end
```

```diff_ruby:app/graphql/types/like_type.rb
module Types
  class LikeType < Types::BaseObject

    field :user, Types::UserType, null: false

+    def user
+      dataloader.with(::Sources::UserById).load(object.user_id)
+    end
  end
end
```

これによって、Userの取得が以下のように1つのクエリで完結します。

```sql
  User Load (0.2ms)  SELECT "users".* FROM "users" WHERE "users"."id" IN (?, ?, ?, ?)  [["id", 2], ["id", 3], ["id", 4], ["id", 5]]
```

### 解説

まずSourceについてです。

https://graphql-ruby.org/dataloader/sources.html

Sourceはバッチロードする処理を記述するクラスで、`GraphQL::Dataloader::Source`を継承し、`fetch`メソッドを実装します。
このSourceは`GraphQL::Dataloader`によってインスタンス生成され、`fetch`メソッドが呼ばれるのですが、引数としては取得するためのkey(今回でいうと`ids`)が渡されます。今回はこのkeyからUserのデータを取得します。
返り値は渡された引数のkeyと同じ順番でオブジェクトを返す必要があります。

そしてresolver側では以下のように取得します。

```rb
def user
  dataloader.with(::Sources::UserById).load(object.user_id)
end
```

loadで裏側でFiberのキューに入れて遅延ロードを行っているようです。この例でいくと`object.user_id`がSourceのfetchメソッドの引数として渡されます。

### マルチレコードを取得

こちらも先にコードを載せます。
基本的にはシングルレコードとあまり変わりません。

```ruby:app/graphql/sources/likes_by_user_id.rb
module Sources
  class LikesByUserId < GraphQL::Dataloader::Source
    def initialize
      @model_class = ::Like
    end

    def fetch(keys)
      records = @model_class.where(article_id: keys)
                            .group_by { |record| record.article_id }
      keys.map { |key| records[key] || [] }
    end
  end
end
```

```diff_ruby:app/graphql/types/article_type.rb
module Types
  class ArticleType < Types::BaseObject

    field :liked, Types::LikeConnection, null: false

+    def liked
+      dataloader.with(::Sources::LikesByUserId).load(object.id)
+    end
  end
end
```

少し変わったのは`fetch`メソッドくらいですね。
こちらは1つのkeyに対して複数レコードが返ってくる可能性があるので`group_by`をしています。
返り値は同様にkeyの順番に合わせて配列を返しています。
これだけでマルチレコードも遅延ロードができました。

結果以下のようになり、かなり効率的なクエリになりました。

```sql
  User Load (0.2ms)  SELECT "users".* FROM "users" WHERE "users"."id" = ? LIMIT ?  [["id", 1], ["LIMIT", 1]]
  Article Load (0.1ms)  SELECT "articles".* FROM "articles" WHERE "articles"."author_id" = ? LIMIT ?  [["author_id", 1], ["LIMIT", 20]]
  Like Load (0.1ms)  SELECT "likes".* FROM "likes" WHERE "likes"."article_id" IN (?, ?, ?, ?, ?)  [["article_id", 1], ["article_id", 2], ["article_id", 3], ["article_id", 4], ["article_id", 5]]
  User Load (0.2ms)  SELECT "users".* FROM "users" WHERE "users"."id" IN (?, ?, ?, ?)  [["id", 2], ["id", 3], ["id", 4], ["id", 5]]
```

## まだ分かっていないこと

GraphQL::Batchなどに比べるとGraphQL Dataloaderはまだ導入されたばかりです。そのため、今回はかなりシンプルなものしか試していないので、実際の運用で出てくるような複雑なものまで実用に耐えられるのかは不明です。

また、ドキュメントを見ていると1Source 1クラスで処理は似たものでもクラスが増えていきそうな印象でした。Sourceが増えてきたときに効率的に実装して管理できるかはまだ不明です。

## 最後に

GraphQL::Batchなど他のGemはまだ触っていないのですが、GraphQL Rubyに入っているというだけあって、導入はめちゃくちゃ楽でした。お手軽に使いたいのであればとてもいいと思います。
また他のLoaderも触ってみて比較したいです。