JCS-pub/coordinator/internal/db/db.go

package db

import (
	_ "github.com/go-sql-driver/mysql"
	"github.com/sirupsen/logrus"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
)

type DB struct {
	db *gorm.DB
}

func NewDB(cfg *Config) (*DB, error) {
	mydb, err := gorm.Open(mysql.Open(cfg.MakeSourceString()), &gorm.Config{})
	if err != nil {
		logrus.Fatalf("failed to connect to database: %v", err)
	}

	return &DB{
		db: mydb,
	}, nil
}

func (db *DB) DoTx(do func(tx SQLContext) error) error {
	return db.db.Transaction(func(tx *gorm.DB) error {
		return do(SQLContext{tx})
	})
}

func DoTx01[R any](db *DB, do func(tx SQLContext) (R, error)) (R, error) {
	var ret R
	err := db.db.Transaction(func(tx *gorm.DB) error {
		var err error
		ret, err = do(SQLContext{tx})
		return err
	})
	return ret, err
}

func DoTx02[R1, R2 any](db *DB, do func(tx SQLContext) (R1, R2, error)) (R1, R2, error) {
	var r1 R1
	var r2 R2
	err := db.db.Transaction(func(tx *gorm.DB) error {
		var err error
		r1, r2, err = do(SQLContext{tx})
		return err
	})
	return r1, r2, err
}

func DoTx12[T any, R any](db *DB, do func(tx SQLContext, t T) (R, error), t T) (R, error) {
	var ret R
	err := db.db.Transaction(func(tx *gorm.DB) error {
		var err error
		ret, err = do(SQLContext{tx}, t)
		return err
	})
	return ret, err
}

func DoTx22[T1 any, T2 any, R any](db *DB, do func(tx SQLContext, t1 T1, t2 T2) (R, error), t1 T1, t2 T2) (R, error) {
	var ret R
	err := db.db.Transaction(func(tx *gorm.DB) error {
		var err error
		ret, err = do(SQLContext{tx}, t1, t2)
		return err
	})
	return ret, err
}

func DoTx32[T1 any, T2 any, T3 any, R any](db *DB, do func(tx SQLContext, t1 T1, t2 T2, t3 T3) (R, error), t1 T1, t2 T2, t3 T3) (R, error) {
	var ret R
	err := db.db.Transaction(func(tx *gorm.DB) error {
		var err error
		ret, err = do(SQLContext{tx}, t1, t2, t3)
		return err
	})
	return ret, err
}

type SQLContext struct {
	*gorm.DB
}

func (db *DB) DefCtx() SQLContext {
	return SQLContext{db.db}
}