database/sqlutil/logdriver.go

   1 package sqlutil
   2
   3 import (
   4         "context"
   5         "database/sql/driver"
   6         "fmt"
   7
   8         "chain/log"
   9 )
  10
  11 // TODO(kr): many databases—Postgres in particular—report the
  12 // execution time of each query or statement as measured on the
  13 // database backend. Find a way to record that timing info in
  14 // the trace.
  15
  16 const maxArgsLogLen = 20 // bytes
  17
  18 func logQuery(ctx context.Context, query string, args interface{}) {
  19         s := fmt.Sprint(args)
  20         if len(s) > maxArgsLogLen {
  21                 s = s[:maxArgsLogLen-3] + "..."
  22         }
  23         log.Printkv(ctx, "query", query, "args", s)
  24 }
  25
  26 type logDriver struct {
  27         driver driver.Driver
  28 }
  29
  30 // LogDriver returns a Driver that logs each query
  31 // before forwarding it to d.
  32 func LogDriver(d driver.Driver) driver.Driver {
  33         return &logDriver{d}
  34 }
  35
  36 func (ld *logDriver) Open(name string) (driver.Conn, error) {
  37         c, err := ld.driver.Open(name)
  38         return &logConn{c}, err
  39 }
  40
  41 type logConn struct {
  42         driver.Conn
  43 }
  44
  45 func (lc *logConn) Prepare(query string) (driver.Stmt, error) {
  46         stmt, err := lc.Conn.Prepare(query)
  47         return &logStmt{query, stmt}, err
  48 }
  49
  50 func (lc *logConn) Exec(query string, args []driver.Value) (driver.Result, error) {
  51         execer, ok := lc.Conn.(driver.Execer)
  52         if !ok {
  53                 return nil, driver.ErrSkip
  54         }
  55         logQuery(context.Background(), query, args)
  56         return execer.Exec(query, args)
  57 }
  58
  59 func (lc *logConn) Query(query string, args []driver.Value) (driver.Rows, error) {
  60         queryer, ok := lc.Conn.(driver.Queryer)
  61         if !ok {
  62                 return nil, driver.ErrSkip
  63         }
  64         logQuery(context.Background(), query, args)
  65         return queryer.Query(query, args)
  66 }
  67
  68 // TODO(kr): implement context variants
  69 // (but don't bother until lib/pq does first).
  70 //func (lc *logConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error)
  71 //func (lc *logConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error)
  72 //func (lc *logConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error)
  73
  74 type logStmt struct {
  75         query string
  76         driver.Stmt
  77 }
  78
  79 func (ls *logStmt) Exec(args []driver.Value) (driver.Result, error) {
  80         logQuery(context.Background(), ls.query, args)
  81         return ls.Stmt.Exec(args)
  82 }
  83
  84 func (ls *logStmt) Query(args []driver.Value) (driver.Rows, error) {
  85         logQuery(context.Background(), ls.query, args)
  86         return ls.Stmt.Query(args)
  87 }