/
db.go
72 lines (57 loc) · 1.7 KB
/
db.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
package testingcommon
import (
"fmt"
"github.com/TheLazarusNetwork/go-helpers/logo"
"github.com/TheLazarusNetwork/superiad/pkg/store"
"gorm.io/gorm"
)
// Referred from https://medium.com/@jarifibrahim/using-gorm-hooks-to-clean-up-test-fixtures-in-golang-99b0fcb04354
func DeleteCreatedEntities() func() {
db := store.DB
type entity struct {
table string
keyname string
key interface{}
}
var entries []entity
hookName := "cleanupHook"
db.Callback().Create().After("gorm:create").Register(hookName, func(tx *gorm.DB) {
statement := tx.Statement.Statement
schema := tx.Statement.Schema
if len(schema.PrimaryFields) == 0 {
return
}
primaryFieldKey := ""
primaryFieldIndex := 0
for i, v := range schema.Fields {
if v.PrimaryKey {
primaryFieldKey = v.DBName
primaryFieldIndex = i
break
}
}
primaryFieldValue := statement.Vars[primaryFieldIndex]
fmt.Printf("Inserted entities of %s with %s=%v\n", statement.Table, primaryFieldKey, primaryFieldValue)
entries = append(entries, entity{table: statement.Table, keyname: primaryFieldKey, key: primaryFieldValue})
})
return func() {
// Remove the hook once we're done
defer db.Callback().Create().Remove(hookName)
for i := len(entries) - 1; i >= 0; i-- {
entry := entries[i]
var deleteValue string
switch v := entry.key.(type) {
case int:
deleteValue = fmt.Sprint(v)
case string:
deleteValue = fmt.Sprintf("'%v'", v)
default:
logo.Fatal("not implemented")
}
q := fmt.Sprintf(`DELETE FROM %v WHERE %v=%v`, entry.table, entry.keyname, deleteValue)
if err := db.Exec(q).Error; err != nil {
logo.Warnf("failed to exec query: %s, in clean up hook, %s", q, err)
}
}
}
}