package repository import ( "context" "fmt" "github.com/glebarez/sqlite" "github.com/go-nunu/nunu-layout-advanced/pkg/log" "github.com/go-nunu/nunu-layout-advanced/pkg/mongo" "github.com/go-nunu/nunu-layout-advanced/pkg/zapgorm2" "github.com/redis/go-redis/v9" "github.com/spf13/viper" "gorm.io/driver/mysql" "gorm.io/driver/postgres" "gorm.io/gorm" gormlogger "gorm.io/gorm/logger" "time" ) const ctxTxKey = "TxKey" type Repository struct { db *gorm.DB //rdb *redis.Client mongodb *mongo.MongoDB logger *log.Logger } func NewRepository( logger *log.Logger, db *gorm.DB, // rdb *redis.Client, mongodb *mongo.MongoDB, ) *Repository { return &Repository{ db: db, //rdb: rdb, mongodb: mongodb, logger: logger, } } type Transaction interface { Transaction(ctx context.Context, fn func(ctx context.Context) error) error } func NewTransaction(r *Repository) Transaction { return r } // DB return tx // If you need to create a Transaction, you must call DB(ctx) and Transaction(ctx,fn) func (r *Repository) DB(ctx context.Context) *gorm.DB { v := ctx.Value(ctxTxKey) if v != nil { if tx, ok := v.(*gorm.DB); ok { return tx } } return r.db.WithContext(ctx) } func (r *Repository) Transaction(ctx context.Context, fn func(ctx context.Context) error) error { return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { ctx = context.WithValue(ctx, ctxTxKey, tx) return fn(ctx) }) } func NewDB(conf *viper.Viper, l *log.Logger) *gorm.DB { var ( db *gorm.DB err error ) driver := conf.GetString("data.db.user.driver") dsn := conf.GetString("data.db.user.dsn") // 读取日志级别配置 logLevelStr := conf.GetString("data.db.user.logLevel") var logLevel gormlogger.LogLevel switch logLevelStr { case "silent": logLevel = gormlogger.Silent case "error": logLevel = gormlogger.Error case "warn": logLevel = gormlogger.Warn case "info": logLevel = gormlogger.Info default: // MySQL 默认只记录警告和错误 if driver == "mysql" { logLevel = gormlogger.Warn } else { logLevel = gormlogger.Info } } logger := zapgorm2.New(l.Logger).LogMode(logLevel) // GORM doc: https://gorm.io/docs/connecting_to_the_database.html switch driver { case "mysql": db, err = gorm.Open(mysql.Open(dsn), &gorm.Config{ Logger: logger, }) case "postgres": db, err = gorm.Open(postgres.New(postgres.Config{ DSN: dsn, PreferSimpleProtocol: true, }), &gorm.Config{ Logger: logger, }) case "sqlite": db, err = gorm.Open(sqlite.Open(dsn), &gorm.Config{ Logger: logger, }) default: panic("unknown db driver") } if err != nil { panic(err) } // Connection Pool config sqlDB, err := db.DB() if err != nil { panic(err) } sqlDB.SetMaxIdleConns(10) sqlDB.SetMaxOpenConns(100) sqlDB.SetConnMaxLifetime(time.Hour) return db } func NewRedis(conf *viper.Viper) *redis.Client { rdb := redis.NewClient(&redis.Options{ Addr: conf.GetString("data.redis.addr"), Password: conf.GetString("data.redis.password"), DB: conf.GetInt("data.redis.db"), }) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() _, err := rdb.Ping(ctx).Result() if err != nil { panic(fmt.Sprintf("redis error: %s", err.Error())) } return rdb } func NewMongoDB(conf *viper.Viper) *mongo.MongoDB { config := &mongo.Config{ URI: conf.GetString("data.mongodb.uri"), Database: conf.GetString("data.mongodb.database"), Timeout: conf.GetDuration("data.mongodb.timeout"), MaxPoolSize: conf.GetUint64("data.mongodb.max_pool_size"), } mongoDB, err := mongo.New(config) if err != nil { panic(fmt.Sprintf("mongodb error: %s", err.Error())) } return mongoDB }