spark-server/internal/infrastructure/mysql/mysql.go

// Package mysql opens GORM and writes SQL logs through zap.
package mysql

import (
	"context"
	"fmt"
	"strings"
	"time"

	"go.uber.org/zap"
	mysqldriver "gorm.io/driver/mysql"
	"gorm.io/gorm"
)

// DB is a GORM handle. Connection pool is the standard library pool behind GORM v2.
type DB struct {
	gdb *gorm.DB
}

// Options configure Open, including the underlying sql.DB pool.
type Options struct {
	DSN             string
	SlowThreshold   time.Duration
	LogLevel        string
	ZapLevel        string
	MaxOpenConns    int
	MaxIdleConns    int
	ConnMaxLifetime time.Duration
	ConnMaxIdleTime time.Duration
}

const (
	defaultMaxOpen     = 64
	defaultMaxIdle     = 16
	defaultConnLife    = 30 * time.Minute
	defaultConnIdle    = 10 * time.Minute
	defaultPingTimeout = 10 * time.Second
)

// Open connects to MySQL. An empty DSN skips the connection for local bring-up.
func Open(log *zap.Logger, o *Options) (*DB, error) {
	if o == nil {
		o = &Options{}
	}
	if o.DSN == "" {
		return &DB{}, nil
	}
	applyPoolDefaults(o)
	gdb, err := gorm.Open(mysqldriver.Open(withDialTimeout(o.DSN)), &gorm.Config{
		Logger: newGormLogger(log, o.SlowThreshold, parseGormLevel(o.LogLevel, o.ZapLevel)),
	})
	if err != nil {
		return nil, fmt.Errorf("mysql open: %w", err)
	}
	sqlDB, err := gdb.DB()
	if err != nil {
		return nil, fmt.Errorf("mysql pool: %w", err)
	}
	sqlDB.SetMaxOpenConns(o.MaxOpenConns)
	sqlDB.SetMaxIdleConns(o.MaxIdleConns)
	sqlDB.SetConnMaxLifetime(o.ConnMaxLifetime)
	sqlDB.SetConnMaxIdleTime(o.ConnMaxIdleTime)
	ctx, cancel := context.WithTimeout(context.Background(), defaultPingTimeout)
	defer cancel()
	if err := sqlDB.PingContext(ctx); err != nil {
		_ = sqlDB.Close()
		return nil, fmt.Errorf("mysql ping: %w", err)
	}
	return &DB{gdb: gdb}, nil
}

func applyPoolDefaults(o *Options) {
	if o.MaxOpenConns <= 0 {
		o.MaxOpenConns = defaultMaxOpen
	}
	if o.MaxIdleConns <= 0 {
		o.MaxIdleConns = defaultMaxIdle
	}
	if o.ConnMaxLifetime <= 0 {
		o.ConnMaxLifetime = defaultConnLife
	}
	if o.ConnMaxIdleTime <= 0 {
		o.ConnMaxIdleTime = defaultConnIdle
	}
}

func withDialTimeout(dsn string) string {
	if strings.Contains(dsn, "timeout=") {
		return dsn
	}
	sep := "?"
	if strings.Contains(dsn, "?") {
		sep = "&"
	}
	return dsn + sep + "timeout=10s"
}

// Gorm returns the ORM handle (nil when DSN was empty).
func (d *DB) Gorm() *gorm.DB {
	if d == nil {
		return nil
	}
	return d.gdb
}

func (d *DB) Close() error {
	if d == nil || d.gdb == nil {
		return nil
	}
	sqlDB, err := d.gdb.DB()
	if err != nil {
		return err
	}
	return sqlDB.Close()
}

func (d *DB) Ready() bool { return d != nil && d.gdb != nil }

func (d *DB) Ping(ctx context.Context) error {
	if !d.Ready() {
		return nil
	}
	sqlDB, err := d.gdb.DB()
	if err != nil {
		return err
	}
	return sqlDB.PingContext(ctx)
}