create.go 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375
  1. package callbacks
  2. import (
  3. "fmt"
  4. "reflect"
  5. "gorm.io/gorm"
  6. "gorm.io/gorm/clause"
  7. "gorm.io/gorm/schema"
  8. )
  9. func BeforeCreate(db *gorm.DB) {
  10. if db.Error == nil && db.Statement.Schema != nil && !db.Statement.SkipHooks && (db.Statement.Schema.BeforeSave || db.Statement.Schema.BeforeCreate) {
  11. callMethod(db, func(value interface{}, tx *gorm.DB) (called bool) {
  12. if db.Statement.Schema.BeforeSave {
  13. if i, ok := value.(BeforeSaveInterface); ok {
  14. called = true
  15. db.AddError(i.BeforeSave(tx))
  16. }
  17. }
  18. if db.Statement.Schema.BeforeCreate {
  19. if i, ok := value.(BeforeCreateInterface); ok {
  20. called = true
  21. db.AddError(i.BeforeCreate(tx))
  22. }
  23. }
  24. return called
  25. })
  26. }
  27. }
  28. func Create(config *Config) func(db *gorm.DB) {
  29. if config.WithReturning {
  30. return CreateWithReturning
  31. }
  32. return func(db *gorm.DB) {
  33. if db.Error != nil {
  34. return
  35. }
  36. if db.Statement.Schema != nil && !db.Statement.Unscoped {
  37. for _, c := range db.Statement.Schema.CreateClauses {
  38. db.Statement.AddClause(c)
  39. }
  40. }
  41. if db.Statement.SQL.String() == "" {
  42. db.Statement.SQL.Grow(180)
  43. db.Statement.AddClauseIfNotExists(clause.Insert{})
  44. db.Statement.AddClause(ConvertToCreateValues(db.Statement))
  45. db.Statement.Build(db.Statement.BuildClauses...)
  46. }
  47. if !db.DryRun && db.Error == nil {
  48. result, err := db.Statement.ConnPool.ExecContext(db.Statement.Context, db.Statement.SQL.String(), db.Statement.Vars...)
  49. if err != nil {
  50. db.AddError(err)
  51. return
  52. }
  53. db.RowsAffected, _ = result.RowsAffected()
  54. if db.RowsAffected != 0 && db.Statement.Schema != nil &&
  55. db.Statement.Schema.PrioritizedPrimaryField != nil && db.Statement.Schema.PrioritizedPrimaryField.HasDefaultValue {
  56. if insertID, err := result.LastInsertId(); err == nil && insertID > 0 {
  57. switch db.Statement.ReflectValue.Kind() {
  58. case reflect.Slice, reflect.Array:
  59. if config.LastInsertIDReversed {
  60. for i := db.Statement.ReflectValue.Len() - 1; i >= 0; i-- {
  61. rv := db.Statement.ReflectValue.Index(i)
  62. if reflect.Indirect(rv).Kind() != reflect.Struct {
  63. break
  64. }
  65. _, isZero := db.Statement.Schema.PrioritizedPrimaryField.ValueOf(rv)
  66. if isZero {
  67. db.Statement.Schema.PrioritizedPrimaryField.Set(rv, insertID)
  68. insertID -= db.Statement.Schema.PrioritizedPrimaryField.AutoIncrementIncrement
  69. }
  70. }
  71. } else {
  72. for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
  73. rv := db.Statement.ReflectValue.Index(i)
  74. if reflect.Indirect(rv).Kind() != reflect.Struct {
  75. break
  76. }
  77. if _, isZero := db.Statement.Schema.PrioritizedPrimaryField.ValueOf(rv); isZero {
  78. db.Statement.Schema.PrioritizedPrimaryField.Set(rv, insertID)
  79. insertID += db.Statement.Schema.PrioritizedPrimaryField.AutoIncrementIncrement
  80. }
  81. }
  82. }
  83. case reflect.Struct:
  84. if _, isZero := db.Statement.Schema.PrioritizedPrimaryField.ValueOf(db.Statement.ReflectValue); isZero {
  85. db.Statement.Schema.PrioritizedPrimaryField.Set(db.Statement.ReflectValue, insertID)
  86. }
  87. }
  88. } else {
  89. db.AddError(err)
  90. }
  91. }
  92. }
  93. }
  94. }
  95. func CreateWithReturning(db *gorm.DB) {
  96. if db.Error == nil {
  97. if db.Statement.Schema != nil && !db.Statement.Unscoped {
  98. for _, c := range db.Statement.Schema.CreateClauses {
  99. db.Statement.AddClause(c)
  100. }
  101. }
  102. if db.Statement.SQL.String() == "" {
  103. db.Statement.AddClauseIfNotExists(clause.Insert{})
  104. db.Statement.AddClause(ConvertToCreateValues(db.Statement))
  105. db.Statement.Build(db.Statement.BuildClauses...)
  106. }
  107. if sch := db.Statement.Schema; sch != nil && len(sch.FieldsWithDefaultDBValue) > 0 {
  108. db.Statement.WriteString(" RETURNING ")
  109. var (
  110. fields = make([]*schema.Field, len(sch.FieldsWithDefaultDBValue))
  111. values = make([]interface{}, len(sch.FieldsWithDefaultDBValue))
  112. )
  113. for idx, field := range sch.FieldsWithDefaultDBValue {
  114. if idx > 0 {
  115. db.Statement.WriteByte(',')
  116. }
  117. fields[idx] = field
  118. db.Statement.WriteQuoted(field.DBName)
  119. }
  120. if !db.DryRun && db.Error == nil {
  121. db.RowsAffected = 0
  122. rows, err := db.Statement.ConnPool.QueryContext(db.Statement.Context, db.Statement.SQL.String(), db.Statement.Vars...)
  123. if err == nil {
  124. defer rows.Close()
  125. switch db.Statement.ReflectValue.Kind() {
  126. case reflect.Slice, reflect.Array:
  127. c := db.Statement.Clauses["ON CONFLICT"]
  128. onConflict, _ := c.Expression.(clause.OnConflict)
  129. for rows.Next() {
  130. BEGIN:
  131. reflectValue := db.Statement.ReflectValue.Index(int(db.RowsAffected))
  132. if reflect.Indirect(reflectValue).Kind() != reflect.Struct {
  133. break
  134. }
  135. for idx, field := range fields {
  136. fieldValue := field.ReflectValueOf(reflectValue)
  137. if onConflict.DoNothing && !fieldValue.IsZero() {
  138. db.RowsAffected++
  139. if int(db.RowsAffected) >= db.Statement.ReflectValue.Len() {
  140. return
  141. }
  142. goto BEGIN
  143. }
  144. values[idx] = fieldValue.Addr().Interface()
  145. }
  146. db.RowsAffected++
  147. if err := rows.Scan(values...); err != nil {
  148. db.AddError(err)
  149. }
  150. }
  151. case reflect.Struct:
  152. for idx, field := range fields {
  153. values[idx] = field.ReflectValueOf(db.Statement.ReflectValue).Addr().Interface()
  154. }
  155. if rows.Next() {
  156. db.RowsAffected++
  157. db.AddError(rows.Scan(values...))
  158. }
  159. }
  160. } else {
  161. db.AddError(err)
  162. }
  163. }
  164. } else if !db.DryRun && db.Error == nil {
  165. if result, err := db.Statement.ConnPool.ExecContext(db.Statement.Context, db.Statement.SQL.String(), db.Statement.Vars...); err == nil {
  166. db.RowsAffected, _ = result.RowsAffected()
  167. } else {
  168. db.AddError(err)
  169. }
  170. }
  171. }
  172. }
  173. func AfterCreate(db *gorm.DB) {
  174. if db.Error == nil && db.Statement.Schema != nil && !db.Statement.SkipHooks && (db.Statement.Schema.AfterSave || db.Statement.Schema.AfterCreate) {
  175. callMethod(db, func(value interface{}, tx *gorm.DB) (called bool) {
  176. if db.Statement.Schema.AfterSave {
  177. if i, ok := value.(AfterSaveInterface); ok {
  178. called = true
  179. db.AddError(i.AfterSave(tx))
  180. }
  181. }
  182. if db.Statement.Schema.AfterCreate {
  183. if i, ok := value.(AfterCreateInterface); ok {
  184. called = true
  185. db.AddError(i.AfterCreate(tx))
  186. }
  187. }
  188. return called
  189. })
  190. }
  191. }
  192. // ConvertToCreateValues convert to create values
  193. func ConvertToCreateValues(stmt *gorm.Statement) (values clause.Values) {
  194. switch value := stmt.Dest.(type) {
  195. case map[string]interface{}:
  196. values = ConvertMapToValuesForCreate(stmt, value)
  197. case *map[string]interface{}:
  198. values = ConvertMapToValuesForCreate(stmt, *value)
  199. case []map[string]interface{}:
  200. values = ConvertSliceOfMapToValuesForCreate(stmt, value)
  201. case *[]map[string]interface{}:
  202. values = ConvertSliceOfMapToValuesForCreate(stmt, *value)
  203. default:
  204. var (
  205. selectColumns, restricted = stmt.SelectAndOmitColumns(true, false)
  206. _, updateTrackTime = stmt.Get("gorm:update_track_time")
  207. curTime = stmt.DB.NowFunc()
  208. isZero bool
  209. )
  210. stmt.Settings.Delete("gorm:update_track_time")
  211. values = clause.Values{Columns: make([]clause.Column, 0, len(stmt.Schema.DBNames))}
  212. for _, db := range stmt.Schema.DBNames {
  213. if field := stmt.Schema.FieldsByDBName[db]; !field.HasDefaultValue || field.DefaultValueInterface != nil {
  214. if v, ok := selectColumns[db]; (ok && v) || (!ok && (!restricted || field.AutoCreateTime > 0 || field.AutoUpdateTime > 0)) {
  215. values.Columns = append(values.Columns, clause.Column{Name: db})
  216. }
  217. }
  218. }
  219. switch stmt.ReflectValue.Kind() {
  220. case reflect.Slice, reflect.Array:
  221. stmt.SQL.Grow(stmt.ReflectValue.Len() * 18)
  222. values.Values = make([][]interface{}, stmt.ReflectValue.Len())
  223. defaultValueFieldsHavingValue := map[*schema.Field][]interface{}{}
  224. if stmt.ReflectValue.Len() == 0 {
  225. stmt.AddError(gorm.ErrEmptySlice)
  226. return
  227. }
  228. for i := 0; i < stmt.ReflectValue.Len(); i++ {
  229. rv := reflect.Indirect(stmt.ReflectValue.Index(i))
  230. if !rv.IsValid() {
  231. stmt.AddError(fmt.Errorf("slice data #%v is invalid: %w", i, gorm.ErrInvalidData))
  232. return
  233. }
  234. values.Values[i] = make([]interface{}, len(values.Columns))
  235. for idx, column := range values.Columns {
  236. field := stmt.Schema.FieldsByDBName[column.Name]
  237. if values.Values[i][idx], isZero = field.ValueOf(rv); isZero {
  238. if field.DefaultValueInterface != nil {
  239. values.Values[i][idx] = field.DefaultValueInterface
  240. field.Set(rv, field.DefaultValueInterface)
  241. } else if field.AutoCreateTime > 0 || field.AutoUpdateTime > 0 {
  242. field.Set(rv, curTime)
  243. values.Values[i][idx], _ = field.ValueOf(rv)
  244. }
  245. } else if field.AutoUpdateTime > 0 && updateTrackTime {
  246. field.Set(rv, curTime)
  247. values.Values[i][idx], _ = field.ValueOf(rv)
  248. }
  249. }
  250. for _, field := range stmt.Schema.FieldsWithDefaultDBValue {
  251. if v, ok := selectColumns[field.DBName]; (ok && v) || (!ok && !restricted) {
  252. if v, isZero := field.ValueOf(rv); !isZero {
  253. if len(defaultValueFieldsHavingValue[field]) == 0 {
  254. defaultValueFieldsHavingValue[field] = make([]interface{}, stmt.ReflectValue.Len())
  255. }
  256. defaultValueFieldsHavingValue[field][i] = v
  257. }
  258. }
  259. }
  260. }
  261. for field, vs := range defaultValueFieldsHavingValue {
  262. values.Columns = append(values.Columns, clause.Column{Name: field.DBName})
  263. for idx := range values.Values {
  264. if vs[idx] == nil {
  265. values.Values[idx] = append(values.Values[idx], stmt.Dialector.DefaultValueOf(field))
  266. } else {
  267. values.Values[idx] = append(values.Values[idx], vs[idx])
  268. }
  269. }
  270. }
  271. case reflect.Struct:
  272. values.Values = [][]interface{}{make([]interface{}, len(values.Columns))}
  273. for idx, column := range values.Columns {
  274. field := stmt.Schema.FieldsByDBName[column.Name]
  275. if values.Values[0][idx], isZero = field.ValueOf(stmt.ReflectValue); isZero {
  276. if field.DefaultValueInterface != nil {
  277. values.Values[0][idx] = field.DefaultValueInterface
  278. field.Set(stmt.ReflectValue, field.DefaultValueInterface)
  279. } else if field.AutoCreateTime > 0 || field.AutoUpdateTime > 0 {
  280. field.Set(stmt.ReflectValue, curTime)
  281. values.Values[0][idx], _ = field.ValueOf(stmt.ReflectValue)
  282. }
  283. } else if field.AutoUpdateTime > 0 && updateTrackTime {
  284. field.Set(stmt.ReflectValue, curTime)
  285. values.Values[0][idx], _ = field.ValueOf(stmt.ReflectValue)
  286. }
  287. }
  288. for _, field := range stmt.Schema.FieldsWithDefaultDBValue {
  289. if v, ok := selectColumns[field.DBName]; (ok && v) || (!ok && !restricted) {
  290. if v, isZero := field.ValueOf(stmt.ReflectValue); !isZero {
  291. values.Columns = append(values.Columns, clause.Column{Name: field.DBName})
  292. values.Values[0] = append(values.Values[0], v)
  293. }
  294. }
  295. }
  296. default:
  297. stmt.AddError(gorm.ErrInvalidData)
  298. }
  299. }
  300. if c, ok := stmt.Clauses["ON CONFLICT"]; ok {
  301. if onConflict, _ := c.Expression.(clause.OnConflict); onConflict.UpdateAll {
  302. if stmt.Schema != nil && len(values.Columns) >= 1 {
  303. selectColumns, restricted := stmt.SelectAndOmitColumns(true, true)
  304. columns := make([]string, 0, len(values.Columns)-1)
  305. for _, column := range values.Columns {
  306. if field := stmt.Schema.LookUpField(column.Name); field != nil {
  307. if v, ok := selectColumns[field.DBName]; (ok && v) || (!ok && !restricted) {
  308. if !field.PrimaryKey && (!field.HasDefaultValue || field.DefaultValueInterface != nil) && field.AutoCreateTime == 0 {
  309. columns = append(columns, column.Name)
  310. }
  311. }
  312. }
  313. }
  314. onConflict.DoUpdates = clause.AssignmentColumns(columns)
  315. // use primary fields as default OnConflict columns
  316. if len(onConflict.Columns) == 0 {
  317. for _, field := range stmt.Schema.PrimaryFields {
  318. onConflict.Columns = append(onConflict.Columns, clause.Column{Name: field.DBName})
  319. }
  320. }
  321. stmt.AddClause(onConflict)
  322. }
  323. }
  324. }
  325. return values
  326. }