初始化实现

mysqlx.go

@@ -0,0 +1,80 @@
+package mysqlx
+
+import (
+	drivermysql "gorm.io/driver/mysql"
+	"gorm.io/gorm"
+)
+
+// Client is the wrapper around a gorm DB instance.
+type Client struct {
+	db *gorm.DB
+}
+
+// New creates a Client using functional options.
+// It requires either a DSN or a custom gorm Dialector.
+func NewDB(opts ...Option) (*Client, error) {
+	cfg := defaultConfig()
+	for _, opt := range opts {
+		opt(cfg)
+	}
+
+	if cfg.Dialector == nil {
+		dsn, err := buildDSN(cfg)
+		if err != nil {
+			return nil, err
+		}
+		cfg.Dialector = drivermysql.Open(dsn)
+	}
+
+	gormConfig := &gorm.Config{
+		Logger:                                   cfg.Logger,
+		NamingStrategy:                           cfg.NamingStrategy,
+		DisableForeignKeyConstraintWhenMigrating: cfg.DisableForeignKeyConstraintWhenMigrating,
+		SkipDefaultTransaction:                   cfg.SkipDefaultTransaction,
+		DryRun:                                   cfg.DryRun,
+	}
+
+	db, err := gorm.Open(cfg.Dialector, gormConfig)
+	if err != nil {
+		return nil, err
+	}
+
+	sqlDB, err := db.DB()
+	if err != nil {
+		return nil, err
+	}
+
+	if cfg.MaxIdleConns > 0 {
+		sqlDB.SetMaxIdleConns(cfg.MaxIdleConns)
+	}
+	if cfg.MaxOpenConns > 0 {
+		sqlDB.SetMaxOpenConns(cfg.MaxOpenConns)
+	}
+	if cfg.ConnMaxIdleTime > 0 {
+		sqlDB.SetConnMaxIdleTime(cfg.ConnMaxIdleTime)
+	}
+	if cfg.ConnMaxLifetime > 0 {
+		sqlDB.SetConnMaxLifetime(cfg.ConnMaxLifetime)
+	}
+
+	return &Client{db: db}, nil
+}
+
+// DB returns the underlying gorm DB.
+func (c *Client) DB() *gorm.DB {
+	return c.db
+}
+
+// Close closes the underlying sql.DB connection pool.
+func (c *Client) Close() error {
+	sqlDB, err := c.db.DB()
+	if err != nil {
+		return err
+	}
+	return sqlDB.Close()
+}
+
+// AutoMigrate runs gorm AutoMigrate for the provided models.
+func (c *Client) AutoMigrate(models ...interface{}) error {
+	return c.db.AutoMigrate(models...)
+}