122 lines
4.5 KiB
Go
122 lines
4.5 KiB
Go
// Package gormx optionally attaches GORM to a pool owned by GoOmni. It never
|
|
// creates a pool, changes pool limits, migrates tables or closes the owner.
|
|
package gormx
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"maps"
|
|
"reflect"
|
|
"time"
|
|
|
|
"git.yuntuo.hk/k2safe/go-omni/database"
|
|
"git.yuntuo.hk/k2safe/go-omni/logx"
|
|
gormmysql "gorm.io/driver/mysql"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
// DialectorFactory must use the supplied connection pool. It must not create a
|
|
// pool or retain/use the ORM before Open succeeds. For example, return
|
|
// mysql.New(mysql.Config{Conn: conn}). Other dialects can use the same contract.
|
|
type DialectorFactory func(conn gorm.ConnPool) gorm.Dialector
|
|
|
|
type options struct {
|
|
config gorm.Config
|
|
logger *logx.Logger
|
|
logConfig LogConfig
|
|
timeout time.Duration
|
|
}
|
|
type Option func(*options)
|
|
|
|
// WithConfig preserves native GORM behavior, including naming, hooks, prepared
|
|
// statements and transactions. ConnPool/Dialector are supplied by this adapter;
|
|
// automatic Ping is disabled because the owning pool already passed startup Ping.
|
|
// A custom Config.Logger takes precedence over the logx adapter.
|
|
func WithConfig(config gorm.Config) Option { return func(o *options) { o.config = config } }
|
|
func WithLogger(logger *logx.Logger) Option { return func(o *options) { o.logger = logger } }
|
|
func WithLogConfig(config LogConfig) Option { return func(o *options) { o.logConfig = config } }
|
|
func WithInitTimeout(timeout time.Duration) Option { return func(o *options) { o.timeout = timeout } }
|
|
|
|
// Error exposes a safe initialization summary while preserving errors.Is/As.
|
|
type Error struct{ cause error }
|
|
|
|
func (e *Error) Error() string { return "gormx: initialization failed" }
|
|
func (e *Error) Unwrap() error { return e.cause }
|
|
|
|
// MySQL attaches GORM's MySQL dialect, including version negotiation, to the
|
|
// existing pool. Configure TLS, timeouts and connection limits on that pool.
|
|
func MySQL(ctx context.Context, pool *database.DB, opts ...Option) (*gorm.DB, error) {
|
|
return Open(ctx, pool, func(conn gorm.ConnPool) gorm.Dialector {
|
|
return gormmysql.New(gormmysql.Config{Conn: conn})
|
|
}, opts...)
|
|
}
|
|
|
|
// Open borrows pool. On failure (including dialect initialization failure),
|
|
// the original pool remains owned by the caller. No ORM Close is required.
|
|
// The returned ORM does not retain the startup context: each request should use
|
|
// WithContext(ctx). Never close the sql.DB returned by orm.DB() independently.
|
|
func Open(ctx context.Context, pool *database.DB, factory DialectorFactory, opts ...Option) (*gorm.DB, error) {
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return nil, &Error{cause: err}
|
|
}
|
|
if pool == nil || pool.DB == nil || factory == nil {
|
|
return nil, errors.New("gormx: pool and dialector factory are required")
|
|
}
|
|
timeout := pool.PoolConfig().PingTimeout
|
|
if timeout == 0 {
|
|
timeout = 5 * time.Second
|
|
}
|
|
o := options{timeout: timeout, logConfig: DefaultLogConfig(), logger: logx.FromContext(ctx)}
|
|
for _, option := range opts {
|
|
if option != nil {
|
|
option(&o)
|
|
}
|
|
}
|
|
if o.timeout <= 0 {
|
|
return nil, errors.New("gormx: initialization timeout must be positive")
|
|
}
|
|
if o.config.ConnPool != nil || o.config.Dialector != nil {
|
|
return nil, errors.New("gormx: configure the existing pool through the dialector factory")
|
|
}
|
|
logger, err := NewLogger(o.logger, o.logConfig)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
config := o.config
|
|
config.Plugins = maps.Clone(config.Plugins)
|
|
config.ClauseBuilders = maps.Clone(config.ClauseBuilders)
|
|
if config.Logger == nil {
|
|
config.Logger = logger
|
|
}
|
|
config.DisableAutomaticPing = true
|
|
startup, cancel := context.WithTimeout(ctx, o.timeout)
|
|
defer cancel()
|
|
borrowed := &borrowedPool{db: pool.DB}
|
|
borrowed.startup.Store(&startupContext{startup})
|
|
config.ConnPool = borrowed
|
|
dialect := factory(borrowed)
|
|
if dialect == nil || (reflect.ValueOf(dialect).Kind() == reflect.Pointer && reflect.ValueOf(dialect).IsNil()) {
|
|
return nil, errors.New("gormx: dialector factory returned nil")
|
|
}
|
|
orm, err := gorm.Open(dialect, &config)
|
|
if err != nil {
|
|
return nil, &Error{cause: err}
|
|
}
|
|
if err = startup.Err(); err != nil {
|
|
return nil, &Error{cause: err}
|
|
}
|
|
// Prepared statements are GORM's own wrapper over the very same pool.
|
|
conn := orm.ConnPool
|
|
if prepared, ok := conn.(*gorm.PreparedStmtDB); ok {
|
|
conn = prepared.ConnPool
|
|
}
|
|
if actual, ok := conn.(*borrowedPool); !ok || actual != borrowed {
|
|
return nil, errors.New("gormx: dialector or plugin replaced the supplied pool")
|
|
}
|
|
borrowed.startup.Store(nil)
|
|
return orm, nil
|
|
}
|