在Go中包装一个db对象，并在同一事务中运行两个方法

Question

In the effort of learning Go a bit better, I am trying to refactor a series of functions which accept a DB connection as the first argument into struct methods and something a bit more "idiomatically" Go.

Right now my "data store" methods are something like this:

func CreateA(db orm.DB, a *A) error {
    db.Exec("INSERT...")
}

func CreateB(db orm.DB, b *B) error {
    db.Exec("INSERT...")
}

These the functions work perfectly fine. orm.DB is the DB interface of go-pg.

Since the two functions accept a db connection I can either pass an actual connection or a transaction (which implements the same interface). I can be sure that both functions issuing SQL INSERTs run in the same transaction, avoiding having inconsistent state in the DB in case either one of them fails.

The trouble started when I decided to read more about how to structure the code a little better and to make it "mockable" in case I need to.

So I googled a bit, read the article Practical Persistence in Go: Organising Database Access and tried to refactor the code to use proper interfaces.

The result is something like this:

type Store {
    CreateA(a *A) error
    CreateB(a *A) error
}

type DB struct {
    orm.DB
}

func NewDBConnection(p *ConnParams) (*DB, error) {
    .... create db connection ...
    return &DB{db}, nil
}

func (db *DB) CreateA(a *A) error {
...
}

func (db *DB) CreateB(b *B) error {
...
}

which allows me to write code like:

db := NewDBConnection()
DB.CreateA(a)
DB.CreateB(b)

instead of:

db := NewDBConnection()
CreateA(db, a)
CreateB(db, b)

The actual issue is that I lost the ability to run the two functions in the same transaction. Before I could do:

pgDB := DB.DB.(*pg.DB) // convert the interface to an actual connection
pgDB.RunInTransaction(func(tx *pg.Tx) error {
    CreateA(tx, a)
    CreateB(tx, b)
})

or something like:

tx := db.DB.Begin()

err = CreateA(tx, a)
err = CreateB(tx, b)

if err != nil {
  tx.Rollback()
} else {
  tx.Commit()
}

which is more or less the same thing.

Since the functions were accepting the common interface between a connection and a transaction I could abstract from my model layer the transaction logic sending down either a full connection or a transaction. This allowed me to decide in the "HTTP handler" when to create a trasaction and when I didn't need to.

Keep in mind that the connection is a global object representing a pool of connections handled automatically by go, so the hack I tried:

pgDB := DB.DB.(*pg.DB) // convert the interface to an actual connection
err = pgDB.RunInTransaction(func(tx *pg.Tx) error {
    DB.DB = tx // replace the connection with a transaction
    DB.CreateA(a)
    DB.CreateB(a)
})

it's clearly a bad idea, because although it works, it works only once because we replace the global connection with a transaction. The following request breaks the server.

Any ideas? I can't find information about this around, probably because I don't know the right keywords being a noob.

Answer 1

I've done something like this in the past (using the standard sql package, you may need to adapt it to your needs):

var ErrNestedTransaction = errors.New("nested transactions are not supported")

// abstraction over sql.TX and sql.DB
// a similar interface seems to be already defined in go-pg. So you may not need this. 
type executor interface {
    Exec(query string, args ...interface{}) (sql.Result, error)
    Query(query string, args ...interface{}) (*sql.Rows, error)
    QueryRow(query string, args ...interface{}) *sql.Row
}

type Store struct {
    // this is the actual connection(pool) to the db which has the Begin() method
    db       *sql.DB
    executor executor
}

func NewStore(dsn string) (*Store, error) {
    db, err := sql.Open("sqlite3", dsn)
    if err != nil {
         return nil, err
    }      
    // the initial store contains just the connection(pool)
    return &Store{db, db}, nil
}

func (s *Store) RunInTransaction(f func(store *Store) error) error {
    if _, ok := s.executor.(*sql.Tx); ok {
        // nested transactions are not supported!
        return ErrNestedTransaction
    }

    tx, err := s.db.Begin()
    if err != nil {
        return err
    }

    transactedStore := &Store{
        s.db,
        tx,
    }

    err = f(transactedStore)
    if err != nil {
        tx.Rollback()
        return err
    }

    return tx.Commit()
}

func (s *Store) CreateA(thing A) error {
    // your implementation
    _, err := s.executor.Exec("INSERT INTO ...", ...)
    return err
}

And then you use it like

// store is a global object
store.RunInTransaction(func(store *Store) error { 
    // this instance of Store uses a transaction to execute the methods
    err := store.CreateA(a)
    if err != nil {
        return err
    }
    return store.CreateB(b)
})

The trick is to use the executor instead of the *sql.DB in your CreateX methods, which allows you to dynamically change the underlying implementation (tx vs. db). However, since there is very little information out there on how to deal with this issue, I can't assure you that this is the "best" solution. Other suggestions are welcome!