auth: add Query method to audit and call it from server handler

This new call logs queries even when they have an error and show how
much time they took.

Signed-off-by: Javi Fontan <jfontan@gmail.com>

Author: jfontan

auth/audit.go

@@ -71,6 +71,15 @@ func (a *Audit) Allowed(ctx *sql.Context, permission Permission) error {
 	return err
 }
 
+// Query implements AuditQuery interface.
+func (a *Audit) Query(ctx *sql.Context, d time.Duration, err error) {
+	if q, ok := a.auth.(*Audit); ok {
+		q.Query(ctx, d, err)
+	}
+
+	a.method.Query(ctx, d, err)
+}
+
 // NewAuditLog creates a new AuditMethod that logs to a logrus.Logger.
 func NewAuditLog(l *logrus.Logger) AuditMethod {
 	la := l.WithField("system", "audit")

server/handler.go

@@ -6,9 +6,11 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"time"
 
 	errors "gopkg.in/src-d/go-errors.v1"
 	"gopkg.in/src-d/go-mysql-server.v0"
+	"gopkg.in/src-d/go-mysql-server.v0/auth"
 	"gopkg.in/src-d/go-mysql-server.v0/sql"
 
 	"github.com/sirupsen/logrus"
@@ -72,7 +74,7 @@ func (h *Handler) ComQuery(
 	c *mysql.Conn,
 	query string,
 	callback func(*sqltypes.Result) error,
-) error {
+) (err error) {
 	ctx := h.sm.NewContextWithQuery(c, query)
 
 	handled, err := h.handleKill(c, query)
@@ -84,7 +86,14 @@ func (h *Handler) ComQuery(
 		return nil
 	}
 
+	start := time.Now()
 	schema, rows, err := h.e.Query(ctx, query)
+	defer func() {
+		if q, ok := h.e.Auth.(*auth.Audit); ok {
+			q.Query(ctx, time.Since(start), err)
+		}
+	}()
+
 	if err != nil {
 		return err
 	}