callbacks.go 8.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360
  1. package gorm
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "reflect"
  7. "sort"
  8. "time"
  9. "gorm.io/gorm/schema"
  10. "gorm.io/gorm/utils"
  11. )
  12. func initializeCallbacks(db *DB) *callbacks {
  13. return &callbacks{
  14. processors: map[string]*processor{
  15. "create": {db: db},
  16. "query": {db: db},
  17. "update": {db: db},
  18. "delete": {db: db},
  19. "row": {db: db},
  20. "raw": {db: db},
  21. },
  22. }
  23. }
  24. // callbacks gorm callbacks manager
  25. type callbacks struct {
  26. processors map[string]*processor
  27. }
  28. type processor struct {
  29. db *DB
  30. Clauses []string
  31. fns []func(*DB)
  32. callbacks []*callback
  33. }
  34. type callback struct {
  35. name string
  36. before string
  37. after string
  38. remove bool
  39. replace bool
  40. match func(*DB) bool
  41. handler func(*DB)
  42. processor *processor
  43. }
  44. func (cs *callbacks) Create() *processor {
  45. return cs.processors["create"]
  46. }
  47. func (cs *callbacks) Query() *processor {
  48. return cs.processors["query"]
  49. }
  50. func (cs *callbacks) Update() *processor {
  51. return cs.processors["update"]
  52. }
  53. func (cs *callbacks) Delete() *processor {
  54. return cs.processors["delete"]
  55. }
  56. func (cs *callbacks) Row() *processor {
  57. return cs.processors["row"]
  58. }
  59. func (cs *callbacks) Raw() *processor {
  60. return cs.processors["raw"]
  61. }
  62. func (p *processor) Execute(db *DB) *DB {
  63. // call scopes
  64. for len(db.Statement.scopes) > 0 {
  65. db = db.executeScopes()
  66. }
  67. var (
  68. curTime = time.Now()
  69. stmt = db.Statement
  70. resetBuildClauses bool
  71. )
  72. if len(stmt.BuildClauses) == 0 {
  73. stmt.BuildClauses = p.Clauses
  74. resetBuildClauses = true
  75. }
  76. if optimizer, ok := db.Statement.Dest.(StatementModifier); ok {
  77. optimizer.ModifyStatement(stmt)
  78. }
  79. // assign model values
  80. if stmt.Model == nil {
  81. stmt.Model = stmt.Dest
  82. } else if stmt.Dest == nil {
  83. stmt.Dest = stmt.Model
  84. }
  85. // parse model values
  86. if stmt.Model != nil {
  87. if err := stmt.Parse(stmt.Model); err != nil && (!errors.Is(err, schema.ErrUnsupportedDataType) || (stmt.Table == "" && stmt.TableExpr == nil && stmt.SQL.Len() == 0)) {
  88. if errors.Is(err, schema.ErrUnsupportedDataType) && stmt.Table == "" && stmt.TableExpr == nil {
  89. db.AddError(fmt.Errorf("%w: Table not set, please set it like: db.Model(&user) or db.Table(\"users\")", err))
  90. } else {
  91. db.AddError(err)
  92. }
  93. }
  94. }
  95. // assign stmt.ReflectValue
  96. if stmt.Dest != nil {
  97. stmt.ReflectValue = reflect.ValueOf(stmt.Dest)
  98. for stmt.ReflectValue.Kind() == reflect.Ptr {
  99. if stmt.ReflectValue.IsNil() && stmt.ReflectValue.CanAddr() {
  100. stmt.ReflectValue.Set(reflect.New(stmt.ReflectValue.Type().Elem()))
  101. }
  102. stmt.ReflectValue = stmt.ReflectValue.Elem()
  103. }
  104. if !stmt.ReflectValue.IsValid() {
  105. db.AddError(ErrInvalidValue)
  106. }
  107. }
  108. for _, f := range p.fns {
  109. f(db)
  110. }
  111. if stmt.SQL.Len() > 0 {
  112. db.Logger.Trace(stmt.Context, curTime, func() (string, int64) {
  113. sql, vars := stmt.SQL.String(), stmt.Vars
  114. if filter, ok := db.Logger.(ParamsFilter); ok {
  115. sql, vars = filter.ParamsFilter(stmt.Context, stmt.SQL.String(), stmt.Vars...)
  116. }
  117. return db.Dialector.Explain(sql, vars...), db.RowsAffected
  118. }, db.Error)
  119. }
  120. if !stmt.DB.DryRun {
  121. stmt.SQL.Reset()
  122. stmt.Vars = nil
  123. }
  124. if resetBuildClauses {
  125. stmt.BuildClauses = nil
  126. }
  127. return db
  128. }
  129. func (p *processor) Get(name string) func(*DB) {
  130. for i := len(p.callbacks) - 1; i >= 0; i-- {
  131. if v := p.callbacks[i]; v.name == name && !v.remove {
  132. return v.handler
  133. }
  134. }
  135. return nil
  136. }
  137. func (p *processor) Before(name string) *callback {
  138. return &callback{before: name, processor: p}
  139. }
  140. func (p *processor) After(name string) *callback {
  141. return &callback{after: name, processor: p}
  142. }
  143. func (p *processor) Match(fc func(*DB) bool) *callback {
  144. return &callback{match: fc, processor: p}
  145. }
  146. func (p *processor) Register(name string, fn func(*DB)) error {
  147. return (&callback{processor: p}).Register(name, fn)
  148. }
  149. func (p *processor) Remove(name string) error {
  150. return (&callback{processor: p}).Remove(name)
  151. }
  152. func (p *processor) Replace(name string, fn func(*DB)) error {
  153. return (&callback{processor: p}).Replace(name, fn)
  154. }
  155. func (p *processor) compile() (err error) {
  156. var callbacks []*callback
  157. removedMap := map[string]bool{}
  158. for _, callback := range p.callbacks {
  159. if callback.match == nil || callback.match(p.db) {
  160. callbacks = append(callbacks, callback)
  161. }
  162. if callback.remove {
  163. removedMap[callback.name] = true
  164. }
  165. }
  166. if len(removedMap) > 0 {
  167. callbacks = removeCallbacks(callbacks, removedMap)
  168. }
  169. p.callbacks = callbacks
  170. if p.fns, err = sortCallbacks(p.callbacks); err != nil {
  171. p.db.Logger.Error(context.Background(), "Got error when compile callbacks, got %v", err)
  172. }
  173. return
  174. }
  175. func (c *callback) Before(name string) *callback {
  176. c.before = name
  177. return c
  178. }
  179. func (c *callback) After(name string) *callback {
  180. c.after = name
  181. return c
  182. }
  183. func (c *callback) Register(name string, fn func(*DB)) error {
  184. c.name = name
  185. c.handler = fn
  186. c.processor.callbacks = append(c.processor.callbacks, c)
  187. return c.processor.compile()
  188. }
  189. func (c *callback) Remove(name string) error {
  190. c.processor.db.Logger.Warn(context.Background(), "removing callback `%s` from %s\n", name, utils.FileWithLineNum())
  191. c.name = name
  192. c.remove = true
  193. c.processor.callbacks = append(c.processor.callbacks, c)
  194. return c.processor.compile()
  195. }
  196. func (c *callback) Replace(name string, fn func(*DB)) error {
  197. c.processor.db.Logger.Info(context.Background(), "replacing callback `%s` from %s\n", name, utils.FileWithLineNum())
  198. c.name = name
  199. c.handler = fn
  200. c.replace = true
  201. c.processor.callbacks = append(c.processor.callbacks, c)
  202. return c.processor.compile()
  203. }
  204. // getRIndex get right index from string slice
  205. func getRIndex(strs []string, str string) int {
  206. for i := len(strs) - 1; i >= 0; i-- {
  207. if strs[i] == str {
  208. return i
  209. }
  210. }
  211. return -1
  212. }
  213. func sortCallbacks(cs []*callback) (fns []func(*DB), err error) {
  214. var (
  215. names, sorted []string
  216. sortCallback func(*callback) error
  217. )
  218. sort.SliceStable(cs, func(i, j int) bool {
  219. if cs[j].before == "*" && cs[i].before != "*" {
  220. return true
  221. }
  222. if cs[j].after == "*" && cs[i].after != "*" {
  223. return true
  224. }
  225. return false
  226. })
  227. for _, c := range cs {
  228. // show warning message the callback name already exists
  229. if idx := getRIndex(names, c.name); idx > -1 && !c.replace && !c.remove && !cs[idx].remove {
  230. c.processor.db.Logger.Warn(context.Background(), "duplicated callback `%s` from %s\n", c.name, utils.FileWithLineNum())
  231. }
  232. names = append(names, c.name)
  233. }
  234. sortCallback = func(c *callback) error {
  235. if c.before != "" { // if defined before callback
  236. if c.before == "*" && len(sorted) > 0 {
  237. if curIdx := getRIndex(sorted, c.name); curIdx == -1 {
  238. sorted = append([]string{c.name}, sorted...)
  239. }
  240. } else if sortedIdx := getRIndex(sorted, c.before); sortedIdx != -1 {
  241. if curIdx := getRIndex(sorted, c.name); curIdx == -1 {
  242. // if before callback already sorted, append current callback just after it
  243. sorted = append(sorted[:sortedIdx], append([]string{c.name}, sorted[sortedIdx:]...)...)
  244. } else if curIdx > sortedIdx {
  245. return fmt.Errorf("conflicting callback %s with before %s", c.name, c.before)
  246. }
  247. } else if idx := getRIndex(names, c.before); idx != -1 {
  248. // if before callback exists
  249. cs[idx].after = c.name
  250. }
  251. }
  252. if c.after != "" { // if defined after callback
  253. if c.after == "*" && len(sorted) > 0 {
  254. if curIdx := getRIndex(sorted, c.name); curIdx == -1 {
  255. sorted = append(sorted, c.name)
  256. }
  257. } else if sortedIdx := getRIndex(sorted, c.after); sortedIdx != -1 {
  258. if curIdx := getRIndex(sorted, c.name); curIdx == -1 {
  259. // if after callback sorted, append current callback to last
  260. sorted = append(sorted, c.name)
  261. } else if curIdx < sortedIdx {
  262. return fmt.Errorf("conflicting callback %s with before %s", c.name, c.after)
  263. }
  264. } else if idx := getRIndex(names, c.after); idx != -1 {
  265. // if after callback exists but haven't sorted
  266. // set after callback's before callback to current callback
  267. after := cs[idx]
  268. if after.before == "" {
  269. after.before = c.name
  270. }
  271. if err := sortCallback(after); err != nil {
  272. return err
  273. }
  274. if err := sortCallback(c); err != nil {
  275. return err
  276. }
  277. }
  278. }
  279. // if current callback haven't been sorted, append it to last
  280. if getRIndex(sorted, c.name) == -1 {
  281. sorted = append(sorted, c.name)
  282. }
  283. return nil
  284. }
  285. for _, c := range cs {
  286. if err = sortCallback(c); err != nil {
  287. return
  288. }
  289. }
  290. for _, name := range sorted {
  291. if idx := getRIndex(names, name); !cs[idx].remove {
  292. fns = append(fns, cs[idx].handler)
  293. }
  294. }
  295. return
  296. }
  297. func removeCallbacks(cs []*callback, nameMap map[string]bool) []*callback {
  298. callbacks := make([]*callback, 0, len(cs))
  299. for _, callback := range cs {
  300. if nameMap[callback.name] {
  301. continue
  302. }
  303. callbacks = append(callbacks, callback)
  304. }
  305. return callbacks
  306. }