gocraft/dbrでPostgreSQLのinsert IDを取得する

Go言語でのORMライブラリもいろいろありますが、今はgocraft/dbrを使ってみています。

PostgreSQLでinsertをした時にプライマリキーがserialの場合はIDを取得したいケースがよくあると思います。

このライブラリに限ったことではないのかもしれませんが、どうもデフォルトの機能ではうまく取得ができないので、やり方を調べてみました。

1 標準の機能
2 RETURNING句
3 成果物

標準の機能

dbrでのinsert実行は以下のコードで実行されます。

func (b *InsertBuilder) Exec() (sql.Result, error) {
	result, err := exec(b.runner, b.EventReceiver, b, b.Dialect)
	if err != nil {
		return nil, err
	}

	if b.RecordID.IsValid() {
		if id, err := result.LastInsertId(); err == nil {
			b.RecordID.SetInt(id)
		}
	}

	return result, nil
}

これを見ると、sql.ResultのLastInsertId()を実行してRecordIDとやらにセットしようとしているのが分かります。

なんだ、対応してるじゃん！...と思っていた時が私にもありました。。

実は、そもそもdbrのPostgreSQL用ドライバとして使っているlib/pqがLastInsertId()に対応していないらしいのです。。lib/pgのGoDocにも、

pq does not support the LastInsertId() method of the Result type in database/sql. To return the identifier of an INSERT (or UPDATE or DELETE), use the Postgres RETURNING clause with a standard Query or QueryRow call

と書いてあります。

さらに、取得するための方法も書いてあって、"RETURNING"句を使えとあります。

RETURNING句

では、RETURNING句とはなんでしょう？PostgreSQLのINSERT文のドキュメントを見てみます。

RETURNING句を指定すると、INSERTは実際に挿入された（あるいはON CONFLICT DO UPDATE句によって更新された）各行に基づいて計算された値を返すようになります。これは、通番のシーケンス番号など、デフォルトで与えられた値を取り出す時に主に便利です。しかし、そのテーブルの列を使用した任意の式を指定することができます。 RETURNINGリストの構文はSELECTの出力リストと同一です。挿入または更新に成功した行だけが返されます。例えば、行がロックされていて、ON CONFLICT DO UPDATE ... WHERE句の conditionが満たされなかったために更新されなかった行は返されません。

これは、通番のシーケンス番号など、デフォルトで与えられた値を取り出す時に主に便利です。 とありますので、まさに今回の目的にぴったりのようです。

さらに調査を進めると、dbrのGitHubのissueに同様の内容が上がっていました。
その中で、RETURNINGを使ってIDを取得するコードを上げてくれている人がいました。

It looks kinda hacky and it does not support multiple records. This logic should be somewhere inside the dialect handler.
I have no better option right now though, I currently use this function as a workaround:

func execInsertGetIDs(r dbr.SessionRunner, b *dbr.InsertBuilder, ids []*int) error {
    var rids []int
    q, vals := b.ToSql()
    _, err := r.SelectBySql(q+" RETURNING \"id\"", vals...).Load(&rids)
    if err != nil {
        return err
    }

    for i, id := range rids {
        *ids[i] = id
    }

    return nil
}

これを参考に実装したいと思いますが、上記のままだと通常のdbrによるinsertの処理の書き方とだいぶ変わってしまうので、その辺りをあまり変えずに実装したいと思います。

例えば、標準のinsertの処理は以下のように書けます。

_, err := sess.InsertInto("muser").
	Columns("nickname", "mail_address").
	Record(u).
	Exec()

この書き方をほとんど変えずに作るのが目標です。

成果物

では、早速今回の成果物を見てみましょう。

package db

import (
    "myapp/common"

    "github.com/Sirupsen/logrus"
    "github.com/gocraft/dbr"
    _ "github.com/lib/pq"
)

func Init() *PgSession {
    session := getSession()
    return session
}

func getSession() *PgSession {
    db, err := dbr.Open("postgres",
        "postgres://"+common.DbUser+":"+common.DbPassword+"@"+common.DbHost+":"+common.DbPort+"/"+common.DbName+"?sslmode=disable",
        nil)

    if err != nil {
        logrus.Error(err)
        return nil
    }

    return &PgSession{Session: db.NewSession(nil)}
}

type kvs map[string]string

type PgSessionRunner interface {
    Select(column ...string) *dbr.SelectBuilder
    SelectBySql(query string, value ...interface{}) *dbr.SelectBuilder

    InsertInto(table string) *PgInsertBuilder
    InsertBySql(query string, value ...interface{}) *dbr.InsertBuilder

    Update(table string) *dbr.UpdateBuilder
    UpdateBySql(query string, value ...interface{}) *dbr.UpdateBuilder

    DeleteFrom(table string) *dbr.DeleteBuilder
    DeleteBySql(query string, value ...interface{}) *dbr.DeleteBuilder
}

type PgSession struct {
    *dbr.Session
}

type PgTx struct {
    *dbr.Tx
}

var (
    _ PgSessionRunner = (*PgTx)(nil)
    _ PgSessionRunner = (*PgSession)(nil)
)

func (sess *PgSession) Begin() (*PgTx, error) {
    tx, err := sess.Session.Begin()
    return &PgTx{Tx: tx}, err
}

func (sess *PgSession) InsertInto(table string) *PgInsertBuilder {
    return &PgInsertBuilder{InsertBuilder: sess.Session.InsertInto(table), SessionRunner: sess.Session}
}

func (tx *PgTx) InsertInto(table string) *PgInsertBuilder {
    return &PgInsertBuilder{InsertBuilder: tx.Tx.InsertInto(table), SessionRunner: tx.Tx}
}

type PgInsertBuilder struct {
    *dbr.InsertBuilder
    dbr.SessionRunner
}

func (b *PgInsertBuilder) Pair(column string, value interface{}) *PgInsertBuilder {
    b.InsertBuilder.Pair(column, value)
    return b
}

func (b *PgInsertBuilder) Columns(column ...string) *PgInsertBuilder {
    b.InsertBuilder.Columns(column...)
    return b
}

func (b *PgInsertBuilder) Record(structValue interface{}) *PgInsertBuilder {
    b.InsertBuilder.Record(structValue)
    return b
}

func (b *PgInsertBuilder) Values(value ...interface{}) *PgInsertBuilder {
    b.InsertBuilder.Values(value)
    return b
}

func (b *PgInsertBuilder) Exec() (int64, error) {
    var rids []int64
    q, vals := b.ToSql()

    _, err := b.SessionRunner.SelectBySql(q+" RETURNING \"id\"", vals...).Load(&rids)
    if err != nil {
        return 0, err
    }

    for _, id := range rids {
        return id, nil
    }

    return 0, nil
}

dbrのSession、Tx、InsertBuilderをそれぞれ埋め込んだstructを新しく定義してみました。
InsertBuilderのExec以外のメソッドは特に処理を変更する必要ないので、それぞれ移譲しています。

RETURNING句のところですが、"id"に固定しているのでプライマリキーの名前がidじゃない場合には使えません。

実際の使い方ですが、例えばモデルクラスに以下のようにメソッドを定義して使えます。

func (u *User) Insert(sess db.PgSessionRunner) error {
	lastId, err := sess.InsertInto("muser").
		Columns("nickname", "mail_address").
		Record(u).
		Exec()

	if err != nil {
		logrus.Error("error: " + err.Error())
	}

	if lastId > 0 {
		u.Id = lastId
	}

	return err
}

どうでしょう。もし参考にできる方がいらっしゃれば幸いです。