Skip to content

Commit

Permalink
Fix panic bug in migrator due to lack of nil check for stmt.Schema (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
pixelmaxQm authored Apr 26, 2024
1 parent ac59252 commit 7892019
Showing 1 changed file with 47 additions and 19 deletions.
66 changes: 47 additions & 19 deletions migrator/migrator.go
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,11 @@ func (m Migrator) AutoMigrate(values ...interface{}) error {
}
} else {
if err := m.RunWithValue(value, func(stmt *gorm.Statement) error {

if stmt.Schema == nil {
return errors.New("failed to get schema")
}

columnTypes, err := queryTx.Migrator().ColumnTypes(value)
if err != nil {
return err
Expand Down Expand Up @@ -211,6 +216,11 @@ func (m Migrator) CreateTable(values ...interface{}) error {
for _, value := range m.ReorderModels(values, false) {
tx := m.DB.Session(&gorm.Session{})
if err := m.RunWithValue(value, func(stmt *gorm.Statement) (err error) {

if stmt.Schema == nil {
return errors.New("failed to get schema")
}

var (
createTableSQL = "CREATE TABLE ? ("
values = []interface{}{m.CurrentTable(stmt)}
Expand Down Expand Up @@ -363,6 +373,9 @@ func (m Migrator) RenameTable(oldName, newName interface{}) error {
func (m Migrator) AddColumn(value interface{}, name string) error {
return m.RunWithValue(value, func(stmt *gorm.Statement) error {
// avoid using the same name field
if stmt.Schema == nil {
return errors.New("failed to get schema")
}
f := stmt.Schema.LookUpField(name)
if f == nil {
return fmt.Errorf("failed to look up field with name: %s", name)
Expand All @@ -382,8 +395,10 @@ func (m Migrator) AddColumn(value interface{}, name string) error {
// DropColumn drop value's `name` column
func (m Migrator) DropColumn(value interface{}, name string) error {
return m.RunWithValue(value, func(stmt *gorm.Statement) error {
if field := stmt.Schema.LookUpField(name); field != nil {
name = field.DBName
if stmt.Schema != nil {
if field := stmt.Schema.LookUpField(name); field != nil {
name = field.DBName
}
}

return m.DB.Exec(
Expand All @@ -395,13 +410,15 @@ func (m Migrator) DropColumn(value interface{}, name string) error {
// AlterColumn alter value's `field` column' type based on schema definition
func (m Migrator) AlterColumn(value interface{}, field string) error {
return m.RunWithValue(value, func(stmt *gorm.Statement) error {
if field := stmt.Schema.LookUpField(field); field != nil {
fileType := m.FullDataTypeOf(field)
return m.DB.Exec(
"ALTER TABLE ? ALTER COLUMN ? TYPE ?",
m.CurrentTable(stmt), clause.Column{Name: field.DBName}, fileType,
).Error
if stmt.Schema != nil {
if field := stmt.Schema.LookUpField(field); field != nil {
fileType := m.FullDataTypeOf(field)
return m.DB.Exec(
"ALTER TABLE ? ALTER COLUMN ? TYPE ?",
m.CurrentTable(stmt), clause.Column{Name: field.DBName}, fileType,
).Error

}
}
return fmt.Errorf("failed to look up field with name: %s", field)
})
Expand All @@ -413,8 +430,10 @@ func (m Migrator) HasColumn(value interface{}, field string) bool {
m.RunWithValue(value, func(stmt *gorm.Statement) error {
currentDatabase := m.DB.Migrator().CurrentDatabase()
name := field
if field := stmt.Schema.LookUpField(field); field != nil {
name = field.DBName
if stmt.Schema != nil {
if field := stmt.Schema.LookUpField(field); field != nil {
name = field.DBName
}
}

return m.DB.Raw(
Expand All @@ -429,12 +448,14 @@ func (m Migrator) HasColumn(value interface{}, field string) bool {
// RenameColumn rename value's field name from oldName to newName
func (m Migrator) RenameColumn(value interface{}, oldName, newName string) error {
return m.RunWithValue(value, func(stmt *gorm.Statement) error {
if field := stmt.Schema.LookUpField(oldName); field != nil {
oldName = field.DBName
}
if stmt.Schema != nil {
if field := stmt.Schema.LookUpField(oldName); field != nil {
oldName = field.DBName
}

if field := stmt.Schema.LookUpField(newName); field != nil {
newName = field.DBName
if field := stmt.Schema.LookUpField(newName); field != nil {
newName = field.DBName
}
}

return m.DB.Exec(
Expand Down Expand Up @@ -794,6 +815,9 @@ type BuildIndexOptionsInterface interface {
// CreateIndex create index `name`
func (m Migrator) CreateIndex(value interface{}, name string) error {
return m.RunWithValue(value, func(stmt *gorm.Statement) error {
if stmt.Schema == nil {
return errors.New("failed to get schema")
}
if idx := stmt.Schema.LookIndex(name); idx != nil {
opts := m.DB.Migrator().(BuildIndexOptionsInterface).BuildIndexOptions(idx.Fields, stmt)
values := []interface{}{clause.Column{Name: idx.Name}, m.CurrentTable(stmt), opts}
Expand Down Expand Up @@ -826,8 +850,10 @@ func (m Migrator) CreateIndex(value interface{}, name string) error {
// DropIndex drop index `name`
func (m Migrator) DropIndex(value interface{}, name string) error {
return m.RunWithValue(value, func(stmt *gorm.Statement) error {
if idx := stmt.Schema.LookIndex(name); idx != nil {
name = idx.Name
if stmt.Schema != nil {
if idx := stmt.Schema.LookIndex(name); idx != nil {
name = idx.Name
}
}

return m.DB.Exec("DROP INDEX ? ON ?", clause.Column{Name: name}, m.CurrentTable(stmt)).Error
Expand All @@ -839,8 +865,10 @@ func (m Migrator) HasIndex(value interface{}, name string) bool {
var count int64
m.RunWithValue(value, func(stmt *gorm.Statement) error {
currentDatabase := m.DB.Migrator().CurrentDatabase()
if idx := stmt.Schema.LookIndex(name); idx != nil {
name = idx.Name
if stmt.Schema != nil {
if idx := stmt.Schema.LookIndex(name); idx != nil {
name = idx.Name
}
}

return m.DB.Raw(
Expand Down

0 comments on commit 7892019

Please sign in to comment.