はじめに

本記事の確認環境は以下
- go: 1.17.2
- sqlx: v1.3.5
サンプルコードについて
- サンプルコードに対するエラーハンドリングは省いています。実際に取り扱う際にはエラーハンドリングしてください。
- 実際に取り扱う際にエラーを扱わない場合は、MustExecなどのMustがついた関数を利用ください。

基本的な使い方

context.WithTimeout()を使います。以下のようにtimer付きのcontextを取得。扱っているcontextがあるなら、WithTimeoutの第一引数に指定します。
また、cancel()は処理の終了時に実行しておく必要があります。

    ctx, cancel := context.WithTimeout(context.Background(), 1 * time.Minute)
    defer cancel()

sqlxでは、通常使うSQL実行用の関数に対して、その関数にcontextの引数を追加した関数が用意されています。例えば NamedExecなら、[NamedExecContext]。
こちらに変更し、上記でtimeoutを設定するわけです。

例えばこんなコードなら

before

    db, _ := sql.Open("postgres", "your connection option")
    db.NamedExec("INSERT INTO person (first_name, last_name, email) VALUES ($1, $2, $3)", "Jason", "Moiron", "jmoiron@jmoiron.net")

こうなります。

after

    db, _ := sql.Open("postgres", "your connection option")

    ctx, cancel := context.WithTimeout(context.Background(), 1 * time.Minute)
    defer cancel()
    db.NamedExecContext(ctx, "INSERT INTO person (first_name, last_name, email) VALUES ($1, $2, $3)", "Jason", "Moiron", "jmoiron@jmoiron.net")

Transactionを使う場合

基本的なTransactionの扱い方

Beginxを使います。Openの後にBeginxを実行、そのreturn値に対してSQL実行用関数を実行するだけです。

basic transaction

    db, _ := sql.Open("postgres", "your connection option")
    tx, _ := db.Beginx()
    //ここがtx.に変わる
    tx.NamedExec("INSERT INTO person (first_name, last_name, email) VALUES ($1, $2, $3)", "Jason", "Moiron", "jmoiron@jmoiron.net")

    //最後にCommit or Rollback。異常終了時もpackage内でRollback処理が走る。
    tx.Commit()

他のサンプルが見たい場合は公式のサンプルをどうぞ。

Timeoutを使いたい場合

context.WithTimeout()を使い、Beginxの代わりにBeginTxxでcontextを渡します。
ただし、BeginTxxに記載のあるとおり、cancelをCommit前に実行してしまうと、処理はrollbackされCommitが失敗しますので注意してください。

The provided context is used until the transaction is committed or rolled back. If the context is canceled, the sql package will roll back the transaction. Tx.Commit will return an error if the context provided to BeginxContext is canceled.

transaction with timeout

    db := sql.Open("postgres", "your connection option")

    ctx, cancel := context.WithTimeout(context.Background(), 1 * time.Minute)
    defer cancel() //commit後に実行するよう注意
    tx, _ := db.Beginxx(ctx, &sql.TxOptions{Isolation: sql.LevelDefault, ReadOnly:false})

    tx.NamedExec("INSERT INTO person (first_name, last_name, email) VALUES ($1, $2, $3)", "Jason", "Moiron", "jmoiron@jmoiron.net")
    tx.Commit()

コードの内容をみた調査結果も置いときます。興味ある方は開いてみてください

Commit時処理が以下。
<-tx.ctx.Done() にcontextの終了時イベントが飛んでくるようになっているので、各所でこれをチラ見してtimeoutに対応しています。

この中のtx.cancel()がContextのcancelを保存する場所になっているので、cancelはCommit時に実行してくれそうです。
こちらを信用してもいいんだけど、公式ドキュメントで内部でcancel実行するでーって書いてるわけでもないので、commit後にcancelという流れが良さそう

database/sql/sql.go

// Commit commits the transaction.
func (tx *Tx) Commit() error {
        // Check context first to avoid transaction leak.
        // If put it behind tx.done CompareAndSwap statement, we can't ensure
        // the consistency between tx.done and the real COMMIT operation.
        select {
        default:
        case <-tx.ctx.Done():
                if atomic.LoadInt32(&tx.done) == 1 {
                        return ErrTxDone
                }
                return tx.ctx.Err()
        }
        if !atomic.CompareAndSwapInt32(&tx.done, 0, 1) {
                return ErrTxDone
        }

        // Cancel the Tx to release any active R-closemu locks.
        // This is safe to do because tx.done has already transitioned
        // from 0 to 1. Hold the W-closemu lock prior to rollback
        // to ensure no other connection has an active query.
        tx.cancel()
        tx.closemu.Lock()
        tx.closemu.Unlock()

        var err error
        withLock(tx.dc, func() {
                err = tx.txi.Commit()
        })
        if err != driver.ErrBadConn {
                tx.closePrepared()
        }
        tx.close(err)
        return err
}

ちなみにこちらがbegin時の終了待ち受け。selectじゃないので無限待ちになっています。 (struct{}へのメッセージにすると、受け取り先なくただの待ち受けchannelとして使えるんですね。地味に便利だな)

database/sql/sql.go

// awaitDone blocks until the context in Tx is canceled and rolls back
// the transaction if it's not already done.
func (tx *Tx) awaitDone() {
        // Wait for either the transaction to be committed or rolled
        // back, or for the associated context to be closed.
        <-tx.ctx.Done()

        // Discard and close the connection used to ensure the
        // transaction is closed and the resources are released.  This
        // rollback does nothing if the transaction has already been
        // committed or rolled back.
        // Do not discard the connection if the connection knows
        // how to reset the session.
        discardConnection := !tx.keepConnOnRollback
        tx.rollback(discardConnection)
}

ちなみに私はgo1.12 -> go1.18へのupdateを行なっている際に、cancelの挙動でcommitできなくなってハマってました。
今までの実装だと ctx, cancel = context.WithTimeout(ctx, timeout) を繰り返してSQL1本1本に対してtimeoutを設定していました。

今回合わせてcontextの仕様を確認したところ、context.WithValue()やWithTimeoutなどで生成した引数とreturn値のcontextが親子関係になるんですね。公式の定義でも引数がparentとなっています。

本来親のcontextがcancelされると子もcancelされるようになる挙動となっているのですが、偶然cancelされない状況にはまってくれていたようです。
おそらくこちらのcommitで修正されたようです。
https://github.com/golang/go/commit/0ad368675bae1e3228c9146e092cd00cfb29ac27

コード差分や現象を見ても、上記で正しそう。ハマった。。。

参考

cancelの記載も含め、色々と参考情報があります。
https://www.alexedwards.net/blog/how-to-manage-database-timeouts-and-cancellations-in-go