Test migrate constraints, close #3986

This commit is contained in:
Jinzhu 2021-01-26 13:39:34 +08:00
parent 59c01b7943
commit 916338a9e1
3 changed files with 97 additions and 34 deletions

View File

@ -451,50 +451,80 @@ func buildConstraint(constraint *schema.Constraint) (sql string, results []inter
return return
} }
func (m Migrator) GuessConstraintAndTable(stmt *gorm.Statement, name string) (_ *schema.Constraint, _ *schema.Check, table string) {
if stmt.Schema == nil {
return nil, nil, stmt.Table
}
checkConstraints := stmt.Schema.ParseCheckConstraints()
if chk, ok := checkConstraints[name]; ok {
return nil, &chk, stmt.Table
}
getTable := func(rel *schema.Relationship) string {
switch rel.Type {
case schema.HasOne, schema.HasMany:
return rel.FieldSchema.Table
case schema.Many2Many:
return rel.JoinTable.Table
}
return stmt.Table
}
for _, rel := range stmt.Schema.Relationships.Relations {
if constraint := rel.ParseConstraint(); constraint != nil && constraint.Name == name {
return constraint, nil, getTable(rel)
}
}
if field := stmt.Schema.LookUpField(name); field != nil {
for _, cc := range checkConstraints {
if cc.Field == field {
return nil, &cc, stmt.Table
}
}
for _, rel := range stmt.Schema.Relationships.Relations {
if constraint := rel.ParseConstraint(); constraint != nil && rel.Field == field {
return constraint, nil, getTable(rel)
}
}
}
return nil, nil, ""
}
func (m Migrator) CreateConstraint(value interface{}, name string) error { func (m Migrator) CreateConstraint(value interface{}, name string) error {
return m.RunWithValue(value, func(stmt *gorm.Statement) error { return m.RunWithValue(value, func(stmt *gorm.Statement) error {
checkConstraints := stmt.Schema.ParseCheckConstraints() constraint, chk, table := m.GuessConstraintAndTable(stmt, name)
if chk, ok := checkConstraints[name]; ok { if chk != nil {
return m.DB.Exec( return m.DB.Exec(
"ALTER TABLE ? ADD CONSTRAINT ? CHECK (?)", "ALTER TABLE ? ADD CONSTRAINT ? CHECK (?)",
m.CurrentTable(stmt), clause.Column{Name: chk.Name}, clause.Expr{SQL: chk.Constraint}, m.CurrentTable(stmt), clause.Column{Name: chk.Name}, clause.Expr{SQL: chk.Constraint},
).Error ).Error
} }
for _, rel := range stmt.Schema.Relationships.Relations { if constraint != nil {
if constraint := rel.ParseConstraint(); constraint != nil && constraint.Name == name { var vars = []interface{}{clause.Table{Name: table}}
sql, values := buildConstraint(constraint) if stmt.TableExpr != nil {
return m.DB.Exec("ALTER TABLE ? ADD "+sql, append([]interface{}{m.CurrentTable(stmt)}, values...)...).Error vars[0] = stmt.TableExpr
} }
sql, values := buildConstraint(constraint)
return m.DB.Exec("ALTER TABLE ? ADD "+sql, append(vars, values...)...).Error
} }
err := fmt.Errorf("failed to create constraint with name %v", name) return nil
if field := stmt.Schema.LookUpField(name); field != nil {
for _, cc := range checkConstraints {
if err = m.DB.Migrator().CreateIndex(value, cc.Name); err != nil {
return err
}
}
for _, rel := range stmt.Schema.Relationships.Relations {
if constraint := rel.ParseConstraint(); constraint != nil && constraint.Field == field {
if err = m.DB.Migrator().CreateIndex(value, constraint.Name); err != nil {
return err
}
}
}
}
return err
}) })
} }
func (m Migrator) DropConstraint(value interface{}, name string) error { func (m Migrator) DropConstraint(value interface{}, name string) error {
return m.RunWithValue(value, func(stmt *gorm.Statement) error { return m.RunWithValue(value, func(stmt *gorm.Statement) error {
return m.DB.Exec( constraint, chk, table := m.GuessConstraintAndTable(stmt, name)
"ALTER TABLE ? DROP CONSTRAINT ?", if constraint != nil {
m.CurrentTable(stmt), clause.Column{Name: name}, name = constraint.Name
).Error } else if chk != nil {
name = chk.Name
}
return m.DB.Exec("ALTER TABLE ? DROP CONSTRAINT ?", clause.Table{Name: table}, clause.Column{Name: name}).Error
}) })
} }
@ -502,9 +532,16 @@ func (m Migrator) HasConstraint(value interface{}, name string) bool {
var count int64 var count int64
m.RunWithValue(value, func(stmt *gorm.Statement) error { m.RunWithValue(value, func(stmt *gorm.Statement) error {
currentDatabase := m.DB.Migrator().CurrentDatabase() currentDatabase := m.DB.Migrator().CurrentDatabase()
constraint, chk, table := m.GuessConstraintAndTable(stmt, name)
if constraint != nil {
name = constraint.Name
} else if chk != nil {
name = chk.Name
}
return m.DB.Raw( return m.DB.Raw(
"SELECT count(*) FROM INFORMATION_SCHEMA.table_constraints WHERE constraint_schema = ? AND table_name = ? AND constraint_name = ?", "SELECT count(*) FROM INFORMATION_SCHEMA.table_constraints WHERE constraint_schema = ? AND table_name = ? AND constraint_name = ?",
currentDatabase, stmt.Table, name, currentDatabase, table, name,
).Row().Scan(&count) ).Row().Scan(&count)
}) })

View File

@ -519,7 +519,7 @@ func (rel *Relationship) ParseConstraint() *Constraint {
} }
for _, ref := range rel.References { for _, ref := range rel.References {
if ref.PrimaryKey != nil { if ref.PrimaryKey != nil && (rel.JoinTable == nil || ref.OwnPrimaryKey) {
constraint.ForeignKeys = append(constraint.ForeignKeys, ref.ForeignKey) constraint.ForeignKeys = append(constraint.ForeignKeys, ref.ForeignKey)
constraint.References = append(constraint.References, ref.PrimaryKey) constraint.References = append(constraint.References, ref.PrimaryKey)
@ -533,10 +533,6 @@ func (rel *Relationship) ParseConstraint() *Constraint {
} }
} }
if rel.JoinTable != nil {
return nil
}
return &constraint return &constraint
} }

View File

@ -323,3 +323,33 @@ func TestMigrateColumns(t *testing.T) {
t.Fatalf("Found deleted column") t.Fatalf("Found deleted column")
} }
} }
func TestMigrateConstraint(t *testing.T) {
if DB.Dialector.Name() == "sqlite" {
t.Skip()
}
names := []string{"Account", "fk_users_account", "Pets", "fk_users_pets", "Company", "fk_users_company", "Manager", "fk_users_manager", "Team", "fk_users_team", "Languages", "fk_users_languages"}
for _, name := range names {
if !DB.Migrator().HasConstraint(&User{}, name) {
DB.Migrator().CreateConstraint(&User{}, name)
}
if err := DB.Migrator().DropConstraint(&User{}, name); err != nil {
t.Fatalf("failed to drop constraint %v, got error %v", name, err)
}
if DB.Migrator().HasConstraint(&User{}, name) {
t.Fatalf("constraint %v should been deleted", name)
}
if err := DB.Migrator().CreateConstraint(&User{}, name); err != nil {
t.Fatalf("failed to create constraint %v, got error %v", name, err)
}
if !DB.Migrator().HasConstraint(&User{}, name) {
t.Fatalf("failed to found constraint %v", name)
}
}
}