package pg_test
import (
"context"
"time"
. "github.com/onsi/ginkgo"
. "github.com/onsi/gomega"
"github.com/go-pg/pg"
"github.com/go-pg/pg/orm"
)
type HookTest struct {
Id int
Value string
afterQuery int
afterSelect int
beforeInsert int
afterInsert int
beforeUpdate int
afterUpdate int
beforeDelete int
afterDelete int
}
func (t *HookTest) AfterQuery(c context.Context, db orm.DB) error {
t.afterQuery++
return nil
}
func (t *HookTest) AfterSelect(c context.Context, db orm.DB) error {
t.afterSelect++
return nil
}
func (t *HookTest) BeforeInsert(c context.Context, db orm.DB) error {
t.beforeInsert++
return nil
}
func (t *HookTest) AfterInsert(c context.Context, db orm.DB) error {
t.afterInsert++
return nil
}
func (t *HookTest) BeforeUpdate(c context.Context, db orm.DB) error {
t.beforeUpdate++
return nil
}
func (t *HookTest) AfterUpdate(c context.Context, db orm.DB) error {
t.afterUpdate++
return nil
}
func (t *HookTest) BeforeDelete(c context.Context, db orm.DB) error {
t.beforeDelete++
return nil
}
func (t *HookTest) AfterDelete(c context.Context, db orm.DB) error {
t.afterDelete++
return nil
}
type eventHookTest struct {
beforeQueryMethod func(*pg.QueryEvent)
afterQueryMethod func(*pg.QueryEvent)
}
func (e eventHookTest) BeforeQuery(event *pg.QueryEvent) {
e.beforeQueryMethod(event)
}
func (e eventHookTest) AfterQuery(event *pg.QueryEvent) {
e.afterQueryMethod(event)
}
var _ = Describe("HookTest", func() {
var db *pg.DB
BeforeEach(func() {
db = pg.Connect(pgOptions())
qs := []string{
"CREATE TEMP TABLE hook_tests (id int, value text)",
"INSERT INTO hook_tests VALUES (1, '')",
}
for _, q := range qs {
_, err := db.Exec(q)
Expect(err).NotTo(HaveOccurred())
}
})
AfterEach(func() {
Expect(db.Close()).NotTo(HaveOccurred())
})
It("calls AfterQuery for a struct", func() {
var hook HookTest
_, err := db.QueryOne(&hook, "SELECT 1 AS id")
Expect(err).NotTo(HaveOccurred())
Expect(hook.afterQuery).To(Equal(1))
Expect(hook.afterSelect).To(Equal(0))
})
It("calls AfterQuery and AfterSelect for a struct model", func() {
var hook HookTest
err := db.Model(&hook).Select()
Expect(err).NotTo(HaveOccurred())
Expect(hook.afterQuery).To(Equal(1))
Expect(hook.afterSelect).To(Equal(1))
})
It("calls AfterQuery for a slice", func() {
var hooks []HookTest
_, err := db.Query(&hooks, "SELECT 1 AS id")
Expect(err).NotTo(HaveOccurred())
Expect(hooks).To(HaveLen(1))
Expect(hooks[0].afterQuery).To(Equal(1))
Expect(hooks[0].afterSelect).To(Equal(0))
})
It("calls AfterQuery and AfterSelect for a slice model", func() {
var hooks []HookTest
err := db.Model(&hooks).Select()
Expect(err).NotTo(HaveOccurred())
Expect(hooks).To(HaveLen(1))
Expect(hooks[0].afterQuery).To(Equal(1))
Expect(hooks[0].afterSelect).To(Equal(1))
})
It("calls BeforeInsert and AfterInsert", func() {
hook := HookTest{
Id: 1,
Value: "value",
}
err := db.Insert(&hook)
Expect(err).NotTo(HaveOccurred())
Expect(hook.afterQuery).To(Equal(0))
Expect(hook.beforeInsert).To(Equal(1))
Expect(hook.afterInsert).To(Equal(1))
})
It("calls BeforeUpdate and AfterUpdate", func() {
hook := HookTest{
Id: 1,
}
err := db.Update(&hook)
Expect(err).NotTo(HaveOccurred())
Expect(hook.afterQuery).To(Equal(0))
Expect(hook.beforeUpdate).To(Equal(1))
Expect(hook.afterUpdate).To(Equal(1))
})
It("does not call BeforeUpdate and AfterUpdate for nil model", func() {
_, err := db.Model((*HookTest)(nil)).
Set("value = 'new'").
Where("id = 123").
Update()
Expect(err).NotTo(HaveOccurred())
})
It("calls BeforeDelete and AfterDelete", func() {
hook := HookTest{
Id: 1,
}
err := db.Delete(&hook)
Expect(err).NotTo(HaveOccurred())
Expect(hook.afterQuery).To(Equal(0))
Expect(hook.beforeDelete).To(Equal(1))
Expect(hook.afterDelete).To(Equal(1))
})
It("does not call BeforeDelete and AfterDelete for nil model", func() {
_, err := db.Model((*HookTest)(nil)).
Where("id = 123").
Delete()
Expect(err).NotTo(HaveOccurred())
})
})
var _ = Describe("OnQueryEvent", func() {
var db *pg.DB
var count int
BeforeEach(func() {
db = pg.Connect(pgOptions())
count = 0
})
AfterEach(func() {
Expect(db.Close()).NotTo(HaveOccurred())
})
Describe("Query/Exec", func() {
afterQuery := func(event *pg.QueryEvent) {
Expect(event.DB).To(Equal(db))
Expect(event.Query).To(Equal("SELECT ?"))
Expect(event.Params).To(Equal([]interface{}{1}))
Expect(event.Result).NotTo(BeNil())
Expect(event.Error).NotTo(HaveOccurred())
q, err := event.UnformattedQuery()
Expect(err).NotTo(HaveOccurred())
Expect(q).To(Equal("SELECT ?"))
q, err = event.FormattedQuery()
Expect(err).NotTo(HaveOccurred())
Expect(q).To(Equal("SELECT 1"))
Expect(event.Data["data"]).To(Equal(1))
count++
}
beforeQuery := func(event *pg.QueryEvent) {
Expect(event.DB).To(Equal(db))
Expect(event.Query).To(Equal("SELECT ?"))
Expect(event.Params).To(Equal([]interface{}{1}))
Expect(event.Result).To(BeNil())
Expect(event.Error).To(BeNil())
q, err := event.UnformattedQuery()
Expect(err).NotTo(HaveOccurred())
Expect(q).To(Equal("SELECT ?"))
q, err = event.FormattedQuery()
Expect(err).NotTo(HaveOccurred())
Expect(q).To(Equal("SELECT 1"))
event.Data["data"] = 1
}
BeforeEach(func() {
hookImpl := struct{ eventHookTest }{}
hookImpl.afterQueryMethod = afterQuery
hookImpl.beforeQueryMethod = beforeQuery
db.AddQueryHook(hookImpl)
})
It("is called for Query", func() {
_, err := db.Query(pg.Discard, "SELECT ?", 1)
Expect(err).NotTo(HaveOccurred())
Expect(count).To(Equal(1))
})
It("is called for Exec", func() {
_, err := db.Exec("SELECT ?", 1)
Expect(err).NotTo(HaveOccurred())
Expect(count).To(Equal(1))
})
})
Describe("Model", func() {
afterQuery := func(event *pg.QueryEvent) {
Expect(event.DB).To(Equal(db))
Expect(event.Query).NotTo(BeNil())
Expect(event.Params).To(HaveLen(1))
Expect(event.Result).NotTo(BeNil())
Expect(event.Error).NotTo(HaveOccurred())
q, err := event.UnformattedQuery()
Expect(err).NotTo(HaveOccurred())
Expect(q).To(Equal("SELECT ?"))
q, err = event.FormattedQuery()
Expect(err).NotTo(HaveOccurred())
Expect(q).To(Equal("SELECT 1"))
Expect(event.Data["data"]).To(Equal(1))
count++
}
beforeQuery := func(event *pg.QueryEvent) {
Expect(event.DB).To(Equal(db))
Expect(event.Query).NotTo(BeNil())
Expect(event.Params).To(HaveLen(1))
Expect(event.Result).To(BeNil())
Expect(event.Error).To(BeNil())
q, err := event.UnformattedQuery()
Expect(err).NotTo(HaveOccurred())
Expect(q).To(Equal("SELECT ?"))
q, err = event.FormattedQuery()
Expect(err).NotTo(HaveOccurred())
Expect(q).To(Equal("SELECT 1"))
event.Data["data"] = 1
}
BeforeEach(func() {
hookImpl := struct{ eventHookTest }{}
hookImpl.afterQueryMethod = afterQuery
hookImpl.beforeQueryMethod = beforeQuery
db.AddQueryHook(hookImpl)
})
It("is called for Model", func() {
var n int
err := db.Model().ColumnExpr("?", 1).Select(&n)
Expect(err).NotTo(HaveOccurred())
Expect(count).To(Equal(1))
})
})
})
type BeforeSelectQueryModel struct {
Id int
DeletedAt time.Time
}
func (BeforeSelectQueryModel) BeforeSelectQuery(
c context.Context, db orm.DB, q *orm.Query,
) (*orm.Query, error) {
q = q.Where("?TableAlias.deleted_at IS NULL")
return q, nil
}
var _ = Describe("BeforeSelectQueryModel", func() {
var db *pg.DB
BeforeEach(func() {
db = pg.Connect(pgOptions())
err := db.CreateTable((*BeforeSelectQueryModel)(nil), &orm.CreateTableOptions{
Temp: true,
})
Expect(err).NotTo(HaveOccurred())
models := []BeforeSelectQueryModel{
{Id: 1},
{Id: 2, DeletedAt: time.Now()},
}
_, err = db.Model(&models).Insert()
Expect(err).NotTo(HaveOccurred())
})
AfterEach(func() {
Expect(db.Close()).NotTo(HaveOccurred())
})
It("applies BeforeSelectQuery hook", func() {
var models []BeforeSelectQueryModel
err := db.Model(&models).Select()
Expect(err).NotTo(HaveOccurred())
Expect(models).To(Equal([]BeforeSelectQueryModel{
{Id: 1},
}))
})
})