go orm with transaction.

go-orm-with-transaction.go

package main

import (
	"errors"
	"fmt"
	"sync"

	_ "github.com/go-sql-driver/mysql"
	"github.com/jinzhu/gorm"
)

var db *gorm.DB

// Transaction .
type Transaction struct {
	once     sync.Once
	rollback bool
	DB       *gorm.DB
}

// NewTransaction .
func NewTransaction(db *gorm.DB) (*gorm.DB, *Transaction) {
	tx := &Transaction{DB: db}
	return db, tx
}

// Close .
func (t *Transaction) Close() {
	t.once.Do(func() {
		if t.rollback {
			t.DB.Rollback()
		} else {
			t.DB.Commit()
		}
	})
}

// Fail .
func (t *Transaction) Fail() {
	t.rollback = true
}

// User .
type User struct {
	ID       int `gorm:"primary_key"`
	Username string
	Password string
}

func init() {
	db, _ = gorm.Open("mysql", "root:@tcp(localhost:3306)/test?charset=utf8mb4&parseTime=true&loc=Local")
	db.SingularTable(true)
	db.LogMode(true)
}

// CreateUser .
func CreateUser(db *gorm.DB, user User) error {
	_, tx := NewTransaction(db.Begin())
	defer tx.Close()

	err := tx.DB.Create(&user).Error
	return err
}

// CreateUserWithConn .
func CreateUserWithConn(db *gorm.DB, user User) error {
	conn, tx := NewTransaction(db.Begin())
	defer tx.Close()

	var u User
	conn.Model(User{}).Where("username = ?", user.Username).First(&u)
	fmt.Println("found user without transaction conn: ", u)

	tx.Fail()
	return errors.New("create user faild")
}

func main() {
	user := User{Username: "foo", Password: "bar"}

	_ = CreateUser(db, user)

	err := CreateUserWithConn(db, user)
	if err != nil {
		fmt.Println("create user failed, db rollback.")
	}
}