53 lines
1.8 KiB
Go
53 lines
1.8 KiB
Go
package gormx
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"sync/atomic"
|
|
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
// GORM can call DB().Close on failed initialization. Do not expose the owner's
|
|
// sql.DB until initialization has succeeded; SQL methods remain fully usable.
|
|
// Dialects using context.Background during version probing are bounded by the
|
|
// startup context, which stops affecting queries as soon as Open succeeds.
|
|
type borrowedPool struct {
|
|
db *sql.DB
|
|
startup atomic.Pointer[startupContext]
|
|
}
|
|
|
|
type startupContext struct{ context.Context }
|
|
|
|
func (p *borrowedPool) operationContext(ctx context.Context) context.Context {
|
|
if startup := p.startup.Load(); startup != nil {
|
|
return startup.Context
|
|
}
|
|
return ctx
|
|
}
|
|
func (p *borrowedPool) GetDBConn() (*sql.DB, error) {
|
|
if p.startup.Load() != nil {
|
|
return nil, gorm.ErrInvalidDB
|
|
}
|
|
return p.db, nil
|
|
}
|
|
func (p *borrowedPool) PrepareContext(ctx context.Context, query string) (*sql.Stmt, error) {
|
|
return p.db.PrepareContext(p.operationContext(ctx), query)
|
|
}
|
|
func (p *borrowedPool) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
|
|
return p.db.ExecContext(p.operationContext(ctx), query, args...)
|
|
}
|
|
func (p *borrowedPool) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
|
|
return p.db.QueryContext(p.operationContext(ctx), query, args...)
|
|
}
|
|
func (p *borrowedPool) QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row {
|
|
return p.db.QueryRowContext(p.operationContext(ctx), query, args...)
|
|
}
|
|
func (p *borrowedPool) BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error) {
|
|
return p.db.BeginTx(p.operationContext(ctx), opts)
|
|
}
|
|
|
|
var _ gorm.ConnPool = (*borrowedPool)(nil)
|
|
var _ gorm.TxBeginner = (*borrowedPool)(nil)
|
|
var _ gorm.GetDBConnector = (*borrowedPool)(nil)
|