hook.go

package pg

import (
	"context"
	"fmt"

	"github.com/go-pg/pg/orm"
)

type dummyFormatter struct{}

func (dummyFormatter) FormatQuery(b []byte, query string, params ...interface{}) []byte {
	return append(b, query...)
}

type QueryEvent struct {
	Ctx     context.Context
	DB      orm.DB
	Query   interface{}
	Params  []interface{}
	Attempt int
	Result  Result
	Error   error

	Data map[interface{}]interface{}
}

type QueryHook interface {
	BeforeQuery(*QueryEvent)
	AfterQuery(*QueryEvent)
}

func (ev *QueryEvent) UnformattedQuery() (string, error) {
	b, err := queryString(ev.Query)
	if err != nil {
		return "", err
	}
	return string(b), nil
}

func (ev *QueryEvent) FormattedQuery() (string, error) {
	b, err := appendQuery(nil, ev.DB, ev.Query, ev.Params...)
	if err != nil {
		return "", err
	}
	return string(b), nil
}

func queryString(query interface{}) ([]byte, error) {
	switch query := query.(type) {
	case orm.TemplateAppender:
		return query.AppendTemplate(nil)
	case string:
		return dummyFormatter{}.FormatQuery(nil, query), nil
	default:
		return nil, fmt.Errorf("pg: can't append %T", query)
	}
}

// AddQueryHook adds a hook into query processing.
func (db *baseDB) AddQueryHook(hook QueryHook) {
	db.queryHooks = append(db.queryHooks, hook)
}

func (db *baseDB) queryStarted(
	c context.Context,
	ormDB orm.DB,
	query interface{},
	params []interface{},
	attempt int,
) *QueryEvent {
	if len(db.queryHooks) == 0 {
		return nil
	}

	event := &QueryEvent{
		Ctx:     c,
		DB:      ormDB,
		Query:   query,
		Params:  params,
		Attempt: attempt,
		Data:    make(map[interface{}]interface{}),
	}
	for _, hook := range db.queryHooks {
		hook.BeforeQuery(event)
	}
	return event
}

func (db *baseDB) queryProcessed(
	res Result,
	err error,
	event *QueryEvent,
) {
	if event == nil {
		return
	}

	event.Error = err
	event.Result = res
	for _, hook := range db.queryHooks {
		hook.AfterQuery(event)
	}
}

func copyQueryHooks(s []QueryHook) []QueryHook {
	return s[:len(s):len(s)]
}