forked from fabric8-services/fabric8-wit
/
db_clean.go
84 lines (77 loc) · 2.44 KB
/
db_clean.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
73
74
75
76
77
78
79
80
81
82
83
84
package cleaner
import (
"database/sql"
"github.com/fabric8-services/fabric8-wit/log"
"github.com/fabric8-services/fabric8-wit/workitem"
uuid "github.com/satori/go.uuid"
"fmt"
"github.com/jinzhu/gorm"
)
// DeleteCreatedEntities records all created entities on the gorm.DB connection
// and returns a function which can be called on defer to delete created
// entities in reverse order on function exit.
//
// In addition to that, the WIT cache is cleared as well in order to respect any
// deletions made to the db.
//
// Usage:
//
// func TestDatabaseActions(t *testing.T) {
//
// // setup database connection
// db := ....
// // setup auto clean up of created entities
// defer DeleteCreatedEntities(db)()
//
// repo := NewRepo(db)
// repo.Create(X)
// repo.Create(X)
// repo.Create(X)
// }
//
// Output:
//
// 2017/01/31 12:08:08 Deleting from x 6d143405-1232-40de-bc73-835b543cd972
// 2017/01/31 12:08:08 Deleting from x 0685068d-4934-4d9a-bac2-91eebbca9575
// 2017/01/31 12:08:08 Deleting from x 2d20944e-7952-40c1-bd15-f3fa1a70026d
func DeleteCreatedEntities(db *gorm.DB) func() {
hookName := "mighti:record"
type entity struct {
table string
keyname string
key interface{}
}
var entires []entity
hookRegistered := db.Callback().Create().Get(hookName) != nil
if hookRegistered {
hookName += "-" + uuid.NewV4().String()
}
db.Callback().Create().After("gorm:create").Register(hookName, func(scope *gorm.Scope) {
log.Logger().Debugln(fmt.Sprintf("Inserted entities from %s with %s=%v", scope.TableName(), scope.PrimaryKey(), scope.PrimaryKeyValue()))
entires = append(entires, entity{table: scope.TableName(), keyname: scope.PrimaryKey(), key: scope.PrimaryKeyValue()})
})
return func() {
defer db.Callback().Create().Remove(hookName)
// Find out if the current db object is already a transaction
_, inTransaction := db.CommonDB().(*sql.Tx)
tx := db
if !inTransaction {
tx = db.Begin()
}
for i := len(entires) - 1; i >= 0; i-- {
entry := entires[i]
log.Debug(nil, map[string]interface{}{
"table": entry.table,
"key": entry.key,
"hook_name": hookName,
}, "Deleting entities from '%s' table with key %v", entry.table, entry.key)
tx.Table(entry.table).Where(entry.keyname+" = ?", entry.key).Delete("")
}
// Delete the work item cache as well
// NOTE: Feel free to add more cache freeing calls here as needed.
workitem.ClearGlobalWorkItemTypeCache()
if !inTransaction {
tx.Commit()
}
}
}