githaven-fork/vendor/xorm.io/xorm/contexts/hook.go

// Copyright 2020 The Xorm Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package contexts

import (
	"context"
	"database/sql"
	"time"
)

// ContextHook represents a hook context
type ContextHook struct {
	start       time.Time
	Ctx         context.Context
	SQL         string        // log content or SQL
	Args        []interface{} // if it's a SQL, it's the arguments
	Result      sql.Result
	ExecuteTime time.Duration
	Err         error // SQL executed error
}

// NewContextHook return context for hook
func NewContextHook(ctx context.Context, sql string, args []interface{}) *ContextHook {
	return &ContextHook{
		start: time.Now(),
		Ctx:   ctx,
		SQL:   sql,
		Args:  args,
	}
}

func (c *ContextHook) End(ctx context.Context, result sql.Result, err error) {
	c.Ctx = ctx
	c.Result = result
	c.Err = err
	c.ExecuteTime = time.Now().Sub(c.start)
}

type Hook interface {
	BeforeProcess(c *ContextHook) (context.Context, error)
	AfterProcess(c *ContextHook) error
}

type Hooks struct {
	hooks []Hook
}

func (h *Hooks) AddHook(hooks ...Hook) {
	h.hooks = append(h.hooks, hooks...)
}

func (h *Hooks) BeforeProcess(c *ContextHook) (context.Context, error) {
	ctx := c.Ctx
	for _, h := range h.hooks {
		var err error
		ctx, err = h.BeforeProcess(c)
		if err != nil {
			return nil, err
		}
	}
	return ctx, nil
}

func (h *Hooks) AfterProcess(c *ContextHook) error {
	firstErr := c.Err
	for _, h := range h.hooks {
		err := h.AfterProcess(c)
		if err != nil && firstErr == nil {
			firstErr = err
		}
	}
	return firstErr
}