mirror of
https://github.com/gophish/gophish
synced 2024-11-14 00:07:19 +00:00
Performance Improvements for Campaign and Group Creation (#1686)
This commit significantly improves the performance of campaign and group creation by changing database access to use transactions. It should also make things more consistent with campaign creation. Specifically, this will ensure that the entire campaign gets created before emails start sending, while I anticipate this will fix #1643, #1080, (possibly) #1677, and #1552.
This commit is contained in:
parent
c2f579a2c5
commit
44f88401bb
7 changed files with 383 additions and 60 deletions
|
@ -491,6 +491,7 @@ func PostCampaign(c *Campaign, uid int64) error {
|
|||
// Insert all the results
|
||||
resultMap := make(map[string]bool)
|
||||
recipientIndex := 0
|
||||
tx := db.Begin()
|
||||
for _, g := range c.Groups {
|
||||
// Insert a result for each target in the group
|
||||
for _, t := range g.Targets {
|
||||
|
@ -515,24 +516,30 @@ func PostCampaign(c *Campaign, uid int64) error {
|
|||
Reported: false,
|
||||
ModifiedDate: c.CreatedDate,
|
||||
}
|
||||
err = r.GenerateId()
|
||||
err = r.GenerateId(tx)
|
||||
if err != nil {
|
||||
log.Error(err)
|
||||
continue
|
||||
tx.Rollback()
|
||||
return err
|
||||
}
|
||||
processing := false
|
||||
if r.SendDate.Before(c.CreatedDate) || r.SendDate.Equal(c.CreatedDate) {
|
||||
r.Status = StatusSending
|
||||
processing = true
|
||||
}
|
||||
err = db.Save(r).Error
|
||||
err = tx.Save(r).Error
|
||||
if err != nil {
|
||||
log.WithFields(logrus.Fields{
|
||||
"email": t.Email,
|
||||
}).Error(err)
|
||||
}).Errorf("error creating result: %v", err)
|
||||
tx.Rollback()
|
||||
return err
|
||||
}
|
||||
c.Results = append(c.Results, *r)
|
||||
log.Infof("Creating maillog for %s to send at %s\n", r.Email, sendDate)
|
||||
log.WithFields(logrus.Fields{
|
||||
"email": r.Email,
|
||||
"send_date": sendDate,
|
||||
}).Debug("creating maillog")
|
||||
m := &MailLog{
|
||||
UserId: c.UserId,
|
||||
CampaignId: c.Id,
|
||||
|
@ -540,16 +547,18 @@ func PostCampaign(c *Campaign, uid int64) error {
|
|||
SendDate: sendDate,
|
||||
Processing: processing,
|
||||
}
|
||||
err = db.Save(m).Error
|
||||
err = tx.Save(m).Error
|
||||
if err != nil {
|
||||
log.Error(err)
|
||||
continue
|
||||
log.WithFields(logrus.Fields{
|
||||
"email": t.Email,
|
||||
}).Errorf("error creating maillog entry: %v", err)
|
||||
tx.Rollback()
|
||||
return err
|
||||
}
|
||||
recipientIndex++
|
||||
}
|
||||
}
|
||||
err = db.Save(c).Error
|
||||
return err
|
||||
return tx.Commit().Error
|
||||
}
|
||||
|
||||
//DeleteCampaign deletes the specified campaign
|
||||
|
|
|
@ -1,6 +1,8 @@
|
|||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
check "gopkg.in/check.v1"
|
||||
|
@ -14,6 +16,11 @@ func (s *ModelsSuite) TestGenerateSendDate(c *check.C) {
|
|||
c.Assert(err, check.Equals, nil)
|
||||
c.Assert(campaign.LaunchDate, check.Equals, campaign.CreatedDate)
|
||||
|
||||
// For comparing the dates, we need to fetch the campaign again. This is
|
||||
// to solve an issue where the campaign object right now has time down to
|
||||
// the microsecond, while in MySQL it's rounded down to the second.
|
||||
campaign, _ = GetCampaign(campaign.Id, campaign.UserId)
|
||||
|
||||
ms, err := GetMailLogsByCampaign(campaign.Id)
|
||||
c.Assert(err, check.Equals, nil)
|
||||
for _, m := range ms {
|
||||
|
@ -27,6 +34,8 @@ func (s *ModelsSuite) TestGenerateSendDate(c *check.C) {
|
|||
err = PostCampaign(&campaign, campaign.UserId)
|
||||
c.Assert(err, check.Equals, nil)
|
||||
|
||||
campaign, _ = GetCampaign(campaign.Id, campaign.UserId)
|
||||
|
||||
ms, err = GetMailLogsByCampaign(campaign.Id)
|
||||
c.Assert(err, check.Equals, nil)
|
||||
for _, m := range ms {
|
||||
|
@ -41,6 +50,8 @@ func (s *ModelsSuite) TestGenerateSendDate(c *check.C) {
|
|||
err = PostCampaign(&campaign, campaign.UserId)
|
||||
c.Assert(err, check.Equals, nil)
|
||||
|
||||
campaign, _ = GetCampaign(campaign.Id, campaign.UserId)
|
||||
|
||||
ms, err = GetMailLogsByCampaign(campaign.Id)
|
||||
c.Assert(err, check.Equals, nil)
|
||||
sendingOffset := 2 / float64(len(ms))
|
||||
|
@ -133,3 +144,128 @@ func (s *ModelsSuite) TestCompleteCampaignAlsoDeletesMailLogs(c *check.C) {
|
|||
c.Assert(err, check.Equals, nil)
|
||||
c.Assert(len(ms), check.Equals, 0)
|
||||
}
|
||||
|
||||
func (s *ModelsSuite) TestCampaignGetResults(c *check.C) {
|
||||
campaign := s.createCampaign(c)
|
||||
got, err := GetCampaign(campaign.Id, campaign.UserId)
|
||||
c.Assert(err, check.Equals, nil)
|
||||
c.Assert(len(campaign.Results), check.Equals, len(got.Results))
|
||||
}
|
||||
|
||||
func setupCampaignDependencies(b *testing.B, size int) {
|
||||
group := Group{Name: "Test Group"}
|
||||
// Create a large group of 5000 members
|
||||
for i := 0; i < size; i++ {
|
||||
group.Targets = append(group.Targets, Target{BaseRecipient: BaseRecipient{Email: fmt.Sprintf("test%d@example.com", i), FirstName: "User", LastName: fmt.Sprintf("%d", i)}})
|
||||
}
|
||||
group.UserId = 1
|
||||
err := PostGroup(&group)
|
||||
if err != nil {
|
||||
b.Fatalf("error posting group: %v", err)
|
||||
}
|
||||
|
||||
// Add a template
|
||||
template := Template{Name: "Test Template"}
|
||||
template.Subject = "{{.RId}} - Subject"
|
||||
template.Text = "{{.RId}} - Text"
|
||||
template.HTML = "{{.RId}} - HTML"
|
||||
template.UserId = 1
|
||||
err = PostTemplate(&template)
|
||||
if err != nil {
|
||||
b.Fatalf("error posting template: %v", err)
|
||||
}
|
||||
|
||||
// Add a landing page
|
||||
p := Page{Name: "Test Page"}
|
||||
p.HTML = "<html>Test</html>"
|
||||
p.UserId = 1
|
||||
err = PostPage(&p)
|
||||
if err != nil {
|
||||
b.Fatalf("error posting page: %v", err)
|
||||
}
|
||||
|
||||
// Add a sending profile
|
||||
smtp := SMTP{Name: "Test Page"}
|
||||
smtp.UserId = 1
|
||||
smtp.Host = "example.com"
|
||||
smtp.FromAddress = "test@test.com"
|
||||
err = PostSMTP(&smtp)
|
||||
if err != nil {
|
||||
b.Fatalf("error posting smtp: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkCampaign100(b *testing.B) {
|
||||
setupBenchmark(b)
|
||||
setupCampaignDependencies(b, 100)
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
campaign := Campaign{Name: "Test campaign"}
|
||||
campaign.UserId = 1
|
||||
campaign.Template = Template{Name: "Test Template"}
|
||||
campaign.Page = Page{Name: "Test Page"}
|
||||
campaign.SMTP = SMTP{Name: "Test Page"}
|
||||
campaign.Groups = []Group{Group{Name: "Test Group"}}
|
||||
|
||||
b.StartTimer()
|
||||
err := PostCampaign(&campaign, 1)
|
||||
if err != nil {
|
||||
b.Fatalf("error posting campaign: %v", err)
|
||||
}
|
||||
b.StopTimer()
|
||||
db.Delete(Result{})
|
||||
db.Delete(MailLog{})
|
||||
db.Delete(Campaign{})
|
||||
}
|
||||
tearDownBenchmark(b)
|
||||
}
|
||||
|
||||
func BenchmarkCampaign1000(b *testing.B) {
|
||||
setupBenchmark(b)
|
||||
setupCampaignDependencies(b, 1000)
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
campaign := Campaign{Name: "Test campaign"}
|
||||
campaign.UserId = 1
|
||||
campaign.Template = Template{Name: "Test Template"}
|
||||
campaign.Page = Page{Name: "Test Page"}
|
||||
campaign.SMTP = SMTP{Name: "Test Page"}
|
||||
campaign.Groups = []Group{Group{Name: "Test Group"}}
|
||||
|
||||
b.StartTimer()
|
||||
err := PostCampaign(&campaign, 1)
|
||||
if err != nil {
|
||||
b.Fatalf("error posting campaign: %v", err)
|
||||
}
|
||||
b.StopTimer()
|
||||
db.Delete(Result{})
|
||||
db.Delete(MailLog{})
|
||||
db.Delete(Campaign{})
|
||||
}
|
||||
tearDownBenchmark(b)
|
||||
}
|
||||
|
||||
func BenchmarkCampaign10000(b *testing.B) {
|
||||
setupBenchmark(b)
|
||||
setupCampaignDependencies(b, 10000)
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
campaign := Campaign{Name: "Test campaign"}
|
||||
campaign.UserId = 1
|
||||
campaign.Template = Template{Name: "Test Template"}
|
||||
campaign.Page = Page{Name: "Test Page"}
|
||||
campaign.SMTP = SMTP{Name: "Test Page"}
|
||||
campaign.Groups = []Group{Group{Name: "Test Group"}}
|
||||
|
||||
b.StartTimer()
|
||||
err := PostCampaign(&campaign, 1)
|
||||
if err != nil {
|
||||
b.Fatalf("error posting campaign: %v", err)
|
||||
}
|
||||
b.StopTimer()
|
||||
db.Delete(Result{})
|
||||
db.Delete(MailLog{})
|
||||
db.Delete(Campaign{})
|
||||
}
|
||||
tearDownBenchmark(b)
|
||||
}
|
||||
|
|
112
models/group.go
112
models/group.go
|
@ -196,13 +196,26 @@ func PostGroup(g *Group) error {
|
|||
return err
|
||||
}
|
||||
// Insert the group into the DB
|
||||
err := db.Save(g).Error
|
||||
tx := db.Begin()
|
||||
err := tx.Save(g).Error
|
||||
if err != nil {
|
||||
tx.Rollback()
|
||||
log.Error(err)
|
||||
return err
|
||||
}
|
||||
for _, t := range g.Targets {
|
||||
insertTargetIntoGroup(t, g.Id)
|
||||
err = insertTargetIntoGroup(tx, t, g.Id)
|
||||
if err != nil {
|
||||
tx.Rollback()
|
||||
log.Error(err)
|
||||
return err
|
||||
}
|
||||
}
|
||||
err = tx.Commit().Error
|
||||
if err != nil {
|
||||
log.Error(err)
|
||||
tx.Rollback()
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
@ -221,50 +234,65 @@ func PutGroup(g *Group) error {
|
|||
}).Error("Error getting targets from group")
|
||||
return err
|
||||
}
|
||||
// Check existing targets, removing any that are no longer in the group.
|
||||
tExists := false
|
||||
// Preload the caches
|
||||
cacheNew := make(map[string]int64, len(g.Targets))
|
||||
for _, t := range g.Targets {
|
||||
cacheNew[t.Email] = t.Id
|
||||
}
|
||||
|
||||
cacheExisting := make(map[string]int64, len(ts))
|
||||
for _, t := range ts {
|
||||
tExists = false
|
||||
// Is the target still in the group?
|
||||
for _, nt := range g.Targets {
|
||||
if t.Email == nt.Email {
|
||||
tExists = true
|
||||
break
|
||||
}
|
||||
cacheExisting[t.Email] = t.Id
|
||||
}
|
||||
|
||||
tx := db.Begin()
|
||||
// Check existing targets, removing any that are no longer in the group.
|
||||
for _, t := range ts {
|
||||
if _, ok := cacheNew[t.Email]; ok {
|
||||
continue
|
||||
}
|
||||
|
||||
// If the target does not exist in the group any longer, we delete it
|
||||
if !tExists {
|
||||
err := db.Where("group_id=? and target_id=?", g.Id, t.Id).Delete(&GroupTarget{}).Error
|
||||
if err != nil {
|
||||
log.WithFields(logrus.Fields{
|
||||
"email": t.Email,
|
||||
}).Error("Error deleting email")
|
||||
}
|
||||
err := tx.Where("group_id=? and target_id=?", g.Id, t.Id).Delete(&GroupTarget{}).Error
|
||||
if err != nil {
|
||||
tx.Rollback()
|
||||
log.WithFields(logrus.Fields{
|
||||
"email": t.Email,
|
||||
}).Error("Error deleting email")
|
||||
}
|
||||
}
|
||||
// Add any targets that are not in the database yet.
|
||||
for _, nt := range g.Targets {
|
||||
// Check and see if the target already exists in the db
|
||||
tExists = false
|
||||
for _, t := range ts {
|
||||
if t.Email == nt.Email {
|
||||
tExists = true
|
||||
nt.Id = t.Id
|
||||
break
|
||||
// If the target already exists in the database, we should just update
|
||||
// the record with the latest information.
|
||||
if id, ok := cacheExisting[nt.Email]; ok {
|
||||
nt.Id = id
|
||||
err = UpdateTarget(tx, nt)
|
||||
if err != nil {
|
||||
log.Error(err)
|
||||
tx.Rollback()
|
||||
return err
|
||||
}
|
||||
continue
|
||||
}
|
||||
// Add target if not in database, otherwise update target information.
|
||||
if !tExists {
|
||||
insertTargetIntoGroup(nt, g.Id)
|
||||
} else {
|
||||
UpdateTarget(nt)
|
||||
// Otherwise, add target if not in database
|
||||
err = insertTargetIntoGroup(tx, nt, g.Id)
|
||||
if err != nil {
|
||||
log.Error(err)
|
||||
tx.Rollback()
|
||||
return err
|
||||
}
|
||||
}
|
||||
err = db.Save(g).Error
|
||||
err = tx.Save(g).Error
|
||||
if err != nil {
|
||||
log.Error(err)
|
||||
return err
|
||||
}
|
||||
err = tx.Commit().Error
|
||||
if err != nil {
|
||||
tx.Rollback()
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
|
@ -285,28 +313,25 @@ func DeleteGroup(g *Group) error {
|
|||
return err
|
||||
}
|
||||
|
||||
func insertTargetIntoGroup(t Target, gid int64) error {
|
||||
func insertTargetIntoGroup(tx *gorm.DB, t Target, gid int64) error {
|
||||
if _, err := mail.ParseAddress(t.Email); err != nil {
|
||||
log.WithFields(logrus.Fields{
|
||||
"email": t.Email,
|
||||
}).Error("Invalid email")
|
||||
return err
|
||||
}
|
||||
trans := db.Begin()
|
||||
err := trans.Where(t).FirstOrCreate(&t).Error
|
||||
err := tx.Where(t).FirstOrCreate(&t).Error
|
||||
if err != nil {
|
||||
log.WithFields(logrus.Fields{
|
||||
"email": t.Email,
|
||||
}).Error(err)
|
||||
trans.Rollback()
|
||||
return err
|
||||
}
|
||||
err = trans.Where("group_id=? and target_id=?", gid, t.Id).Find(&GroupTarget{}).Error
|
||||
err = tx.Where("group_id=? and target_id=?", gid, t.Id).Find(&GroupTarget{}).Error
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
err = trans.Save(&GroupTarget{GroupId: gid, TargetId: t.Id}).Error
|
||||
err = tx.Save(&GroupTarget{GroupId: gid, TargetId: t.Id}).Error
|
||||
if err != nil {
|
||||
log.Error(err)
|
||||
trans.Rollback()
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
@ -314,26 +339,19 @@ func insertTargetIntoGroup(t Target, gid int64) error {
|
|||
log.WithFields(logrus.Fields{
|
||||
"email": t.Email,
|
||||
}).Error("Error adding many-many mapping")
|
||||
trans.Rollback()
|
||||
return err
|
||||
}
|
||||
err = trans.Commit().Error
|
||||
if err != nil {
|
||||
trans.Rollback()
|
||||
log.Error("Error committing db changes")
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateTarget updates the given target information in the database.
|
||||
func UpdateTarget(target Target) error {
|
||||
func UpdateTarget(tx *gorm.DB, target Target) error {
|
||||
targetInfo := map[string]interface{}{
|
||||
"first_name": target.FirstName,
|
||||
"last_name": target.LastName,
|
||||
"position": target.Position,
|
||||
}
|
||||
err := db.Model(&target).Where("id = ?", target.Id).Updates(targetInfo).Error
|
||||
err := tx.Model(&target).Where("id = ?", target.Id).Updates(targetInfo).Error
|
||||
if err != nil {
|
||||
log.WithFields(logrus.Fields{
|
||||
"email": target.Email,
|
||||
|
|
|
@ -1,6 +1,9 @@
|
|||
package models
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/jinzhu/gorm"
|
||||
"gopkg.in/check.v1"
|
||||
)
|
||||
|
@ -170,3 +173,121 @@ func (s *ModelsSuite) TestPutGroupEmptyAttribute(c *check.C) {
|
|||
c.Assert(targets[1].FirstName, check.Equals, "Second")
|
||||
c.Assert(targets[1].LastName, check.Equals, "Example")
|
||||
}
|
||||
|
||||
func benchmarkPostGroup(b *testing.B, iter, size int) {
|
||||
b.StopTimer()
|
||||
g := &Group{
|
||||
Name: fmt.Sprintf("Group-%d", iter),
|
||||
}
|
||||
for i := 0; i < size; i++ {
|
||||
g.Targets = append(g.Targets, Target{
|
||||
BaseRecipient: BaseRecipient{
|
||||
FirstName: "User",
|
||||
LastName: fmt.Sprintf("%d", i),
|
||||
Email: fmt.Sprintf("test-%d@test.com", i),
|
||||
},
|
||||
})
|
||||
}
|
||||
b.StartTimer()
|
||||
err := PostGroup(g)
|
||||
if err != nil {
|
||||
b.Fatalf("error posting group: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// benchmarkPutGroup modifies half of the group to simulate a large change
|
||||
func benchmarkPutGroup(b *testing.B, iter, size int) {
|
||||
b.StopTimer()
|
||||
// First, we need to create the group
|
||||
g := &Group{
|
||||
Name: fmt.Sprintf("Group-%d", iter),
|
||||
}
|
||||
for i := 0; i < size; i++ {
|
||||
g.Targets = append(g.Targets, Target{
|
||||
BaseRecipient: BaseRecipient{
|
||||
FirstName: "User",
|
||||
LastName: fmt.Sprintf("%d", i),
|
||||
Email: fmt.Sprintf("test-%d@test.com", i),
|
||||
},
|
||||
})
|
||||
}
|
||||
err := PostGroup(g)
|
||||
if err != nil {
|
||||
b.Fatalf("error posting group: %v", err)
|
||||
}
|
||||
// Now we need to change half of the group.
|
||||
for i := 0; i < size/2; i++ {
|
||||
g.Targets[i].Email = fmt.Sprintf("test-modified-%d@test.com", i)
|
||||
}
|
||||
b.StartTimer()
|
||||
err = PutGroup(g)
|
||||
if err != nil {
|
||||
b.Fatalf("error modifying group: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkPostGroup100(b *testing.B) {
|
||||
setupBenchmark(b)
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
benchmarkPostGroup(b, i, 100)
|
||||
b.StopTimer()
|
||||
resetBenchmark(b)
|
||||
}
|
||||
tearDownBenchmark(b)
|
||||
}
|
||||
|
||||
func BenchmarkPostGroup1000(b *testing.B) {
|
||||
setupBenchmark(b)
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
benchmarkPostGroup(b, i, 1000)
|
||||
b.StopTimer()
|
||||
resetBenchmark(b)
|
||||
}
|
||||
tearDownBenchmark(b)
|
||||
}
|
||||
|
||||
func BenchmarkPostGroup10000(b *testing.B) {
|
||||
setupBenchmark(b)
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
benchmarkPostGroup(b, i, 10000)
|
||||
b.StopTimer()
|
||||
resetBenchmark(b)
|
||||
}
|
||||
tearDownBenchmark(b)
|
||||
}
|
||||
|
||||
func BenchmarkPutGroup100(b *testing.B) {
|
||||
setupBenchmark(b)
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
benchmarkPutGroup(b, i, 100)
|
||||
b.StopTimer()
|
||||
resetBenchmark(b)
|
||||
}
|
||||
tearDownBenchmark(b)
|
||||
}
|
||||
|
||||
func BenchmarkPutGroup1000(b *testing.B) {
|
||||
setupBenchmark(b)
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
benchmarkPutGroup(b, i, 1000)
|
||||
b.StopTimer()
|
||||
resetBenchmark(b)
|
||||
}
|
||||
tearDownBenchmark(b)
|
||||
}
|
||||
|
||||
func BenchmarkPutGroup10000(b *testing.B) {
|
||||
setupBenchmark(b)
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
benchmarkPutGroup(b, i, 10000)
|
||||
b.StopTimer()
|
||||
resetBenchmark(b)
|
||||
}
|
||||
tearDownBenchmark(b)
|
||||
}
|
||||
|
|
|
@ -96,5 +96,44 @@ func (s *ModelsSuite) createCampaign(ch *check.C) Campaign {
|
|||
c := s.createCampaignDependencies(ch)
|
||||
// Setup and "launch" our campaign
|
||||
ch.Assert(PostCampaign(&c, c.UserId), check.Equals, nil)
|
||||
|
||||
// For comparing the dates, we need to fetch the campaign again. This is
|
||||
// to solve an issue where the campaign object right now has time down to
|
||||
// the microsecond, while in MySQL it's rounded down to the second.
|
||||
c, _ = GetCampaign(c.Id, c.UserId)
|
||||
return c
|
||||
}
|
||||
|
||||
func setupBenchmark(b *testing.B) {
|
||||
conf := &config.Config{
|
||||
DBName: "sqlite3",
|
||||
DBPath: ":memory:",
|
||||
MigrationsPath: "../db/db_sqlite3/migrations/",
|
||||
}
|
||||
err := Setup(conf)
|
||||
if err != nil {
|
||||
b.Fatalf("Failed creating database: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func tearDownBenchmark(b *testing.B) {
|
||||
err := db.Close()
|
||||
if err != nil {
|
||||
b.Fatalf("error closing database: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func resetBenchmark(b *testing.B) {
|
||||
db.Delete(Group{})
|
||||
db.Delete(Target{})
|
||||
db.Delete(GroupTarget{})
|
||||
db.Delete(SMTP{})
|
||||
db.Delete(Page{})
|
||||
db.Delete(Result{})
|
||||
db.Delete(MailLog{})
|
||||
db.Delete(Campaign{})
|
||||
|
||||
// Reset users table to default state.
|
||||
db.Not("id", 1).Delete(User{})
|
||||
db.Model(User{}).Update("username", "admin")
|
||||
}
|
||||
|
|
|
@ -189,7 +189,7 @@ func generateResultId() (string, error) {
|
|||
|
||||
// GenerateId generates a unique key to represent the result
|
||||
// in the database
|
||||
func (r *Result) GenerateId() error {
|
||||
func (r *Result) GenerateId(tx *gorm.DB) error {
|
||||
// Keep trying until we generate a unique key (shouldn't take more than one or two iterations)
|
||||
for {
|
||||
rid, err := generateResultId()
|
||||
|
@ -197,7 +197,7 @@ func (r *Result) GenerateId() error {
|
|||
return err
|
||||
}
|
||||
r.RId = rid
|
||||
err = db.Table("results").Where("r_id=?", r.RId).First(&Result{}).Error
|
||||
err = tx.Table("results").Where("r_id=?", r.RId).First(&Result{}).Error
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
break
|
||||
}
|
||||
|
|
|
@ -10,7 +10,7 @@ import (
|
|||
|
||||
func (s *ModelsSuite) TestGenerateResultId(c *check.C) {
|
||||
r := Result{}
|
||||
r.GenerateId()
|
||||
r.GenerateId(db)
|
||||
match, err := regexp.Match("[a-zA-Z0-9]{7}", []byte(r.RId))
|
||||
c.Assert(err, check.Equals, nil)
|
||||
c.Assert(match, check.Equals, true)
|
||||
|
|
Loading…
Reference in a new issue