GORM DBResolverで実現するクエリ自動分散

はじめに

こんにちは。技術本部MA部の平井です。
普段はマーケティングシステムのリプレイスプロジェクトに従事しています。
今回は、そのリプレイスプロジェクトにおいてGORM DBResolverを利用し、データベースのクエリ自動分散を実現した内容についてお話ししたいと思います。

リプレイス対象システムについては以下のテックブログに詳細が記載されています。

背景

リプレイス対象システムには、その特徴からデータベースに高可用性が求められます。
そのため、AlloyDBをリプレイス先システムのデータベースとして採用しています。

また、アプリケーションはGoで開発しており、ORMとしてGORMを利用しています。

AlloyDBの構成

AlloyDBでは、クラスターの中に更新および読み取りクエリ用のプライマリーインスタンスを1つ、読み取りクエリ専用の読み取りプールインスタンスを複数立ち上げることが出来ます。

また、プライマリーインスタンスと読み取りプールインスタンスは別のエンドポイントを持っています。

今回のプロジェクトでもこれらの構成でAlloyDBを利用しています。

実現したいこと

先述の通り、ORMにはGORMを利用しているため、GORMを活用して更新系クエリはプライマリーインスタンスへ、読み取りクエリは読み取りプールインスタンスへとクエリを自動分散させたいと考えました。

そこで、GORMのDBResoloverを利用することで、上記の要件を簡単に実現できることがわかりました。

GORMのDBResolverとは

GORMが公式に提供しているパッケージです。以下の機能があります。

• Multiple sources, replicas
• Read/Write Splitting
• Automatic connection switching based on the working table/struct
• Manual connection switching
• Sources/Replicas load balancing
• Works for RAW SQL
• Transaction

今回はRead/Write Splittingを利用して、クエリの自動分散を実現しました。

実装方法

設定方法はドキュメント記載通りです。

	cnfPrimaryDB := cnfDB.Source
	primaryDsn := "host=primaryHost user=primaryUser"

    // ①
	db, err := gorm.Open(postgres.Open(primaryDsn), &gorm.Config{})
 
	if err != nil {
		panic(err.Error())
	}

	cnfReplicaDB := cnfDB.Replica

    replicaDsn := "host=readHost user=readUser"
    // ②
    err = db.Use(dbresolver.Register(dbresolver.Config{
        Replicas:          []gorm.Dialector{postgres.Open(replicaDsn)},
        TraceResolverMode: true,
    })

    if err != nil {
        panic(err.Error())
    }

①でプライマリーインスタンスへ接続しています。②のReplicasに読み取りプールインスタンスの接続情報を渡して設定完了です。

TraceResolverModeをtrueにすることでログにどのデータベースに対してクエリを実行したかが出力されます。

検証

ローカル環境で実際に処理を実行した結果、以下のログが出力されました。
ログから、INSERTクエリはプライマリーインスタンス(source), SELECTクエリは読み取りプールインスタンス(replica)に対して実行されていることが確認できました。

app-1          | [3.778ms] [rows:1] [source] INSERT INTO (以降は省略)
app-1          | [0.881ms] [rows:0] [replica] SELECT * FROM (以降は省略)

ソースコード

DBResolverの内部でどのように分散しているのか気になったので、実際にソースコードを読んでみました。
ドキュメントに記載されている通り、以下のように接続を切り替えるCallBackを設定していることがわかりました。

package dbresolver

import (
	"strings"

	"gorm.io/gorm"
)

func (dr *DBResolver) registerCallbacks(db *gorm.DB) {
	dr.Callback().Create().Before("*").Register("gorm:db_resolver", dr.switchSource)
	dr.Callback().Query().Before("*").Register("gorm:db_resolver", dr.switchReplica)
	dr.Callback().Update().Before("*").Register("gorm:db_resolver", dr.switchSource)
	dr.Callback().Delete().Before("*").Register("gorm:db_resolver", dr.switchSource)
	dr.Callback().Row().Before("*").Register("gorm:db_resolver", dr.switchReplica)
	dr.Callback().Raw().Before("*").Register("gorm:db_resolver", dr.switchGuess)
}

注意点

注意点としては、Rawメソッドを利用する場合、SELECTで始まらない場合はSourceにクエリを実行してしまいます。
実際、Rawメソッドを利用してWITH句から始まるクエリを実行していた箇所があり、そのクエリがプライマリーインスタンスで実行される問題が発生しました。

このような場合は、以下のように対象データベースを明示的に指定することで解決することが出来ます。

db.Clauses(dbresolver.Read).Raw

最後に

GormのDBResolverは比較的シンプルに設定でき、コードの変更も最小限で済みました。内部で自動的にクエリの内容に応じて分散処理を行ってくれるため、とても便利でした。