create.go

package sqlserver

import (
	"reflect"

	"gorm.io/gorm"
	"gorm.io/gorm/callbacks"
	"gorm.io/gorm/clause"
	"gorm.io/gorm/schema"
)

func Create(db *gorm.DB) {
	if db.Error != nil {
		return
	}

	if db.Statement.Schema != nil && !db.Statement.Unscoped {
		for _, c := range db.Statement.Schema.CreateClauses {
			db.Statement.AddClause(c)
		}
	}

	hasOutput := false
	if db.Statement.SQL.String() == "" {
		var (
			values                  = callbacks.ConvertToCreateValues(db.Statement)
			c                       = db.Statement.Clauses["ON CONFLICT"]
			onConflict, hasConflict = c.Expression.(clause.OnConflict)
		)

		if hasConflict {
			if len(onConflict.Columns) > 0 {
				columnsMap := map[string]bool{}
				for _, column := range values.Columns {
					columnsMap[column.Name] = true
				}

				for _, conflictColumn := range onConflict.Columns {
					if _, ok := columnsMap[conflictColumn.Name]; !ok {
						hasConflict = false
						break
					}
				}
			} else if len(db.Statement.Schema.PrimaryFields) > 0 {
				columnsMap := map[string]bool{}
				for _, column := range values.Columns {
					columnsMap[column.Name] = true
				}

				for _, field := range db.Statement.Schema.PrimaryFields {
					if _, ok := columnsMap[field.DBName]; !ok {
						hasConflict = false
						break
					}
				}
			} else {
				hasConflict = false
			}
		}

		if hasConflict {
			hasOutput = MergeCreate(db, onConflict, values)
		} else {
			setIdentityInsert := false

			if db.Statement.Schema != nil {
				if field := db.Statement.Schema.PrioritizedPrimaryField; field != nil && field.AutoIncrement {
					switch db.Statement.ReflectValue.Kind() {
					case reflect.Struct:
						_, isZero := field.ValueOf(db.Statement.Context, db.Statement.ReflectValue)
						setIdentityInsert = !isZero
					case reflect.Slice, reflect.Array:
						for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
							obj := db.Statement.ReflectValue.Index(i)
							if reflect.Indirect(obj).Kind() == reflect.Struct {
								_, isZero := field.ValueOf(db.Statement.Context, db.Statement.ReflectValue.Index(i))
								setIdentityInsert = !isZero
							}
							break
						}
					}

					if setIdentityInsert {
						db.Statement.WriteString("SET IDENTITY_INSERT ")
						db.Statement.WriteQuoted(db.Statement.Table)
						db.Statement.WriteString(" ON;")
					}
				}
			}

			db.Statement.AddClauseIfNotExists(clause.Insert{})
			db.Statement.Build("INSERT")
			db.Statement.WriteByte(' ')

			db.Statement.AddClause(values)
			if values, ok := db.Statement.Clauses["VALUES"].Expression.(clause.Values); ok {
				if len(values.Columns) > 0 {
					db.Statement.WriteByte('(')
					for idx, column := range values.Columns {
						if idx > 0 {
							db.Statement.WriteByte(',')
						}
						db.Statement.WriteQuoted(column)
					}
					db.Statement.WriteByte(')')

					hasOutput = outputInserted(db)

					db.Statement.WriteString(" VALUES ")

					for idx, value := range values.Values {
						if idx > 0 {
							db.Statement.WriteByte(',')
						}

						db.Statement.WriteByte('(')
						db.Statement.AddVar(db.Statement, value...)
						db.Statement.WriteByte(')')
					}

					db.Statement.WriteString(";")
				} else {
					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 && db.Error == nil {
		if db.Statement.Schema != nil && hasOutput {
			rows, err := db.Statement.ConnPool.QueryContext(db.Statement.Context, db.Statement.SQL.String(), db.Statement.Vars...)
			if db.AddError(err) == nil {
				defer rows.Close()
				gorm.Scan(rows, db, gorm.ScanUpdate|gorm.ScanOnConflictDoNothing)
				if db.Statement.Result != nil {
					db.Statement.Result.RowsAffected = db.RowsAffected
				}
			}
		} else {
			result, err := db.Statement.ConnPool.ExecContext(db.Statement.Context, db.Statement.SQL.String(), db.Statement.Vars...)
			if db.AddError(err) == nil {
				db.RowsAffected, _ = result.RowsAffected()
				if db.Statement.Result != nil {
					db.Statement.Result.Result = result
					db.Statement.Result.RowsAffected = db.RowsAffected
				}
			}
		}
	}
}

func MergeCreate(db *gorm.DB, onConflict clause.OnConflict, values clause.Values) bool {
	db.Statement.WriteString("MERGE INTO ")
	db.Statement.WriteQuoted(db.Statement.Table)
	db.Statement.WriteString(" USING (VALUES")
	for idx, value := range values.Values {
		if idx > 0 {
			db.Statement.WriteByte(',')
		}

		db.Statement.WriteByte('(')
		db.Statement.AddVar(db.Statement, value...)
		db.Statement.WriteByte(')')
	}

	db.Statement.WriteString(") AS excluded (")
	for idx, column := range values.Columns {
		if idx > 0 {
			db.Statement.WriteByte(',')
		}
		db.Statement.WriteQuoted(column.Name)
	}
	db.Statement.WriteString(") ON ")

	var where clause.Where
	if len(onConflict.Columns) > 0 {
		for _, conflictColumn := range onConflict.Columns {
			where.Exprs = append(where.Exprs, clause.Eq{
				Column: clause.Column{Table: db.Statement.Table, Name: conflictColumn.Name},
				Value:  clause.Column{Table: "excluded", Name: conflictColumn.Name},
			})
		}
	} else if len(db.Statement.Schema.PrimaryFields) > 0 {
		for _, field := range db.Statement.Schema.PrimaryFields {
			where.Exprs = append(where.Exprs, clause.Eq{
				Column: clause.Column{Table: db.Statement.Table, Name: field.DBName},
				Value:  clause.Column{Table: "excluded", Name: field.DBName},
			})
		}
	}
	where.Build(db.Statement)

	if !onConflict.DoNothing && (len(onConflict.DoUpdates) > 0 || onConflict.UpdateAll) {
		if onConflict.UpdateAll {
			conflictColumnMap := make(map[string]bool)
			for _, conflictColumn := range onConflict.Columns {
				conflictColumnMap[conflictColumn.Name] = true
			}

			// Get select and omit columns with requireUpdate=true to respect field.Updatable
			selectColumns, restricted := db.Statement.SelectAndOmitColumns(true, true)
			curTime := db.Statement.DB.NowFunc()

			// Collect columns to update and check if there are any
			type updateColumn struct {
				name           string
				useCurrentTime bool
				autoUpdateTime schema.TimeType
			}
			var updateColumns []updateColumn

			for _, column := range values.Columns {
				if conflictColumnMap[column.Name] {
					continue
				}

				// Check if the field should be updated based on SelectAndOmitColumns result
				if v, ok := selectColumns[column.Name]; (ok && !v) || (restricted && !ok) {
					continue
				}

				// Get the field to check additional properties
				if field := db.Statement.Schema.LookUpField(column.Name); field != nil {
					// Skip primary keys, autoCreateTime fields
					if field.PrimaryKey || field.AutoCreateTime > 0 {
						continue
					}

					// Handle AutoUpdateTime fields with current time
					if field.AutoUpdateTime > 0 {
						updateColumns = append(updateColumns, updateColumn{
							name:           column.Name,
							useCurrentTime: true,
							autoUpdateTime: field.AutoUpdateTime,
						})
						continue
					}
				}

				updateColumns = append(updateColumns, updateColumn{
					name:           column.Name,
					useCurrentTime: false,
				})
			}

			// Only generate WHEN MATCHED clause if there are columns to update
			if len(updateColumns) > 0 {
				db.Statement.WriteString(" WHEN MATCHED THEN UPDATE SET ")
				for i, col := range updateColumns {
					if i > 0 {
						db.Statement.WriteString(", ")
					}
					db.Statement.WriteQuoted(col.name)
					db.Statement.WriteString(" = ")
					if col.useCurrentTime {
						// Use current time for AutoUpdateTime fields
						var timeValue interface{}
						switch col.autoUpdateTime {
						case schema.UnixNanosecond:
							timeValue = curTime.UnixNano()
						case schema.UnixMillisecond:
							timeValue = curTime.UnixMilli()
						case schema.UnixSecond:
							timeValue = curTime.Unix()
						default:
							timeValue = curTime
						}
						db.Statement.AddVar(db.Statement, timeValue)
					} else {
						db.Statement.WriteString("excluded.")
						db.Statement.WriteQuoted(col.name)
					}
				}
			}
		} else {
			db.Statement.WriteString(" WHEN MATCHED THEN UPDATE SET ")
			onConflict.DoUpdates.Build(db.Statement)
		}
	}

	db.Statement.WriteString(" WHEN NOT MATCHED THEN INSERT (")

	written := false
	for _, column := range values.Columns {
		if db.Statement.Schema.PrioritizedPrimaryField == nil || !db.Statement.Schema.PrioritizedPrimaryField.AutoIncrement || db.Statement.Schema.PrioritizedPrimaryField.DBName != column.Name {
			if written {
				db.Statement.WriteByte(',')
			}
			written = true
			db.Statement.WriteQuoted(column.Name)
		}
	}

	db.Statement.WriteString(") VALUES (")

	written = false
	for _, column := range values.Columns {
		if db.Statement.Schema.PrioritizedPrimaryField == nil || !db.Statement.Schema.PrioritizedPrimaryField.AutoIncrement || db.Statement.Schema.PrioritizedPrimaryField.DBName != column.Name {
			if written {
				db.Statement.WriteByte(',')
			}
			written = true
			db.Statement.WriteQuoted(clause.Column{
				Table: "excluded",
				Name:  column.Name,
			})
		}
	}

	db.Statement.WriteString(")")
	hasOutput := outputInserted(db)
	db.Statement.WriteString(";")
	return hasOutput
}

func outputInserted(db *gorm.DB) (hasOutput bool) {
	if db.Statement.Schema != nil && len(db.Statement.Schema.FieldsWithDefaultDBValue) > 0 {
		for _, field := range db.Statement.Schema.FieldsWithDefaultDBValue {
			if hasOutput {
				db.Statement.WriteString(",")
			}
			if field.Readable {
				if !hasOutput {
					db.Statement.WriteString(" OUTPUT INSERTED.")
					hasOutput = true
				} else {
					db.Statement.WriteString(" INSERTED.")
				}
				db.Statement.AddVar(db.Statement, clause.Column{Name: field.DBName})
			}
		}
	}
	return hasOutput
}