diff --git a/dialects/mssql/create.go b/dialects/mssql/create.go index 6820bb7b..84732427 100644 --- a/dialects/mssql/create.go +++ b/dialects/mssql/create.go @@ -2,6 +2,7 @@ package mssql import ( "reflect" + "sort" "github.com/jinzhu/gorm" "github.com/jinzhu/gorm/callbacks" @@ -17,10 +18,35 @@ func Create(db *gorm.DB) { } if db.Statement.SQL.String() == "" { + setIdentityInsert := false c := db.Statement.Clauses["ON CONFLICT"] onConflict, hasConflict := c.Expression.(clause.OnConflict) - if hasConflict { + if field := db.Statement.Schema.PrioritizedPrimaryField; field != nil { + setIdentityInsert = false + switch db.Statement.ReflectValue.Kind() { + case reflect.Struct: + _, isZero := field.ValueOf(db.Statement.ReflectValue) + setIdentityInsert = !isZero + case reflect.Slice: + for i := 0; i < db.Statement.ReflectValue.Len(); i++ { + _, isZero := field.ValueOf(db.Statement.ReflectValue.Index(i)) + setIdentityInsert = !isZero + break + } + } + + if setIdentityInsert && (field.DataType == schema.Int || field.DataType == schema.Uint) { + setIdentityInsert = true + db.Statement.WriteString("SET IDENTITY_INSERT ") + db.Statement.WriteQuoted(db.Statement.Table) + db.Statement.WriteString(" ON;") + } else { + setIdentityInsert = false + } + } + + if hasConflict && len(db.Statement.Schema.PrimaryFields) > 0 { MergeCreate(db, onConflict) } else { db.Statement.AddClauseIfNotExists(clause.Insert{Table: clause.Table{Name: db.Statement.Table}}) @@ -55,10 +81,16 @@ func Create(db *gorm.DB) { db.Statement.WriteString(";") } else { - db.Statement.WriteString("DEFAULT VALUES") + db.Statement.WriteString("DEFAULT VALUES;") } } } + + if setIdentityInsert { + db.Statement.WriteString("SET IDENTITY_INSERT ") + db.Statement.WriteQuoted(db.Statement.Table) + db.Statement.WriteString(" OFF;") + } } if !db.DryRun { @@ -67,25 +99,32 @@ func Create(db *gorm.DB) { if err == nil { defer rows.Close() - switch db.Statement.ReflectValue.Kind() { - case reflect.Slice, reflect.Array: - if len(db.Statement.Schema.PrimaryFields) > 0 { - values := make([]interface{}, len(db.Statement.Schema.PrimaryFields)) + if len(db.Statement.Schema.FieldsWithDefaultDBValue) > 0 { + sortedKeys := []string{} + for _, field := range db.Statement.Schema.FieldsWithDefaultDBValue { + sortedKeys = append(sortedKeys, field.DBName) + } + sort.Strings(sortedKeys) + returnningFields := make([]*schema.Field, len(sortedKeys)) + for idx, key := range sortedKeys { + returnningFields[idx] = db.Statement.Schema.LookUpField(key) + } + + values := make([]interface{}, len(returnningFields)) + + switch db.Statement.ReflectValue.Kind() { + case reflect.Slice, reflect.Array: for rows.Next() { - for idx, field := range db.Statement.Schema.PrimaryFields { + for idx, field := range returnningFields { values[idx] = field.ReflectValueOf(db.Statement.ReflectValue.Index(int(db.RowsAffected))).Addr().Interface() } db.RowsAffected++ db.AddError(rows.Scan(values...)) } - } - case reflect.Struct: - if len(db.Statement.Schema.PrimaryFields) > 0 { - values := make([]interface{}, len(db.Statement.Schema.PrimaryFields)) - - for idx, field := range db.Statement.Schema.PrimaryFields { + case reflect.Struct: + for idx, field := range returnningFields { values[idx] = field.ReflectValueOf(db.Statement.ReflectValue).Addr().Interface() } @@ -103,16 +142,6 @@ func Create(db *gorm.DB) { func MergeCreate(db *gorm.DB, onConflict clause.OnConflict) { values := callbacks.ConvertToCreateValues(db.Statement) - setIdentityInsert := false - - if field := db.Statement.Schema.PrioritizedPrimaryField; field != nil { - if field.DataType == schema.Int || field.DataType == schema.Uint { - setIdentityInsert = true - db.Statement.WriteString("SET IDENTITY_INSERT ") - db.Statement.WriteQuoted(db.Statement.Table) - db.Statement.WriteString("ON;") - } - } db.Statement.WriteString("MERGE INTO ") db.Statement.WriteQuoted(db.Statement.Table) @@ -174,23 +203,23 @@ func MergeCreate(db *gorm.DB, onConflict clause.OnConflict) { db.Statement.WriteString(")") outputInserted(db) db.Statement.WriteString(";") - - if setIdentityInsert { - db.Statement.WriteString("SET IDENTITY_INSERT ") - db.Statement.WriteQuoted(db.Statement.Table) - db.Statement.WriteString("OFF;") - } } func outputInserted(db *gorm.DB) { - if len(db.Statement.Schema.PrimaryFields) > 0 { - db.Statement.WriteString(" OUTPUT ") - for idx, field := range db.Statement.Schema.PrimaryFields { + if len(db.Statement.Schema.FieldsWithDefaultDBValue) > 0 { + sortedKeys := []string{} + for _, field := range db.Statement.Schema.FieldsWithDefaultDBValue { + sortedKeys = append(sortedKeys, field.DBName) + } + sort.Strings(sortedKeys) + + db.Statement.WriteString(" OUTPUT") + for idx, key := range sortedKeys { if idx > 0 { db.Statement.WriteString(",") } db.Statement.WriteString(" INSERTED.") - db.Statement.AddVar(db.Statement, clause.Column{Name: field.DBName}) + db.Statement.AddVar(db.Statement, clause.Column{Name: key}) } } }