虎の穴開発室ブログ

虎の穴ラボ株式会社所属のエンジニアが書く技術ブログです

MENU

Rust+GraphQLでN+1対策する

技術エントリ

※この記事は予約投稿です。

こんにちは。虎の穴ラボのY.Fです。
この記事は、虎の穴ラボ Advent Calendar 2020 - Qiita12日目の記事になります。
11日目はH.Kさんの Chrome拡張公開方法とChrome拡張として作ったJava APIドキュメントの遷移補助ツールについて についての記事です。ぜひ読んでみてください。

toranoana-lab.hatenablog.com

13日目はいわみーさんの 顧客や非エンジニアに向けてわかりやすい文章を書くときに気をつけていることをまとめてみた の記事です。お楽しみに！

本日の記事は、本ブログで連載していたRust+GraphQLの続きの記事となります。(前回の記事はこちら)

toranoana-lab.hatenablog.com

今回は、関連付けされたデータを取得するときに必ず起きてしまうN＋1問題をどのように防止するか書いていきたいと思います。

(GraphiQLでの実行の様子)

実行の様子
GraphiQLでの実行

前回までの振り返り

前回は、GraphQLで1対多の関連付けがあるデータを取得する方法と、そのときに出力されるSQLについて確認しました。

(クエリ)

query {
  allPosts {
    id
    title
    content
    photos {
      id
      name
      description
      url
    }
  }
}

(実行されるSQL)

[2020-07-30T07:15:45Z INFO  actix_server::builder] Starting 4 workers
[2020-07-30T07:15:45Z INFO  actix_server::builder] Starting "actix-web-service-127.0.0.1:3000" service on 127.0.0.1:3000
[2020-07-30T07:16:04Z DEBUG graphql_sample::db::post_repository] "SELECT \"posts\".\"id\", \"posts\".\"title\", \"posts\".\"content\", \"posts\".\"created_at\", \"posts\".\"updated_at\" FROM \"posts\" WHERE \"posts\".\"id\" = $1 -- binds: [1]"
[2020-07-30T07:16:04Z DEBUG graphql_sample::db::post_repository] "SELECT \"posts\".\"id\", \"posts\".\"title\", \"posts\".\"content\", \"posts\".\"created_at\", \"posts\".\"updated_at\" FROM \"posts\" WHERE \"posts\".\"id\" = $1 -- binds: [1]"
[2020-07-30T07:16:04Z DEBUG graphql_sample::db::post_repository] "SELECT \"posts\".\"id\", \"posts\".\"title\", \"posts\".\"content\", \"posts\".\"created_at\", \"posts\".\"updated_at\" FROM \"posts\" WHERE \"posts\".\"id\" = $1 -- binds: [1]"
[2020-07-30T07:16:04Z DEBUG graphql_sample::db::post_repository] "SELECT \"posts\".\"id\", \"posts\".\"title\", \"posts\".\"content\", \"posts\".\"created_at\", \"posts\".\"updated_at\" FROM \"posts\" WHERE \"posts\".\"id\" = $1 -- binds: [2]"
[2020-07-30T07:16:04Z DEBUG graphql_sample::db::post_repository] "SELECT \"posts\".\"id\", \"posts\".\"title\", \"posts\".\"content\", \"posts\".\"created_at\", \"posts\".\"updated_at\" FROM \"posts\" WHERE \"posts\".\"id\" = $1 -- binds: [2]"
[2020-07-30T07:16:04Z DEBUG graphql_sample::db::post_repository] "SELECT \"posts\".\"id\", \"posts\".\"title\", \"posts\".\"content\", \"posts\".\"created_at\", \"posts\".\"updated_at\" FROM \"posts\" WHERE \"posts\".\"id\" = $1 -- binds: [2]"
[2020-07-30T07:16:04Z DEBUG graphql_sample::db::post_repository] "SELECT \"posts\".\"id\", \"posts\".\"title\", \"posts\".\"content\", \"posts\".\"created_at\", \"posts\".\"updated_at\" FROM \"posts\" WHERE \"posts\".\"id\" = $1 -- binds: [3]"
[2020-07-30T07:16:04Z DEBUG graphql_sample::db::post_repository] "SELECT \"posts\".\"id\", \"posts\".\"title\", \"posts\".\"content\", \"posts\".\"created_at\", \"posts\".\"updated_at\" FROM \"posts\" WHERE \"posts\".\"id\" = $1 -- binds: [3]"
[2020-07-30T07:16:04Z DEBUG graphql_sample::db::post_repository] "SELECT \"posts\".\"id\", \"posts\".\"title\", \"posts\".\"content\", \"posts\".\"created_at\", \"posts\".\"updated_at\" FROM \"posts\" WHERE \"posts\".\"id\" = $1 -- binds: [3]"
[2020-07-30T07:16:04Z DEBUG graphql_sample::db::post_repository] "SELECT \"posts\".\"id\", \"posts\".\"title\", \"posts\".\"content\", \"posts\".\"created_at\", \"posts\".\"updated_at\" FROM \"posts\" WHERE \"posts\".\"id\" = $1 -- binds: [4]"
[2020-07-30T07:16:04Z DEBUG graphql_sample::db::post_repository] "SELECT \"posts\".\"id\", \"posts\".\"title\", \"posts\".\"content\", \"posts\".\"created_at\", \"posts\".\"updated_at\" FROM \"posts\" WHERE \"posts\".\"id\" = $1 -- binds: [4]"
[2020-07-30T07:16:04Z DEBUG graphql_sample::db::post_repository] "SELECT \"posts\".\"id\", \"posts\".\"title\", \"posts\".\"content\", \"posts\".\"created_at\", \"posts\".\"updated_at\" FROM \"posts\" WHERE \"posts\".\"id\" = $1 -- binds: [4]"
[2020-07-30T07:16:04Z INFO  actix_web::middleware::logger] 127.0.0.1:62291 "POST /graphql HTTP/1.1" 200 1512 "http://localhost:3000/graphiql" "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/84.0.4147.89 Safari/537.36" 0.040156

今回の記事では、このレコードの取得ごとに1SQL発行される状態を、SQLのin句を使ってまとめて取得できるように改善してみたいと思います。

必要なライブラリのインストール

基本となるライブラリは前回の記事をご覧ください。
今回は以下を追加します。

(Cargo.toml)

[package]
name = "graphql-sample"
version = "0.1.0"
authors = ["内緒だよ"]
edition = "2018"

[dependencies]
actix-web = "2"
actix-rt = "1"
juniper = { git = "https://github.com/graphql-rust/juniper" }
dotenv = "*"
env_logger = "*"
log = "*"
serde = { version = "*", features = ["derive"] }
serde_json = "*"
actix-cors = "0.2.0"
r2d2 = "*"
r2d2_postgres = "*"
diesel = { version = "1.0.0", features = ["postgres", "r2d2", "chrono"] }
chrono = "0.4.10"
# 追加
dataloader = "0.13"
futures = "0.3"
async-trait = "0.1"

忘れずにビルドしておきます

$ cargo build

今回追加したライブラリは、Facebook社が参照実装を公開している、dataloaderというライブラリのRust実装版になります。

github.com

ちなみに利用しているWebフレームワークである actix-web ですが、v3が出ています。
今回は未検証なためv2で作ります。

dataloaderを使ってDBからデータを読み込む

早速dataloaderを使って見ましょう。まずは、src/dbフォルダ以下にloaders.rsを作ります。

$ touch src/db/loaders.rs

中身の実装は以下のような感じにします。

(loaders.rs)

use self::diesel::prelude::*;
use crate::db::manager::DataPgPool;
use async_trait::async_trait;
use dataloader::cached::Loader;
use dataloader::BatchFn;
use log::error;
use std::collections::HashMap;
use crate::db::post::Post;
use crate::db::post_repository::PostRepository;

extern crate diesel;

pub struct PostsLoadFn {
    /// 非同期関数内でコネクションプールを直接使おうとすると怒られるのでDataをそのまま持ち回す
    pub pool: DataPgPool,
}

impl PostsLoadFn {
    pub fn posts(&self, keys: &[i32]) -> Vec<Post> {
        let result = PostRepository::any_posts(&self.pool, keys);
        match result {
            Ok(t) => t,
            Err(e) => {
                error!("{}", e);
                Vec::new()
            }
        }
    }
}

#[async_trait]
impl BatchFn<i32, Post> for PostsLoadFn {
    async fn load(&mut self, keys: &[i32]) -> HashMap<i32, Post> {
        let res = self.posts(keys);
        res.iter().map(|p| (p.id, p.clone())).collect()
    }
}

pub type PostsLoader = Loader<i32, Post, PostsLoadFn>;

pub fn create_posts_loader(pool: &DataPgPool) -> PostsLoader {
    Loader::new(PostsLoadFn { pool: pool.clone() }).with_yield_count(100)
}

pub struct Loaders {
    pub posts_loader: PostsLoader,
}

impl Loaders {
    pub fn new(pool: &DataPgPool) -> Loaders {
        Loaders {
            posts_loader: create_posts_loader(pool),
        }
    }
}

リポジトリ側に any_posts メソッドがないので追加します。

(post_repository.rs)

// 略
impl PostRepository {
    // 略
    // 追加: 引数が他のメソッドと違うので注意
    pub fn any_posts(
        pool: &DataPgPool,
        keys: &[i32],
    ) -> Result<Vec<Post>, Error> {
        use crate::schema::posts::dsl::*;
        let conn = &pool.get().unwrap();
        let select_query = posts.filter(id.eq_any(keys));
        let sql = debug_query::<Pg, _>(&select_query).to_string();
        debug!("{:?}", sql);
        select_query.get_results::<Post>(conn)
    }
}

このメソッドは主キーであるidの配列を引数にとり、in句を使ってpostを一括で取得するメソッドになります。
次に、このローダーをリクエストの間で引き回すためにcontextに追加します。graphql/schema.rs を変更します。

(schema.rs)

// 略
pub struct Context {
    pub pool: DataPgPool,
    // 追加
    pub loaders: Loaders,
}
// 略

Contextの内容が変わったので作ってる部分も変更します。(main.rs)

(main.rs)

/// actixからGraphQLにアクセスするためのハンドラメソッド
pub async fn graphql(
    st: web::Data<Arc<Schema>>,
    data: web::Json<GraphQLRequest>,
    pool: web::Data<r2d2::Pool<ConnectionManager<PgConnection>>>,
) -> Result<HttpResponse, Error> {
    let user = web::block(move || {
        let mut rt = futures::executor::LocalPool::new();
        let ctx = &Context {
            pool: pool.clone(),
            // 追加
            loaders: Loaders::new(&pool),
        };
        let graphql_res = data.execute(&st, ctx);
        let res = rt.run_until(graphql_res);
        Ok::<_, serde_json::error::Error>(serde_json::to_string(&res)?)
    })
    .await?;
    Ok(HttpResponse::Ok()
        .content_type("application/json")
        .body(user))
}

では、このメソッドをGraphQL側のハンドラーから呼び出してみます。

(schema.rs)

// 略

#[juniper::graphql_object(Context = Context)]
#[graphql(description = "A Photo returns struct")]
impl Photo {
    pub fn id(&self) -> ID {
        ID::new(self.id.to_string())
    }

    pub fn name(&self) -> String {
        self.name.clone()
    }

    pub fn description(&self) -> String {
        self.description.clone()
    }

    pub fn url(&self) -> String {
        format!("http://hogehoge/{}", self.id)
    }

    pub async fn post(&self, context: &Context) -> FieldResult<Post> {
        Ok(context.loaders.posts_loader.load(self.post_id).await.into())
    }
}

// 略

では、ログを確認してみましょう。

2020-12-04T06:38:52Z INFO  actix_server::builder] Starting 8 workers
[2020-12-04T06:38:52Z INFO  actix_server::builder] Starting "actix-web-service-127.0.0.1:3000" service on 127.0.0.1:3000
[2020-12-04T06:46:27Z DEBUG graphql_sample::db::post_repository] "SELECT \"posts\".\"id\", \"posts\".\"title\", \"posts\".\"content\", \"posts\".\"created_at\", \"posts\".\"updated_at\" FROM \"posts\" WHERE \"posts\".\"id\" IN ($1, $2, $3, $4, $5) -- binds: [4, 5, 3, 1, 2]"
[2020-12-04T06:46:28Z INFO  actix_web::middleware::logger] 127.0.0.1:60600 "POST /graphql HTTP/1.1" 200 3386 "http://localhost:3000/" "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/86.0.4240.198 Safari/537.36" 0.057265

逆向きの依存に対応する

逆にPost→Photoに対応します。
Postのloaderを作ったのと同様に、Postに紐づくPhotoを取得するloaderを作ります。

(loader.rs)

// Association用のPostレコードベクタを作る関数
// 要は関連付けあるものを取得するときにdieselの機能を使いたいための処理
fn create_assoc_posts(keys: Vec<i32>) -> Vec<AssocPost> {
    keys.into_iter().map(|k| AssocPost { id: k }).collect()
}

pub struct PostPhotosLoadFn {
    /// 非同期関数内でコネクションプールを直接使おうとすると怒られるのでDataをそのまま持ち回す
    pub pool: DataPgPool,
}

impl PostPhotosLoadFn {
    pub fn post_photos(&self, keys: &[i32]) -> Vec<Photo> {
        let query_result = PhotoRepository::any_post_photos(&self.pool, keys);
        match query_result {
            Ok(t) => t,
            Err(e) => {
                error!("{}", e);
                Vec::new()
            }
        }
    }
}

#[async_trait]
impl BatchFn<i32, Vec<Task>> for PostPhotosLoadFn {
    async fn load(&mut self, keys: &[i32]) -> HashMap<i32, Vec<Photo>> {
        // associationを取るためには構造体のPostが必要なのでidからダミーを作成
        let assoc_posts: Vec<AssocPost> = create_assoc_posts(keys.to_vec());
        let post_photos = self.post_photos(keys).grouped_by(&assoc_posts);
        let result = assoc_posts
            .iter()
            .zip(post_photos)
            .map(|assoc| (assoc.0.id, assoc.1.clone()))
            .collect();
        result
    }
}

pub type PostPhotosLoader = Loader<i32, Photo, PostPhotosLoadFn>;

pub fn create_post_photos_loader(pool: &DataPgPool) -> PostPhotosLoader {
    Loader::new(PostPhotosLoadFn { pool: pool.clone() }).with_yield_count(100)
}

pub struct Loaders {
    pub posts_loader: PostsLoader,
    // 追加
    pub post_photos_loader: PostPhotosLoader
}

impl Loaders {
    pub fn new(pool: &DataPgPool) -> Loaders {
        Loaders {
            posts_loader: create_posts_loader(pool),
            // 追加
            post_photos_loader: create_post_photos_loader(pool)
        }
    }
}

AssocPost構造体が突然出現してるので、これをdb/post.rsに定義します。

(post.rs)

#[derive(Eq, PartialEq, Debug, Queryable, Clone, Identifiable)]
#[table_name = "posts"]
pub struct AssocPost {
    pub id: i32,
}

さらに、Photo構造体の方も変更します。

(photo.rs)

use crate::db::post::{AssocPost, Post};
use crate::schema::photos;
use chrono::NaiveDateTime;

// DBからのデータ取得用構造体
#[derive(Eq, PartialEq, Debug, Queryable, Associations, Identifiable, Clone)]
#[belongs_to(parent = "Post")]
#[belongs_to(parent = "AssocPost", foreign_key = "post_id")]
pub struct Photo {
    pub id: i32,
    pub post_id: i32,
    pub name: String,
    pub description: Option<String>,
    pub created_at: NaiveDateTime,
    pub updated_at: NaiveDateTime,
}

これは、diesel(ORM)の関連付け機能を使いたいがための処理です。IDだけ持つ構造体をダミーで生成することで対応します。

最後に、リゾルバでloaderを使うように変更します。

(schema.rs)

//　略

/// 実際にGraphQLとしての型になるのは以下アトリビュートがついているこちら
#[juniper::graphql_object(Context = Context)]
#[graphql(description = "A Post returns struct")]
impl Post {
    pub fn id(&self) -> ID {
        ID::new(self.id.to_string())
    }

    pub fn title(&self) -> String {
        self.title.clone()
    }

    pub fn content(&self) -> String {
        self.content.clone()
    }

    async fn photos(&self, context: &Context) -> FieldResult<Vec<Photo>> {
        Ok(context
            .loaders
            .post_photos_loader
            .load(self.id)
            .await
            .into_iter()
            .map(|t| t.into())
            .collect())
    }
}

// 略

ここまで書いたら http://localhost:3000/graphiql にアクセスして以下のクエリを実行してみます。

query {
  allPosts {
    id
    title
    content
    photos {
      id
      name
      description
      url
    }
  }
}

ログを確認してみると以下のような出力になっています。

2020-12-04T07:41:57Z INFO  actix_server::builder] Starting 8 workers
[2020-12-04T07:41:57Z INFO  actix_server::builder] Starting "actix-web-service-127.0.0.1:3000" service on 127.0.0.1:3000
[2020-12-04T07:42:40Z DEBUG graphql_sample::db::photo_repository] "SELECT \"photos\".\"id\", \"photos\".\"post_id\", \"photos\".\"name\", \"photos\".\"description\", \"photos\".\"created_at\", \"photos\".\"updated_at\" FROM \"photos\" WHERE \"photos\".\"post_id\" IN ($1, $2, $3, $4, $5) -- binds: [3, 1, 2, 4, 5]"
[2020-12-04T07:42:40Z INFO  actix_web::middleware::logger] 127.0.0.1:61719 "POST /graphql HTTP/1.1" 200 2313 "http://localhost:3000/" "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/86.0.4240.198 Safari/537.36" 0.039096

in句のSQLが発行されてることが確認できました。
さらに、相互なデータの取り合いに対応します。post→photo→postといった形です。schema.rsを修正します。

/// GraphQLのクエリ系リゾルバ
#[juniper::graphql_object(Context = Context)]
impl Query {
    fn all_photos(&self, context: &Context) -> FieldResult<Vec<Photo>> {
        PhotoRepository::all_photos(context)
            .and_then(|photos| Ok(photos.into_iter().map(|t| t.into()).collect()))
            .map_err(Into::into)
    }

    async fn all_posts(&self, context: &Context) -> FieldResult<Vec<Post>> {
        let posts = PostRepository::all_posts(context)?;
        let mut result = Vec::new();
        for post in posts {
            context.loaders.posts_loader.prime(post.id,post.clone()).await;
            result.push(post.into());
        }
        Ok(result)
    }
}

all_postsを変更しました。リポジトリ経由で取得したデータをloaderに登録するようにしています。all_photosメソッドも同様にすれば良いので、今回は割愛します。
では、クエリを実行してみましょう。

query {
  allPosts {
    id
    title
    content
    photos {
      id
      name
      description
      url
      post {
        id
        title
        content
      }
    }
  }
}

ログを確認します。

2020-12-04T08:00:28Z INFO  actix_server::builder] Starting 8 workers
[2020-12-04T08:00:28Z INFO  actix_server::builder] Starting "actix-web-service-127.0.0.1:3000" service on 127.0.0.1:3000
[2020-12-04T08:00:34Z DEBUG graphql_sample::db::photo_repository] "SELECT \"photos\".\"id\", \"photos\".\"post_id\", \"photos\".\"name\", \"photos\".\"description\", \"photos\".\"created_at\", \"photos\".\"updated_at\" FROM \"photos\" WHERE \"photos\".\"post_id\" IN ($1, $2, $3, $4, $5) -- binds: [5, 4, 3, 1, 2]"
[2020-12-04T08:00:34Z INFO  actix_web::middleware::logger] 127.0.0.1:61944 "POST /graphql HTTP/1.1" 200 3665 "http://localhost:3000/" "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/86.0.4240.198 Safari/537.36" 0.044010

postのログは出てこないことが確認できたかと思います。loaderに一度登録された値はそのコンテキストが有効な間はキャッシュされます。従って、post→photo→post→photoなどを繰り返す場合でもSQLが実行されるのは一回で済みます。

まとめ

今回は、GraphQLでいかにN+1問題を解消するかについて紹介しました。実際にはdataloaderの実装の中身などを見てみると、非同期処理をどう待ち受けるのかなど勉強になると思います。また、機会があれば紹介してみたいと思います。
一方で困ったこととしては、ドキュメントに書いて無いことがいくつかあった点があります。例えば、公式ドキュメントを見に行くと、 juniper::graphql_object アトリビュートはDEPRECATIONとされています。

docs.rs

一方で、以下のプルリクを見ると違ったことが書いてあります。

github.com

Hello! The doc is actually wrong, the name of the proc macro is graphql_object. Would love a PR to fix.

ドキュメントが間違ってる&async awaitを使いたいならmasterブランチでないとだめ、といったことが書かれていました。

などなど、枯れてない技術特有の部分で結構困る部分はあるかなと思いました。

P.S.

カジュアル面談

弊社エンジニアと1on1で話せます、カジュアル面談も現在受付中です！こちらも是非ご検討ください。 yumenosora.connpass.com

その他採用情報

虎の穴ラボでの開発に少しでも興味を持っていただけた方は、採用説明会やカジュアル面談という場でもっと深くお話しすることもできます。ぜひお気軽に申し込みいただければ幸いです。
カジュアル面談では虎の穴ラボのエンジニアが、開発プロセスの内容であったり、「今期何見ました？」といったオタクトークから業務の話まで何でもお応えします。
カジュアル面談や採用情報はこちらをご確認ください。
yumenosora.co.jp