mirror of https://github.com/go-gorm/gorm.git
312 lines
6.8 KiB
Go
312 lines
6.8 KiB
Go
package gorm
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
|
|
"gorm.io/gorm/clause"
|
|
"gorm.io/gorm/logger"
|
|
"gorm.io/gorm/schema"
|
|
)
|
|
|
|
// Config GORM config
|
|
type Config struct {
|
|
// GORM perform single create, update, delete operations in transactions by default to ensure database data integrity
|
|
// You can disable it by setting `SkipDefaultTransaction` to true
|
|
SkipDefaultTransaction bool
|
|
// NamingStrategy tables, columns naming strategy
|
|
NamingStrategy schema.Namer
|
|
// Logger
|
|
Logger logger.Interface
|
|
// NowFunc the function to be used when creating a new timestamp
|
|
NowFunc func() time.Time
|
|
// DryRun generate sql without execute
|
|
DryRun bool
|
|
// PrepareStmt executes the given query in cached statement
|
|
PrepareStmt bool
|
|
// DisableAutomaticPing
|
|
DisableAutomaticPing bool
|
|
// DisableForeignKeyConstraintWhenMigrating
|
|
DisableForeignKeyConstraintWhenMigrating bool
|
|
|
|
// ClauseBuilders clause builder
|
|
ClauseBuilders map[string]clause.ClauseBuilder
|
|
// ConnPool db conn pool
|
|
ConnPool ConnPool
|
|
// Dialector database dialector
|
|
Dialector
|
|
|
|
callbacks *callbacks
|
|
cacheStore *sync.Map
|
|
}
|
|
|
|
// DB GORM DB definition
|
|
type DB struct {
|
|
*Config
|
|
Error error
|
|
RowsAffected int64
|
|
Statement *Statement
|
|
clone int
|
|
}
|
|
|
|
// Session session config when create session with Session() method
|
|
type Session struct {
|
|
DryRun bool
|
|
PrepareStmt bool
|
|
WithConditions bool
|
|
Context context.Context
|
|
Logger logger.Interface
|
|
NowFunc func() time.Time
|
|
}
|
|
|
|
// Open initialize db session based on dialector
|
|
func Open(dialector Dialector, config *Config) (db *DB, err error) {
|
|
if config == nil {
|
|
config = &Config{}
|
|
}
|
|
|
|
if config.NamingStrategy == nil {
|
|
config.NamingStrategy = schema.NamingStrategy{}
|
|
}
|
|
|
|
if config.Logger == nil {
|
|
config.Logger = logger.Default
|
|
}
|
|
|
|
if config.NowFunc == nil {
|
|
config.NowFunc = func() time.Time { return time.Now().Local() }
|
|
}
|
|
|
|
if dialector != nil {
|
|
config.Dialector = dialector
|
|
}
|
|
|
|
if config.cacheStore == nil {
|
|
config.cacheStore = &sync.Map{}
|
|
}
|
|
|
|
db = &DB{Config: config, clone: 1}
|
|
|
|
db.callbacks = initializeCallbacks(db)
|
|
|
|
if config.ClauseBuilders == nil {
|
|
config.ClauseBuilders = map[string]clause.ClauseBuilder{}
|
|
}
|
|
|
|
if config.Dialector != nil {
|
|
err = config.Dialector.Initialize(db)
|
|
}
|
|
|
|
if config.PrepareStmt {
|
|
db.ConnPool = &PreparedStmtDB{
|
|
ConnPool: db.ConnPool,
|
|
Stmts: map[string]*sql.Stmt{},
|
|
}
|
|
}
|
|
|
|
db.Statement = &Statement{
|
|
DB: db,
|
|
ConnPool: db.ConnPool,
|
|
Context: context.Background(),
|
|
Clauses: map[string]clause.Clause{},
|
|
}
|
|
|
|
if err == nil && !config.DisableAutomaticPing {
|
|
if pinger, ok := db.ConnPool.(interface{ Ping() error }); ok {
|
|
err = pinger.Ping()
|
|
}
|
|
}
|
|
|
|
if err != nil {
|
|
config.Logger.Error(context.Background(), "failed to initialize database, got error %v", err)
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
// Session create new db session
|
|
func (db *DB) Session(config *Session) *DB {
|
|
var (
|
|
txConfig = *db.Config
|
|
tx = &DB{
|
|
Config: &txConfig,
|
|
Statement: db.Statement,
|
|
clone: 1,
|
|
}
|
|
)
|
|
|
|
if config.Context != nil {
|
|
tx.Statement = tx.Statement.clone()
|
|
tx.Statement.DB = tx
|
|
tx.Statement.Context = config.Context
|
|
}
|
|
|
|
if config.PrepareStmt {
|
|
tx.Statement.ConnPool = &PreparedStmtDB{
|
|
ConnPool: db.Config.ConnPool,
|
|
Stmts: map[string]*sql.Stmt{},
|
|
}
|
|
}
|
|
|
|
if config.WithConditions {
|
|
tx.clone = 2
|
|
}
|
|
|
|
if config.DryRun {
|
|
tx.Config.DryRun = true
|
|
}
|
|
|
|
if config.Logger != nil {
|
|
tx.Config.Logger = config.Logger
|
|
}
|
|
|
|
if config.NowFunc != nil {
|
|
tx.Config.NowFunc = config.NowFunc
|
|
}
|
|
|
|
return tx
|
|
}
|
|
|
|
// WithContext change current instance db's context to ctx
|
|
func (db *DB) WithContext(ctx context.Context) *DB {
|
|
return db.Session(&Session{WithConditions: true, Context: ctx})
|
|
}
|
|
|
|
// Debug start debug mode
|
|
func (db *DB) Debug() (tx *DB) {
|
|
return db.Session(&Session{
|
|
WithConditions: true,
|
|
Logger: db.Logger.LogMode(logger.Info),
|
|
})
|
|
}
|
|
|
|
// Set store value with key into current db instance's context
|
|
func (db *DB) Set(key string, value interface{}) *DB {
|
|
tx := db.getInstance()
|
|
tx.Statement.Settings.Store(key, value)
|
|
return tx
|
|
}
|
|
|
|
// Get get value with key from current db instance's context
|
|
func (db *DB) Get(key string) (interface{}, bool) {
|
|
return db.Statement.Settings.Load(key)
|
|
}
|
|
|
|
// InstanceSet store value with key into current db instance's context
|
|
func (db *DB) InstanceSet(key string, value interface{}) *DB {
|
|
tx := db.getInstance()
|
|
tx.Statement.Settings.Store(fmt.Sprintf("%p", tx.Statement)+key, value)
|
|
return tx
|
|
}
|
|
|
|
// InstanceGet get value with key from current db instance's context
|
|
func (db *DB) InstanceGet(key string) (interface{}, bool) {
|
|
return db.Statement.Settings.Load(fmt.Sprintf("%p", db.Statement) + key)
|
|
}
|
|
|
|
// Callback returns callback manager
|
|
func (db *DB) Callback() *callbacks {
|
|
return db.callbacks
|
|
}
|
|
|
|
// AddError add error to db
|
|
func (db *DB) AddError(err error) error {
|
|
if db.Error == nil {
|
|
db.Error = err
|
|
} else if err != nil {
|
|
db.Error = fmt.Errorf("%v; %w", db.Error, err)
|
|
}
|
|
return db.Error
|
|
}
|
|
|
|
// DB returns `*sql.DB`
|
|
func (db *DB) DB() (*sql.DB, error) {
|
|
connPool := db.ConnPool
|
|
|
|
if stmtDB, ok := connPool.(*PreparedStmtDB); ok {
|
|
connPool = stmtDB.ConnPool
|
|
}
|
|
|
|
if sqldb, ok := connPool.(*sql.DB); ok {
|
|
return sqldb, nil
|
|
}
|
|
|
|
return nil, errors.New("invalid db")
|
|
}
|
|
|
|
func (db *DB) getInstance() *DB {
|
|
if db.clone > 0 {
|
|
tx := &DB{Config: db.Config}
|
|
|
|
if db.clone == 1 {
|
|
// clone with new statement
|
|
tx.Statement = &Statement{
|
|
DB: tx,
|
|
ConnPool: db.Statement.ConnPool,
|
|
Context: db.Statement.Context,
|
|
Clauses: map[string]clause.Clause{},
|
|
}
|
|
} else {
|
|
// with clone statement
|
|
tx.Statement = db.Statement.clone()
|
|
tx.Statement.DB = tx
|
|
}
|
|
|
|
return tx
|
|
}
|
|
|
|
return db
|
|
}
|
|
|
|
func Expr(expr string, args ...interface{}) clause.Expr {
|
|
return clause.Expr{SQL: expr, Vars: args}
|
|
}
|
|
|
|
func (db *DB) SetupJoinTable(model interface{}, field string, joinTable interface{}) error {
|
|
var (
|
|
tx = db.getInstance()
|
|
stmt = tx.Statement
|
|
modelSchema, joinSchema *schema.Schema
|
|
)
|
|
|
|
if err := stmt.Parse(model); err == nil {
|
|
modelSchema = stmt.Schema
|
|
} else {
|
|
return err
|
|
}
|
|
|
|
if err := stmt.Parse(joinTable); err == nil {
|
|
joinSchema = stmt.Schema
|
|
} else {
|
|
return err
|
|
}
|
|
|
|
if relation, ok := modelSchema.Relationships.Relations[field]; ok && relation.JoinTable != nil {
|
|
for _, ref := range relation.References {
|
|
if f := joinSchema.LookUpField(ref.ForeignKey.DBName); f != nil {
|
|
f.DataType = ref.ForeignKey.DataType
|
|
ref.ForeignKey = f
|
|
} else {
|
|
return fmt.Errorf("missing field %v for join table", ref.ForeignKey.DBName)
|
|
}
|
|
}
|
|
|
|
for name, rel := range relation.JoinTable.Relationships.Relations {
|
|
if _, ok := joinSchema.Relationships.Relations[name]; !ok {
|
|
rel.Schema = joinSchema
|
|
joinSchema.Relationships.Relations[name] = rel
|
|
}
|
|
}
|
|
|
|
relation.JoinTable = joinSchema
|
|
} else {
|
|
return fmt.Errorf("failed to found relation: %v", field)
|
|
}
|
|
|
|
return nil
|
|
}
|