123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366 |
- package callbacks
- import (
- "fmt"
- "reflect"
- "gorm.io/gorm"
- "gorm.io/gorm/clause"
- "gorm.io/gorm/schema"
- )
- func BeforeCreate(db *gorm.DB) {
- if db.Error == nil && db.Statement.Schema != nil && !db.Statement.SkipHooks && (db.Statement.Schema.BeforeSave || db.Statement.Schema.BeforeCreate) {
- callMethod(db, func(value interface{}, tx *gorm.DB) (called bool) {
- if db.Statement.Schema.BeforeSave {
- if i, ok := value.(BeforeSaveInterface); ok {
- called = true
- db.AddError(i.BeforeSave(tx))
- }
- }
- if db.Statement.Schema.BeforeCreate {
- if i, ok := value.(BeforeCreateInterface); ok {
- called = true
- db.AddError(i.BeforeCreate(tx))
- }
- }
- return called
- })
- }
- }
- func Create(config *Config) func(db *gorm.DB) {
- if config.WithReturning {
- return CreateWithReturning
- } else {
- return func(db *gorm.DB) {
- if db.Error == nil {
- if db.Statement.Schema != nil && !db.Statement.Unscoped {
- for _, c := range db.Statement.Schema.CreateClauses {
- db.Statement.AddClause(c)
- }
- }
- if db.Statement.SQL.String() == "" {
- db.Statement.SQL.Grow(180)
- db.Statement.AddClauseIfNotExists(clause.Insert{})
- db.Statement.AddClause(ConvertToCreateValues(db.Statement))
- db.Statement.Build("INSERT", "VALUES", "ON CONFLICT")
- }
- if !db.DryRun && db.Error == nil {
- result, err := db.Statement.ConnPool.ExecContext(db.Statement.Context, db.Statement.SQL.String(), db.Statement.Vars...)
- if err == nil {
- db.RowsAffected, _ = result.RowsAffected()
- if db.RowsAffected > 0 {
- if db.Statement.Schema != nil && db.Statement.Schema.PrioritizedPrimaryField != nil && db.Statement.Schema.PrioritizedPrimaryField.HasDefaultValue {
- if insertID, err := result.LastInsertId(); err == nil && insertID > 0 {
- switch db.Statement.ReflectValue.Kind() {
- case reflect.Slice, reflect.Array:
- if config.LastInsertIDReversed {
- for i := db.Statement.ReflectValue.Len() - 1; i >= 0; i-- {
- rv := db.Statement.ReflectValue.Index(i)
- if reflect.Indirect(rv).Kind() != reflect.Struct {
- break
- }
- _, isZero := db.Statement.Schema.PrioritizedPrimaryField.ValueOf(rv)
- if isZero {
- db.Statement.Schema.PrioritizedPrimaryField.Set(rv, insertID)
- insertID -= db.Statement.Schema.PrioritizedPrimaryField.AutoIncrementIncrement
- }
- }
- } else {
- for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
- rv := db.Statement.ReflectValue.Index(i)
- if reflect.Indirect(rv).Kind() != reflect.Struct {
- break
- }
- if _, isZero := db.Statement.Schema.PrioritizedPrimaryField.ValueOf(rv); isZero {
- db.Statement.Schema.PrioritizedPrimaryField.Set(rv, insertID)
- insertID += db.Statement.Schema.PrioritizedPrimaryField.AutoIncrementIncrement
- }
- }
- }
- case reflect.Struct:
- if _, isZero := db.Statement.Schema.PrioritizedPrimaryField.ValueOf(db.Statement.ReflectValue); isZero {
- db.Statement.Schema.PrioritizedPrimaryField.Set(db.Statement.ReflectValue, insertID)
- }
- }
- } else {
- db.AddError(err)
- }
- }
- }
- } else {
- db.AddError(err)
- }
- }
- }
- }
- }
- }
- func CreateWithReturning(db *gorm.DB) {
- if db.Error == nil {
- if db.Statement.Schema != nil && !db.Statement.Unscoped {
- for _, c := range db.Statement.Schema.CreateClauses {
- db.Statement.AddClause(c)
- }
- }
- if db.Statement.SQL.String() == "" {
- db.Statement.AddClauseIfNotExists(clause.Insert{})
- db.Statement.AddClause(ConvertToCreateValues(db.Statement))
- db.Statement.Build("INSERT", "VALUES", "ON CONFLICT")
- }
- if sch := db.Statement.Schema; sch != nil && len(sch.FieldsWithDefaultDBValue) > 0 {
- db.Statement.WriteString(" RETURNING ")
- var (
- fields = make([]*schema.Field, len(sch.FieldsWithDefaultDBValue))
- values = make([]interface{}, len(sch.FieldsWithDefaultDBValue))
- )
- for idx, field := range sch.FieldsWithDefaultDBValue {
- if idx > 0 {
- db.Statement.WriteByte(',')
- }
- fields[idx] = field
- db.Statement.WriteQuoted(field.DBName)
- }
- if !db.DryRun && db.Error == nil {
- db.RowsAffected = 0
- rows, err := db.Statement.ConnPool.QueryContext(db.Statement.Context, db.Statement.SQL.String(), db.Statement.Vars...)
- if err == nil {
- defer rows.Close()
- switch db.Statement.ReflectValue.Kind() {
- case reflect.Slice, reflect.Array:
- c := db.Statement.Clauses["ON CONFLICT"]
- onConflict, _ := c.Expression.(clause.OnConflict)
- for rows.Next() {
- BEGIN:
- reflectValue := db.Statement.ReflectValue.Index(int(db.RowsAffected))
- if reflect.Indirect(reflectValue).Kind() != reflect.Struct {
- break
- }
- for idx, field := range fields {
- fieldValue := field.ReflectValueOf(reflectValue)
- if onConflict.DoNothing && !fieldValue.IsZero() {
- db.RowsAffected++
- if int(db.RowsAffected) >= db.Statement.ReflectValue.Len() {
- return
- }
- goto BEGIN
- }
- values[idx] = fieldValue.Addr().Interface()
- }
- db.RowsAffected++
- if err := rows.Scan(values...); err != nil {
- db.AddError(err)
- }
- }
- case reflect.Struct:
- for idx, field := range fields {
- values[idx] = field.ReflectValueOf(db.Statement.ReflectValue).Addr().Interface()
- }
- if rows.Next() {
- db.RowsAffected++
- db.AddError(rows.Scan(values...))
- }
- }
- } else {
- db.AddError(err)
- }
- }
- } else if !db.DryRun && db.Error == nil {
- if result, err := db.Statement.ConnPool.ExecContext(db.Statement.Context, db.Statement.SQL.String(), db.Statement.Vars...); err == nil {
- db.RowsAffected, _ = result.RowsAffected()
- } else {
- db.AddError(err)
- }
- }
- }
- }
- func AfterCreate(db *gorm.DB) {
- if db.Error == nil && db.Statement.Schema != nil && !db.Statement.SkipHooks && (db.Statement.Schema.AfterSave || db.Statement.Schema.AfterCreate) {
- callMethod(db, func(value interface{}, tx *gorm.DB) (called bool) {
- if db.Statement.Schema.AfterSave {
- if i, ok := value.(AfterSaveInterface); ok {
- called = true
- db.AddError(i.AfterSave(tx))
- }
- }
- if db.Statement.Schema.AfterCreate {
- if i, ok := value.(AfterCreateInterface); ok {
- called = true
- db.AddError(i.AfterCreate(tx))
- }
- }
- return called
- })
- }
- }
- // ConvertToCreateValues convert to create values
- func ConvertToCreateValues(stmt *gorm.Statement) (values clause.Values) {
- switch value := stmt.Dest.(type) {
- case map[string]interface{}:
- values = ConvertMapToValuesForCreate(stmt, value)
- case *map[string]interface{}:
- values = ConvertMapToValuesForCreate(stmt, *value)
- case []map[string]interface{}:
- values = ConvertSliceOfMapToValuesForCreate(stmt, value)
- case *[]map[string]interface{}:
- values = ConvertSliceOfMapToValuesForCreate(stmt, *value)
- default:
- var (
- selectColumns, restricted = stmt.SelectAndOmitColumns(true, false)
- curTime = stmt.DB.NowFunc()
- isZero bool
- )
- values = clause.Values{Columns: make([]clause.Column, 0, len(stmt.Schema.DBNames))}
- for _, db := range stmt.Schema.DBNames {
- if field := stmt.Schema.FieldsByDBName[db]; !field.HasDefaultValue || field.DefaultValueInterface != nil {
- if v, ok := selectColumns[db]; (ok && v) || (!ok && (!restricted || field.AutoCreateTime > 0 || field.AutoUpdateTime > 0)) {
- values.Columns = append(values.Columns, clause.Column{Name: db})
- }
- }
- }
- switch stmt.ReflectValue.Kind() {
- case reflect.Slice, reflect.Array:
- stmt.SQL.Grow(stmt.ReflectValue.Len() * 18)
- values.Values = make([][]interface{}, stmt.ReflectValue.Len())
- defaultValueFieldsHavingValue := map[*schema.Field][]interface{}{}
- if stmt.ReflectValue.Len() == 0 {
- stmt.AddError(gorm.ErrEmptySlice)
- return
- }
- for i := 0; i < stmt.ReflectValue.Len(); i++ {
- rv := reflect.Indirect(stmt.ReflectValue.Index(i))
- if !rv.IsValid() {
- stmt.AddError(fmt.Errorf("slice data #%v is invalid: %w", i, gorm.ErrInvalidData))
- return
- }
- values.Values[i] = make([]interface{}, len(values.Columns))
- for idx, column := range values.Columns {
- field := stmt.Schema.FieldsByDBName[column.Name]
- if values.Values[i][idx], isZero = field.ValueOf(rv); isZero {
- if field.DefaultValueInterface != nil {
- values.Values[i][idx] = field.DefaultValueInterface
- field.Set(rv, field.DefaultValueInterface)
- } else if field.AutoCreateTime > 0 || field.AutoUpdateTime > 0 {
- field.Set(rv, curTime)
- values.Values[i][idx], _ = field.ValueOf(rv)
- }
- } else if field.AutoUpdateTime > 0 {
- if _, ok := stmt.DB.InstanceGet("gorm:update_track_time"); ok {
- field.Set(rv, curTime)
- values.Values[i][idx], _ = field.ValueOf(rv)
- }
- }
- }
- for _, field := range stmt.Schema.FieldsWithDefaultDBValue {
- if v, ok := selectColumns[field.DBName]; (ok && v) || (!ok && !restricted) {
- if v, isZero := field.ValueOf(rv); !isZero {
- if len(defaultValueFieldsHavingValue[field]) == 0 {
- defaultValueFieldsHavingValue[field] = make([]interface{}, stmt.ReflectValue.Len())
- }
- defaultValueFieldsHavingValue[field][i] = v
- }
- }
- }
- }
- for field, vs := range defaultValueFieldsHavingValue {
- values.Columns = append(values.Columns, clause.Column{Name: field.DBName})
- for idx := range values.Values {
- if vs[idx] == nil {
- values.Values[idx] = append(values.Values[idx], stmt.Dialector.DefaultValueOf(field))
- } else {
- values.Values[idx] = append(values.Values[idx], vs[idx])
- }
- }
- }
- case reflect.Struct:
- values.Values = [][]interface{}{make([]interface{}, len(values.Columns))}
- for idx, column := range values.Columns {
- field := stmt.Schema.FieldsByDBName[column.Name]
- if values.Values[0][idx], isZero = field.ValueOf(stmt.ReflectValue); isZero {
- if field.DefaultValueInterface != nil {
- values.Values[0][idx] = field.DefaultValueInterface
- field.Set(stmt.ReflectValue, field.DefaultValueInterface)
- } else if field.AutoCreateTime > 0 || field.AutoUpdateTime > 0 {
- field.Set(stmt.ReflectValue, curTime)
- values.Values[0][idx], _ = field.ValueOf(stmt.ReflectValue)
- }
- }
- }
- for _, field := range stmt.Schema.FieldsWithDefaultDBValue {
- if v, ok := selectColumns[field.DBName]; (ok && v) || (!ok && !restricted) {
- if v, isZero := field.ValueOf(stmt.ReflectValue); !isZero {
- values.Columns = append(values.Columns, clause.Column{Name: field.DBName})
- values.Values[0] = append(values.Values[0], v)
- }
- }
- }
- default:
- stmt.AddError(gorm.ErrInvalidData)
- }
- }
- if c, ok := stmt.Clauses["ON CONFLICT"]; ok {
- if onConflict, _ := c.Expression.(clause.OnConflict); onConflict.UpdateAll {
- if stmt.Schema != nil && len(values.Columns) > 1 {
- columns := make([]string, 0, len(values.Columns)-1)
- for _, column := range values.Columns {
- if field := stmt.Schema.LookUpField(column.Name); field != nil {
- if !field.PrimaryKey && (!field.HasDefaultValue || field.DefaultValueInterface != nil) && field.AutoCreateTime == 0 {
- columns = append(columns, column.Name)
- }
- }
- }
- onConflict := clause.OnConflict{
- Columns: make([]clause.Column, len(stmt.Schema.PrimaryFieldDBNames)),
- DoUpdates: clause.AssignmentColumns(columns),
- }
- for idx, field := range stmt.Schema.PrimaryFields {
- onConflict.Columns[idx] = clause.Column{Name: field.DBName}
- }
- stmt.AddClause(onConflict)
- }
- }
- }
- return values
- }
|