From 51359c9ce8b3975a888d1b97e3ee55797748288c Mon Sep 17 00:00:00 2001 From: yflau <5521134@qq.com> Date: Thu, 9 May 2024 16:52:07 +0800 Subject: [PATCH] add hooks support --- go.mod | 1 + go.sum | 6 ++ grammar/hooks/README.md | 39 ++++++++++ grammar/hooks/hooks.go | 143 +++++++++++++++++++++++++++++++++++ grammar/hooks/log/hook.go | 129 +++++++++++++++++++++++++++++++ grammar/hooks/log/logger.go | 34 +++++++++ grammar/mysql/mysql.go | 8 +- grammar/postgres/postgres.go | 8 +- grammar/sql/sql.go | 15 +++- grammar/sqlite3/sqlite3.go | 8 +- 10 files changed, 381 insertions(+), 10 deletions(-) create mode 100644 grammar/hooks/README.md create mode 100644 grammar/hooks/hooks.go create mode 100644 grammar/hooks/log/hook.go create mode 100644 grammar/hooks/log/logger.go diff --git a/go.mod b/go.mod index a87687f..284c22a 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index f426ac4..b1d39cc 100644 --- a/go.sum +++ b/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,9 +97,12 @@ 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= @@ -105,6 +110,7 @@ github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasO 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= diff --git a/grammar/hooks/README.md b/grammar/hooks/README.md new file mode 100644 index 0000000..5371061 --- /dev/null +++ b/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"，编译启动即可。 diff --git a/grammar/hooks/hooks.go b/grammar/hooks/hooks.go new file mode 100644 index 0000000..5d96072 --- /dev/null +++ b/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) diff --git a/grammar/hooks/log/hook.go b/grammar/hooks/log/hook.go new file mode 100644 index 0000000..7fd9c1d --- /dev/null +++ b/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 "" + } + if len(*val) > n { + return (*val)[:n] + "..." + } + return *val + case []byte: + if val == nil { + return []byte("") + } + 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) diff --git a/grammar/hooks/log/logger.go b/grammar/hooks/log/logger.go new file mode 100644 index 0000000..0534a88 --- /dev/null +++ b/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) diff --git a/grammar/mysql/mysql.go b/grammar/mysql/mysql.go index 3edf061..ac088c8 100644 --- a/grammar/mysql/mysql.go +++ b/grammar/mysql/mysql.go @@ -88,11 +88,13 @@ func (grammarSQL MySQL) OnConnected() error { } // New Create a new MySQL grammar inteface -func New() dbal.Grammar { +func New(opts ...sql.Option) dbal.Grammar { my := MySQL{ - SQL: sql.NewSQL(&Quoter{}), + SQL: sql.NewSQL(&Quoter{}, opts...), + } + if my.Driver == "" { + my.Driver = "mysql" } - my.Driver = "mysql" // set fliptypes flipTypes, ok := utils.MapFilp(my.Types) if ok { diff --git a/grammar/postgres/postgres.go b/grammar/postgres/postgres.go index 26bbaa5..fecd49f 100644 --- a/grammar/postgres/postgres.go +++ b/grammar/postgres/postgres.go @@ -87,11 +87,13 @@ func (grammarSQL Postgres) NewWithRead(write *sqlx.DB, writeConfig *dbal.Config, } // New Create a new mysql grammar inteface -func New() dbal.Grammar { +func New(opts ...sql.Option) dbal.Grammar { pg := Postgres{ - SQL: sql.NewSQL(&Quoter{}), + SQL: sql.NewSQL(&Quoter{}, opts...), + } + if pg.Driver == "" { + pg.Driver = "postgres" } - pg.Driver = "postgres" pg.IndexTypes = map[string]string{ "unique": "UNIQUE INDEX", "index": "INDEX", diff --git a/grammar/sql/sql.go b/grammar/sql/sql.go index a50131d..b9d6ef9 100644 --- a/grammar/sql/sql.go +++ b/grammar/sql/sql.go @@ -29,7 +29,7 @@ type SQL struct { } // NewSQL create a new SQL instance -func NewSQL(quoter dbal.Quoter) SQL { +func NewSQL(quoter dbal.Quoter, opts ...Option) SQL { sql := &SQL{ Driver: "sql", Mode: "production", @@ -71,9 +71,22 @@ func NewSQL(quoter dbal.Quoter) SQL { // "mediumInteger": "mediumInteger", }, } + for _, opt := range opts { + opt(sql) + } return *sql } +// Option used to specify attributes +type Option func(*SQL) + +// WithDriver specify the driver, used in hooks mode +func WithDriver(driver string) Option { + return func(my *SQL) { + my.Driver = driver + } +} + // New Create a new mysql grammar inteface func New(dsn string) dbal.Grammar { sql := NewSQL(&Quoter{}) diff --git a/grammar/sqlite3/sqlite3.go b/grammar/sqlite3/sqlite3.go index c405d3b..c5ee678 100644 --- a/grammar/sqlite3/sqlite3.go +++ b/grammar/sqlite3/sqlite3.go @@ -83,11 +83,13 @@ func (grammarSQL SQLite3) NewWithRead(write *sqlx.DB, writeConfig *dbal.Config, } // New Create a new mysql grammar inteface -func New() dbal.Grammar { +func New(opts ...sql.Option) dbal.Grammar { sqlite := SQLite3{ - SQL: sql.NewSQL(&Quoter{}), + SQL: sql.NewSQL(&Quoter{}, opts...), + } + if sqlite.Driver == "" { + sqlite.Driver = "sqlite3" } - sqlite.Driver = "sqlite3" sqlite.IndexTypes = map[string]string{ "unique": "UNIQUE INDEX", "index": "INDEX",