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

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

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

https://techblog.zozo.com/entry/real-time-marketing-system

# 背景
リプレイス対象システムには、その特徴からデータベースに高可用性が求められます。
そのため、[AlloyDB](https://cloud.google.com/alloydb/docs)をリプレイス先システムのデータベースとして採用しています。

また、アプリケーションはGoで開発しており、ORMとして[GORM](https://gorm.io/)を利用しています。
# AlloyDBの構成
AlloyDBでは、クラスターの中に更新および読み取りクエリ用のプライマリーインスタンスを1つ、読み取りクエリ専用の読み取りプールインスタンスを複数立ち上げることが出来ます。

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

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

![スクリーンショット 2024-12-17 11.53.20.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299065/642d6981-f87e-6489-680d-ced72bcfb73e.png)

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

そこで、GORMの[DBResolover](https://gorm.io/docs/dbresolver.html)を利用することで、上記の要件を簡単に実現できることがわかりました。

# 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`](https://gorm.io/docs/dbresolver.html#Read-x2F-Write-Splitting)を利用して、クエリの自動分散を実現しました。

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

```go
	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の内部でどのように分散しているのか気になったので、実際にソースコードを読んでみました。
[ドキュメント](https://gorm.io/docs/dbresolver.html#Read-x2F-Write-Splitting:~:text=Read/Write%20splitting%20with%20DBResolver%20based%20on%20the%20current%20used%20GORM%20callbacks.)に記載されている通り、以下のように接続を切り替えるCallBackを設定していることがわかりました。

```go
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)
}
```

https://github.com/go-gorm/dbresolver/blob/ece044e659ee1ffc1810267bf25d0749a478258b/callbacks.go#L9-L16

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

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

```go
db.Clauses(dbresolver.Read).Raw
```

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