add hooks support

go.mod

@@ -22,6 +22,7 @@ require (
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/qustavo/sqlhooks/v2 v2.1.0 // indirect
 	github.com/sirupsen/logrus v1.8.1 // indirect
 	golang.org/x/sys v0.6.0 // indirect
 	gopkg.in/yaml.v3 v3.0.0 // indirect

go.sum

@@ -27,6 +27,7 @@ github.com/fatih/color v1.12.0/go.mod h1:ELkj/draVOlAH/xkhN6mQ50Qd0MPOk5AAr3maGE
 github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w=
 github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
 github.com/ghodss/yaml v1.0.0/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
+github.com/go-sql-driver/mysql v1.4.1/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG5ZlKdlhCg5w=
 github.com/go-sql-driver/mysql v1.5.0 h1:ozyZYNQW3x3HtqT1jira07DN2PArx2v7/mN66gGcHOs=
 github.com/go-sql-driver/mysql v1.5.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
 github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
@@ -82,6 +83,7 @@ github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Ky
 github.com/mattn/go-isatty v0.0.13/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU=
 github.com/mattn/go-isatty v0.0.14 h1:yVuAays6BHfxijgZPzw+3Zlu5yQgKGP2/hcQbHb7S9Y=
 github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94=
+github.com/mattn/go-sqlite3 v1.10.0/go.mod h1:FPy6KqzDD04eiIsT53CuJW3U88zkxoIYsOqkbpncsNc=
 github.com/mattn/go-sqlite3 v1.14.6 h1:dNPt6NO46WmLVt2DLNpwczCmdV5boIZ6g/tlDrlRUbg=
 github.com/mattn/go-sqlite3 v1.14.6/go.mod h1:NyWgC/yNuGj7Q9rpYnZvas74GogHl5/Z4A/KQRfk6bU=
 github.com/mitchellh/go-testing-interface v0.0.0-20171004221916-a61a99592b77/go.mod h1:kRemZodwjscx+RGhAo8eIhFbs2+BFgRtFPeD/KE+zxI=
@@ -95,16 +97,20 @@ github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9G
 github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
 github.com/oklog/run v1.0.0/go.mod h1:dlhp/R75TPv97u0XWUtDeV/lRKWPKSdTuV0TZvrmrQA=
 github.com/oklog/run v1.1.0/go.mod h1:sVPdnTZT1zYwAJeCMu2Th4T21pA3FPOQRfWjQlk7DVU=
+github.com/opentracing/opentracing-go v1.1.0/go.mod h1:UkNAQd3GIcIGf0SeVgPpRdFStlNbqXla1AfSYxPUl2o=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
+github.com/qustavo/sqlhooks/v2 v2.1.0 h1:54yBemHnGHp/7xgT+pxwmIlMSDNYKx5JW5dfRAiCZi0=
+github.com/qustavo/sqlhooks/v2 v2.1.0/go.mod h1:aMREyKo7fOKTwiLuWPsaHRXEmtqG4yREztO0idF83AU=
 github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
 github.com/sirupsen/logrus v1.8.1 h1:dJKuHgqk1NNQlqoA6BTlM1Wf9DOH3NBjQyu0h9+AZZE=
 github.com/sirupsen/logrus v1.8.1/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
 github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
+github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY=

grammar/hooks/README.md

@@ -0,0 +1,39 @@
+# hooks
+
+给数据库driver注入sqlhooks.Hooks。
+
+## Usage
+
+```go
+package main
+
+import (
+	"github.com/yaoapp/yao/cmd"
+	"github.com/yaoapp/yao/utils"
+
+	_ "github.com/yaoapp/gou/encoding"
+	_ "github.com/yaoapp/yao/aigc"
+	_ "github.com/yaoapp/yao/crypto"
+	_ "github.com/yaoapp/yao/helper"
+	_ "github.com/yaoapp/yao/openai"
+	_ "github.com/yaoapp/yao/wework"
+
+	_ "github.com/yaoapp/yao/xun/grammar/hooks"
+	_ "github.com/yaoapp/yao/xun/grammar/hooks/log"
+	// your customized hooks
+)
+
+// 主程序
+func main() {
+	// 1. 注入你的driver
+	hooks.RegisterDriver("mysql:log")
+	// 2. 此处可定制log hook的字段，例如从context中获取trace_id
+	loghook.Default.ContextFields = func(ctx context.Context) log.F {
+		return log.F{"trace_id": "-"}
+	}
+	utils.Init()
+	cmd.Execute()
+}
+```
+
+在.env文件中，指定YAO_DB_DRIVER="mysql:log"，编译启动即可。

grammar/hooks/hooks.go

@@ -0,0 +1,143 @@
+package hooks
+
+import (
+	"database/sql"
+	"database/sql/driver"
+	"errors"
+	"fmt"
+	"strings"
+
+	"github.com/go-sql-driver/mysql"
+	"github.com/lib/pq"
+	"github.com/mattn/go-sqlite3"
+	"github.com/qustavo/sqlhooks/v2"
+	"github.com/yaoapp/xun/dbal"
+	gmysql "github.com/yaoapp/xun/grammar/mysql"
+	gpostgres "github.com/yaoapp/xun/grammar/postgres"
+	gsql "github.com/yaoapp/xun/grammar/sql"
+	gsqlite3 "github.com/yaoapp/xun/grammar/sqlite3"
+)
+
+var (
+	// Hooks is the `sqlhooks.Hooks` registry
+	Hooks = map[string]sqlhooks.Hooks{}
+
+	// Drivers is the Driver registry
+	Drivers = map[string]*Driver{}
+)
+
+// NoHooksError no hooks error
+var NoHooksError = fmt.Errorf("no hooks error")
+
+// RegisterHook register a `sqlhooks.Hooks`
+func RegisterHook(name string, hook sqlhooks.Hooks) error {
+	if _, ok := Hooks[name]; ok {
+		return fmt.Errorf("hook %s already registered", name)
+	}
+	Hooks[name] = hook
+	return nil
+}
+
+// RegisterDriver register a driver with the given name
+// driver name format: type:hook1[:hook2][:...], type must be one of [mysql, postgres, sqlite3]
+// it will both register the `database/sql/driver` and `yaoapp/xun/dbal`
+func RegisterDriver(name string) (err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			err = fmt.Errorf("recovered from panic: %v", r)
+		}
+	}()
+	if _, ok := Drivers[name]; ok {
+		return nil
+	}
+	d, err := NewDriver(name)
+	if err != nil {
+		if errors.Is(err, NoHooksError) {
+			return nil // NOTE: no need to register if no hooks!
+		}
+		return err
+	}
+	Drivers[name] = d
+	sql.Register(name, d)
+	dbal.Register(name, d.grammar())
+	return nil
+}
+
+// NewDriver create a new driver with `sqlhooks.Hooks` support
+// Usage:
+// 1. register hook, e.g. `hooks.RegisterHook("log", Default)` in `yaoapp/xun/grammar/hooks/log package`
+// 2. register driver, e.g. `hooks.RegisterDriver("mysql:log")`
+func NewDriver(name string) (*Driver, error) {
+	parts := strings.Split(name, ":")
+	typ := parts[0]
+	if typ != "mysql" && typ != "postgres" && typ != "sqlite3" {
+		return nil, fmt.Errorf("driver type %s not in [mysql, postgres, sqlite3]", typ)
+	}
+	if len(parts) == 1 {
+		return nil, NoHooksError
+	}
+	d := &Driver{
+		name: name,
+		typ:  typ,
+	}
+	for i, hookName := range parts {
+		if i == 0 { // NOTE: 0 is the default type driver, no need to repeat
+			continue
+		}
+		if hook, ok := Hooks[hookName]; ok {
+			d.Driver = sqlhooks.Wrap(d.driver(), hook)
+		} else {
+			return nil, fmt.Errorf("sql hook %s not found", hookName)
+		}
+	}
+	if d.Driver == nil {
+		return nil, fmt.Errorf("no driver hooks")
+	}
+	return d, nil
+}
+
+// Driver is the database driver which supports the `sqlhooks` in specific naming rule
+type Driver struct {
+	name string
+	typ  string // NOTE: "oneof=mysql postgres sqlite3"`
+	driver.Driver
+}
+
+// Name returns the name of the driver
+func (d *Driver) Name() string {
+	return d.name
+}
+
+// Type returns the type of the driver
+func (d *Driver) Type() string {
+	return d.typ
+}
+
+func (d *Driver) driver() driver.Driver {
+	if d.Driver != nil {
+		return d.Driver
+	}
+	switch d.typ {
+	case "mysql":
+		return &mysql.MySQLDriver{}
+	case "postgres":
+		return &pq.Driver{}
+	case "sqlite3":
+		return &sqlite3.SQLiteDriver{}
+	}
+	return nil
+}
+
+func (d *Driver) grammar() dbal.Grammar {
+	switch d.typ {
+	case "mysql":
+		return gmysql.New(gsql.WithDriver(d.name))
+	case "postgres":
+		return gpostgres.New(gsql.WithDriver(d.name))
+	case "sqlite3":
+		return gsqlite3.New(gsql.WithDriver(d.name))
+	}
+	return nil
+}
+
+var _ driver.Driver = (*Driver)(nil)

grammar/hooks/log/hook.go

@@ -0,0 +1,129 @@
+package log
+
+import (
+	"context"
+	"database/sql/driver"
+	"errors"
+	"fmt"
+	"strconv"
+	"time"
+
+	"github.com/qustavo/sqlhooks/v2"
+	"github.com/yaoapp/kun/log"
+	"github.com/yaoapp/xun/grammar/hooks"
+)
+
+func init() {
+	err := hooks.RegisterHook("log", Default)
+	if err != nil {
+		panic(err)
+	}
+}
+
+// Default default log hook instance
+var Default = &Hook{
+	Logger:       &KunLogger{},
+	Level:        log.InfoLevel,
+	MaxFieldSize: 256,
+}
+
+var (
+	ErrorFieldName       = "error"
+	QueryFieldName       = "query"
+	RequestTimeFieldName = "rt"
+	ArgFieldPrefix       = "arg_"
+)
+
+// Hook record sql logs
+type Hook struct {
+	Level         log.Level `json:"level,omitempty"`
+	MaxFieldSize  int       `json:"max_field_size,omitempty"`
+	Logger        Logger
+	ContextFields func(ctx context.Context) log.F
+}
+
+// Before hook will print the query with it's args and return the context with the timestamp
+func (h *Hook) Before(ctx context.Context, query string, args ...interface{}) (context.Context, error) {
+	return context.WithValue(ctx, ctxKeyStartTime, time.Now()), nil
+}
+
+// After hook will get the timestamp registered on the Before hook and print the elapsed time
+func (h *Hook) After(ctx context.Context, query string, args ...interface{}) (context.Context, error) {
+	return h.log(ctx, query, nil, args...)
+}
+
+func (h *Hook) OnError(ctx context.Context, err error, query string, args ...interface{}) error {
+	if errors.Is(err, driver.ErrSkip) {
+		return err
+	}
+	_, _ = h.log(ctx, query, err, args...)
+	return err
+}
+
+func (h *Hook) log(ctx context.Context, query string, err error, args ...interface{}) (context.Context, error) {
+	start := ctx.Value(ctxKeyStartTime).(time.Time)
+	rt := time.Since(start).Nanoseconds() / 1e6 // unit: Ms
+	fields := make(log.F)
+	if err != nil {
+		fields[ErrorFieldName] = err
+	}
+	fields[QueryFieldName] = query
+	fields[RequestTimeFieldName] = rt
+	for k, v := range h.ContextFields(ctx) {
+		fields[k] = v
+	}
+	for i, arg := range args {
+		argName := ArgFieldPrefix + strconv.Itoa(i)
+		if h.MaxFieldSize-ellipsisLength <= 0 {
+			fields[argName] = arg
+		} else {
+			fields[argName] = LimitSize(arg, h.MaxFieldSize-ellipsisLength)
+		}
+	}
+	h.Logger.Log(h.Level, "sql log", fields)
+	return ctx, nil
+}
+
+// LimitSize limit the print size of the value
+func LimitSize(value interface{}, n int) interface{} {
+	switch val := value.(type) {
+	case []bool, []complex128, []complex64, []float64, []float32,
+		[]int, []int64, []int32, []int16, []int8, []string, []uint, []uint64, []uint32, []uint16, []uintptr, []time.Time,
+		[]time.Duration, []error:
+		s := fmt.Sprintf("%v", value)
+		if len(s) > n {
+			return s[:n] + "..."
+		}
+		return value
+	case string:
+		if len(val) > n {
+			return val[:n] + "..."
+		}
+		return val
+	case *string:
+		if val == nil {
+			return "<nil>"
+		}
+		if len(*val) > n {
+			return (*val)[:n] + "..."
+		}
+		return *val
+	case []byte:
+		if val == nil {
+			return []byte("<nil>")
+		}
+		if len(val) > n {
+			return string(append(val[:n], []byte("...")...))
+		}
+		return string(val)
+	default:
+		return value
+	}
+}
+
+var (
+	ctxKeyStartTime = struct{}{}
+	ellipsisLength  = 3
+)
+
+var _ sqlhooks.Hooks = (*Hook)(nil)

grammar/hooks/log/logger.go

@@ -0,0 +1,34 @@
+package log
+
+import (
+	"github.com/yaoapp/kun/log"
+)
+
+// Logger log hook logger interface
+type Logger interface {
+	Log(level log.Level, msg string, fields log.F)
+}
+
+// KunLogger implements Logger with kun/log
+type KunLogger struct{}
+
+func (l *KunLogger) Log(level log.Level, msg string, fields log.F) {
+	switch level {
+	case log.TraceLevel:
+		log.With(fields).Trace(msg)
+	case log.DebugLevel:
+		log.With(fields).Debug(msg)
+	case log.InfoLevel:
+		log.With(fields).Info(msg)
+	case log.WarnLevel:
+		log.With(fields).Warn(msg)
+	case log.ErrorLevel:
+		log.With(fields).Error(msg)
+	case log.PanicLevel:
+		log.With(fields).Panic(msg)
+	case log.FatalLevel:
+		log.With(fields).Fatal(msg)
+	}
+}
+
+var _ Logger = (*KunLogger)(nil)