/
database_utils.go
92 lines (82 loc) · 2.84 KB
/
database_utils.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
85
86
87
88
89
90
91
92
// Copyright (C) Simfiny, Inc. 2022-present.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
package postgres // import "github.com/SolomonAIEngineering/backend-core-library/database/postgres"
import (
"database/sql"
"fmt"
"math/rand"
"time"
"unsafe"
"github.com/jinzhu/gorm"
)
var src = rand.NewSource(time.Now().UnixNano())
const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
const (
letterIdxBits = 6 // 6 bits to represent a letter index
letterIdxMask = 1<<letterIdxBits - 1 // All 1-bits, as many as letterIdxBits
letterIdxMax = 63 / letterIdxBits // # of letter indices fitting in 63 bits
)
// GenerateRandomString generates a random string based on the size specified by the client
func GenerateRandomString(n int) string {
b := make([]byte, n)
// A src.Int63() generates 63 random bits, enough for letterIdxMax characters!
for i, cache, remain := n-1, src.Int63(), letterIdxMax; i >= 0; {
if remain == 0 {
cache, remain = src.Int63(), letterIdxMax
}
if idx := int(cache & letterIdxMask); idx < len(letterBytes) {
b[i] = letterBytes[idx]
i--
}
cache >>= letterIdxBits
remain--
}
return *(*string)(unsafe.Pointer(&b))
}
// DeleteCreatedEntitiesAfterTest sets up GORM `onCreate` hook and return a function that can be deferred to
// remove all the entities created after the hook was set up
// You can use it as
//
// func TestSomething(t *testing.T){
// db, _ := gorm.Open(...)
//
// cleaner := DeleteCreatedEntities(db)
// defer cleaner()
//
// }
func DeleteCreatedEntitiesAfterTest(db *gorm.DB) func() {
type entity struct {
table string
keyname string
key interface{}
}
var entries []entity
hookName := "cleanupHook"
db.Callback().Create().After("gorm:create").Register(hookName, func(scope *gorm.Scope) {
fmt.Printf("Inserted entities of %s with %s=%v\n", scope.TableName(), scope.PrimaryKey(), scope.PrimaryKeyValue())
entries = append(entries, entity{table: scope.TableName(), keyname: scope.PrimaryKey(), key: scope.PrimaryKeyValue()})
})
return func() {
// Remove the hook once we're done
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()
}
// Loop from the end. It is important that we delete the entries in the
// reverse order of their insertion
for i := len(entries) - 1; i >= 0; i-- {
entry := entries[i]
fmt.Printf("Deleting entities from '%s' table with key %v\n", entry.table, entry.key)
tx.Table(entry.table).Where(entry.keyname+" = ?", entry.key).Delete("")
}
if !inTransaction {
tx.Commit()
}
}
}