トップ > Go > sqlcで使うDBドライバをpqからpgx/v5に移行した話

sqlcで使うDBドライバをpqからpgx/v5に移行した話

Go

ogp

この記事は enechain Advent Calendar 2023 の 16 日目の記事です。

はじめに

みなさん、こんにちは。enechainでSWEをしている@kkagurazakaです。

enechainでは最低限のルール (バックエンドサービス間の通信はgRPCを使う、など) を守っていれば、技術選定は各チームに委ねられています。 たとえば、本Advent Calendarの10日目の記事「GoのAtlasとBunを使ったマイグレーション環境を構築する」を書いたJCEXチームでは、ORMにBunを採用しています。

一方、我々のチームではsqlc-dev/sqlcをデータベースへのクエリに使っています。 最近、そのsqlc経由で用いるPostgreSQLのドライバをpqからpgx/v5に移行したのですが、それに際して幾つかの工夫をする必要がありました。 そこで、今回の記事では我々がどのようにして移行を行ったのかについて紹介します。

背景

GoでSQLデータベースにクエリを投げたいときは、標準ライブラリであるdatabase/sqlパッケージを使うのが一般的です。 このパッケージでは、ドライバを切り替えることによって様々なSQLデータベースを統一されたインターフェースで利用でき、たとえばPostgreSQLの場合はpqが有名なドライバです。

我々が採用しているsqlcにおいても、生成コードはデフォルトでdatabase/sqlパッケージ経由でpqドライバを用いるようになっていますが、pqには次の課題があります。

  1. メンテナンスモードに入っているため、新しい機能は追加されず、メンテナは主体的にissueを解決しない
  2. NULLの扱いが sql.NullString といったwrap構造体になっており、扱いづらい

特に1つめの課題が大事で、既に枯れて安定しているソフトウェアとはいえPostgreSQL自体は進化を続けていることを踏まえると、今後のメンテがされないことはリスクとなります。 pqのGitHubページでもそのリスクを気にするユーザはpgxの利用を推奨すると記載されているため、移行を決断しました。

実際の移行プロセス

アダプタの選定

pgxに移行するにあたり、まず決めなければいけないのはGoで使う際のアダプタをどうするかです。 pgxでもpgx/stdlibパッケージを利用すると、これまでと同じくdatabase/sqlパッケージのアダプタを利用できます。一方で、ドライバを切り替えて複数のSQLデータベースに対応するというコンセプト上、たとえばCOPYのようなPostgreSQL独自機能にはアクセスできません。

移行時点でCOPYなどを使う予定はありませんでしたが、せっかく移行するのですから今後の可能性は広くとっておくほうがお得です。 また、sqlcでpgxを採用する場合、標準的な扱い方ではpgxのアダプタを利用するコードが生成されることもあり、database/sqlパッケージとは決別することにしました。

sqlcの設定

sqlc.yaml を次のように設定しました。なお、ドライバの移行に直接的に関係ないところは省いています。

version: "2"
sql:
  - engine: postgresql
    gen:
      go:
        package: db
        out: gen/db
        sql_package: pgx/v5
        emit_pointers_for_null_types: true
        overrides:
          - db_type: pg_catalog.numeric
            go_type: github.com/shopspring/decimal.Decimal
          - db_type: pg_catalog.numeric
            go_type:
              import: github.com/shopspring/decimal
              type: Decimal
              pointer: true
            nullable: true
          - db_type: uuid
            go_type:
              import: github.com/gofrs/uuid/v5
              package: uuid
              type: UUID
          - db_type: uuid
            go_type:
              import: github.com/gofrs/uuid/v5
              package: uuid
              type: UUID
              pointer: true
            nullable: true
          - db_type: timestamptz
            go_type:
              import: time
              type: Time
          - db_type: timestamptz
            go_type:
              import: time
              type: Time
              pointer: true
            nullable: true

sql_packageにpgx/v5を設定することで、生成されるコードがpgxのドライバ・アダプタを利用するようになります。

また、pgxを使う場合に限りemit_pointers_for_null_typesをtrueに設定できます。 これは読んで字の如くNULLになりうるカラムの型をポインタ型として扱うためのオプションです。 pgxではNULLの取り扱いとしてpgtype.Textのようなwrap構造体と*stringのようなポインタ型の両方をサポートしているため、どちらを使ったコードを生成するかをsqlcに指示するオプションになります。

そして次にズラっと並んでいるのがtype overrideの設定です。 pgxはDECIMAL・UUID・TIMESTAMPTZといったPostgreSQLがサポートする型についてはpgx/pgtypeパッケージで対応する型を網羅しています。

しかしながら、我々はレイヤードアーキテクチャでいうところのインフラ層以外の部分において、decimal型にはshopspring/decimalを、uuid型にはgofrs/uuidを使用していたため、変換の手間をなくすためには直接これらの型にマッピングして欲しいところです。 これについてもpgxは独自の型変換処理を登録する機構があり、また前述の2つのパッケージは公式でその変換を行うパッケージが公開されているため、次のようにシンプルに対応できます。

cfg, err := pgxpool.ParseConfig(connString) // connStringはDBへの接続文字列
if err != nil {
    // エラーハンドリング
}

cfg.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
    pgx_decimal.Register(conn.TypeMap())
    pgx_uuid.Register(conn.TypeMap())
    return nil
}

TIMESTAMPTZ型についても、time.Timeへ変換して扱うように設定しています。 ドキュメントによると、time.Timeのポインタ型は使うべきではなく、そのシーンではゼロ値を使うべきだという記載があります。 ただし、我々としては他の型との対称性を保ちたかったことと、sqlcが生成したテーブルに対応する構造体の型を見るだけでそのカラムがNULL許容なのかがわかるというメリットを重視して、NULL許容のTIMESTAMPTZは*time.Timeへマッピングさせています。

これまでsql.NullStringといったdatabase/sqlの型を使っている箇所はすべて置き換える必要がありましたが、元々ユーティリティ関数を用意して変換していたため、使用箇所の特定は容易でした。

アダプタ変更に伴う対応

使用するsql_packageをpgx/v5に変更すると、sqlcが生成するエントリーポイントであるQueriesは次のように変わります(一部を抜粋)。

package db

type DBTX interface {
    Exec(context.Context, string, ...interface{}) (pgconn.CommandTag, error)
    Query(context.Context, string, ...interface{}) (pgx.Rows, error)
    QueryRow(context.Context, string, ...interface{}) pgx.Row
}

func New(db DBTX) *Queries {
    return &Queries{db: db}
}

type Queries struct {
    db DBTX
}

func (q *Queries) WithTx(tx pgx.Tx) *Queries {
    return &Queries{
        db: tx,
    }
}

DBTX interfaceは*pgx.Connまたは*pgx.Poolが満たします。実際のプロダクションでは基本的にコネクションプールの方を利用することになるでしょう。

pool, err := pgxpool.NewWithConfig(ctx, cfg) // cfgは前述のもの
if err != nil {
    // エラーハンドリング
}

pgx.BeginFunc(
  ctx, pool, func(tx pgx.Tx) error {
      q := db.New(pool).WithTx(tx)
      // qを使ってクエリを投げる
  },
)

この書き換えでプロダクションコードについては動作するようになりました。しかしながら、テストにおいて2つの問題が発生しました。

テストデータ (fixture) の投入

我々のチームのテスト方針として、DBはモックせずにDockerで立ち上げたテスト用のコンテナに実際にアクセスさせています。 その際、テストデータの投入にはgo-testfixtures/testfixturesを利用していました。 これはyamlでフィクスチャを定義できるようになるパッケージで、他の言語でもよくある機能かと思います。

問題となったのは、このtestfixturesライブラリはdatabase/sqlパッケージのアダプタである*sql.DBを要求している点でした。 代わりに必要なメソッドのみを定義したinterfaceを受け入れるように変更してPRを出す選択肢も考えましたが、変更量がかなり多くなってしまうので受け入れてもらえなかったり、マージまでの時間がかかってしまう恐れがありました。

また、チームでもフィクスチャの再利用性の観点からコードで表現できたらより扱いやすいという議論もあったため、メルカリさんのGoでテストのフィクスチャをいい感じに書くを参考に、自前でフィクスチャを投入する実装を行いました。 基本的には記事の通りですが、モデルにはsqlcが生成したテーブルに対応する構造体をそのまま使っています。

Golandの補完とGitHub Copilotが優秀なこともあり、現在は愚直にマニュアル実装していますが、いずれsqlcのプラグインで生成できるようすることも考えています。

DBにアクセスするテストの並列実行

Goではテストの実行はパッケージごとに並列実行されます。また、t.Parallel()を呼び出してテストケースごとに並列実行させるのも一般的です。 前述のとおり、テストに使うDBはDockerで立ち上げたコンテナ1つのため、何も考えずにDBを利用するテストを書くと、各テストケースでデータが競合する問題が発生します。

これまではDATA-DOG/go-txdbを利用することでこの問題を解決していました。 go-txdbはDBのコネクションが開くと同時にトランザクションを貼り、閉じるときにロールバックするパッケージです。 各テストケースごとにgo-txdb経由でコネクションを開くことで、データが競合してしまう問題に対処できます。

このgo-txdbもdatabase/sqlパッケージのドライバとして実装されており、実際に使うドライバをwrapする形で実装されています。 移行にあたってdatabase/sqlからはByeByeしているので、こちらも別のソリューションを考えなければなりません。

go-txdbのアイデアは非常に気に入っていたため、pgxのアダプタでも同じような実装を自前で行うことにしました。

func RunWithFixture(t *testing.T, fix *fixture.Fixture, invoke func(ctx context.Context, q *db.Queries)) {
    t.Helper()

    ctx := context.Background()

    // 接続設定
    cfg, err := pgx.ParseConfig("timezone=Asia/Tokyo")
    if err != nil {
        assert.FailNow(t, err.Error())
    }

    // DBに接続
    conn, err := pgx.ConnectConfig(ctx, cfg)
    if err != nil {
        assert.FailNow(t, err.Error())
    }
    defer func() { _ = conn.Close(ctx) }()

    pgx_decimal.Register(conn.TypeMap())
    pgx_uuid.Register(conn.TypeMap())

    // トランザクションを貼る
    tx, err := conn.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.Serializable})
    if err != nil {
        assert.FailNow(t, err.Error())
    }
    defer func() { _ = tx.Rollback(ctx) }()

    // フィクスチャ投入
    fix.Setup(t, tx)

    // 実際のテスト
    invoke(ctx, db.New(tx))
}

testifyが使われているのは気にしないでください。 やっていることは単純で、トランザクションを貼ってフィクスチャを投入し、テストを実行したあとにdeferでトランザクションがロールバックされるようになっています。 各テストケース間でコネクションが使い回されることによるトラブルを避けたかったため、コネクションプールは利用していません。 invokeの中では*db.QueriesをDIの機能で差し替えてテストを実行させています。

1点注意が必要な点は、この方法は並列実行セーフではあるものの、トランザクションの分離を活用しているため、実行自体は並列ではないということです。 現時点ではそこまでプロダクトコードが大きくないので実行時間は問題にはなっていませんが、テストケースが肥大化して時間がかかるようになった場合は、例えばパッケージごとに別のDBコンテナを立てるなどの最適化が必要になるかもしれません。

おわりに

本記事ではsqlcでPostgreSQLを扱う際のドライバをpqからpgx/v5に変更するにあたり行った工夫について紹介しました。 ドライバの移行そのものは特につまることなくスムーズに行えましたが、database/sqlパッケージからの脱却には一定の工夫が求められました。 それでもメンテされないパッケージを使い続けるリスクの解消と、ポインタ型によるNULL表現での開発体験は素晴らしく、移行してよかったと思っています。

明日は同じチームのPdMであるokpさんが「経験が裏目に？！大手IT企業出身PdMが3カ月で学んだこと」を投稿予定です。

enechainでは、日本の電気インフラを支えるプロダクトを共に作っていける仲間を募集しています。

herp.careers herp.careers