Sergiusz Bazanski | 1fad2e5 | 2019-08-01 20:16:27 +0200 | [diff] [blame] | 1 | package mirko |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "database/sql" |
| 6 | "database/sql/driver" |
| 7 | "time" |
| 8 | |
| 9 | "github.com/gchaincl/sqlhooks" |
| 10 | "golang.org/x/net/trace" |
| 11 | ) |
| 12 | |
| 13 | type sqlHooks struct{} |
| 14 | |
| 15 | func (h *sqlHooks) Before(ctx context.Context, query string, args ...interface{}) (context.Context, error) { |
| 16 | tr, ok := trace.FromContext(ctx) |
| 17 | if ok { |
| 18 | tr.LazyPrintf("SQL query: %s", query) |
| 19 | tr.LazyPrintf("SQL args: %+v", args) |
| 20 | } |
| 21 | return context.WithValue(ctx, "begin", time.Now()), nil |
| 22 | } |
| 23 | |
| 24 | func (h *sqlHooks) After(ctx context.Context, query string, args ...interface{}) (context.Context, error) { |
| 25 | begin := ctx.Value("begin").(time.Time) |
| 26 | tr, ok := trace.FromContext(ctx) |
| 27 | if ok { |
| 28 | tr.LazyPrintf("SQL took: %s", time.Since(begin).String()) |
| 29 | } |
| 30 | return ctx, nil |
| 31 | } |
| 32 | |
| 33 | func TraceSQL(driver driver.Driver, wrapped string) { |
| 34 | sql.Register(wrapped, sqlhooks.Wrap(driver, &sqlHooks{})) |
| 35 | } |