Instrumentation: Add counters and histograms for database queries (#28236)

Signed-off-by: bergquist <carl.bergquist@gmail.com>

Co-authored-by: Sofia Papagiannaki <papagian@users.noreply.github.com>
Co-authored-by: Arve Knudsen <arve.knudsen@gmail.com>

go.mod

@@ -30,6 +30,7 @@ require (
 	github.com/facebookgo/structtag v0.0.0-20150214074306-217e25fb9691 // indirect
 	github.com/facebookgo/subset v0.0.0-20150612182917-8dac2c3c4870 // indirect
 	github.com/fatih/color v1.9.0
+	github.com/gchaincl/sqlhooks v1.3.0
 	github.com/go-macaron/binding v0.0.0-20190806013118-0b4f37bab25b
 	github.com/go-macaron/gzip v0.0.0-20160222043647-cad1c6580a07
 	github.com/go-macaron/session v0.0.0-20190805070824-1a3cdc6f5659

go.sum

@@ -317,6 +317,8 @@ github.com/frankban/quicktest v1.7.2/go.mod h1:jaStnuzAqU1AJdCO0l53JDCJrVDKcS03D
 github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I=
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
 github.com/fsouza/fake-gcs-server v1.7.0/go.mod h1:5XIRs4YvwNbNoz+1JF8j6KLAyDh7RHGAyAK3EP2EsNk=
+github.com/gchaincl/sqlhooks v1.3.0 h1:yKPXxW9a5CjXaVf2HkQn6wn7TZARvbAOAelr3H8vK2Y=
+github.com/gchaincl/sqlhooks v1.3.0/go.mod h1:9BypXnereMT0+Ys8WGWHqzgkkOfHIhyeUCqXC24ra34=
 github.com/getkin/kin-openapi v0.2.0/go.mod h1:V1z9xl9oF5Wt7v32ne4FmiF1alpS4dM6mNzoywPOXlk=
 github.com/getkin/kin-openapi v0.13.0/go.mod h1:WGRs2ZMM1Q8LR1QBEwUxC6RJEfaBcD0s+pcEVXFuAjw=
 github.com/ghodss/yaml v0.0.0-20150909031657-73d445a93680/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=

pkg/services/sqlstore/database_wrapper.go

+package sqlstore
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"time"
+
+	"github.com/gchaincl/sqlhooks"
+	"github.com/go-sql-driver/mysql"
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/services/sqlstore/migrator"
+	"github.com/lib/pq"
+	"github.com/mattn/go-sqlite3"
+	"github.com/prometheus/client_golang/prometheus"
+	"xorm.io/core"
+)
+
+var (
+	databaseQueryCounter   *prometheus.CounterVec
+	databaseQueryHistogram *prometheus.HistogramVec
+)
+
+func init() {
+	databaseQueryCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
+		Namespace: "grafana",
+		Name:      "database_queries_total",
+		Help:      "The total amount of Database queries",
+	}, []string{"status"})
+
+	databaseQueryHistogram = prometheus.NewHistogramVec(prometheus.HistogramOpts{
+		Namespace: "grafana",
+		Name:      "database_queries_duration_seconds",
+		Help:      "Database query histogram",
+		Buckets:   prometheus.ExponentialBuckets(0.0001, 4, 9),
+	}, []string{"status"})
+
+	prometheus.MustRegister(databaseQueryCounter, databaseQueryHistogram)
+}
+
+// WrapDatabaseDriverWithHooks creates a fake database driver that
+// executes pre and post functions which we use to gather metrics about
+// database queries.
+func WrapDatabaseDriverWithHooks(dbType string) string {
+	drivers := map[string]driver.Driver{
+		migrator.SQLITE:   &sqlite3.SQLiteDriver{},
+		migrator.MYSQL:    &mysql.MySQLDriver{},
+		migrator.POSTGRES: &pq.Driver{},
+	}
+
+	d, exist := drivers[dbType]
+	if !exist {
+		return dbType
+	}
+
+	driverWithHooks := dbType + "WithHooks"
+	sql.Register(driverWithHooks, sqlhooks.Wrap(d, &databaseQueryWrapper{log: log.New("sqlstore.metrics")}))
+	core.RegisterDriver(driverWithHooks, &databaseQueryWrapperParser{dbType: dbType})
+	return driverWithHooks
+}
+
+// databaseQueryWrapper satisfies the sqlhook.databaseQueryWrapper interface
+// which allow us to wrap all SQL queries with a `Before` & `After` hook.
+type databaseQueryWrapper struct {
+	log log.Logger
+}
+
+// databaseQueryWrapperKey is used as key to save values in `context.Context`
+type databaseQueryWrapperKey struct{}
+
+// Before hook will print the query with its args and return the context with the timestamp
+func (h *databaseQueryWrapper) Before(ctx context.Context, query string, args ...interface{}) (context.Context, error) {
+	return context.WithValue(ctx, databaseQueryWrapperKey{}, time.Now()), nil
+}
+
+// After hook will get the timestamp registered on the Before hook and print the elapsed time
+func (h *databaseQueryWrapper) After(ctx context.Context, query string, args ...interface{}) (context.Context, error) {
+	begin := ctx.Value(databaseQueryWrapperKey{}).(time.Time)
+	elapsed := time.Since(begin)
+	databaseQueryCounter.WithLabelValues("success").Inc()
+	databaseQueryHistogram.WithLabelValues("success").Observe(elapsed.Seconds())
+	h.log.Debug("query finished", "status", "success", "elapsed time", elapsed, "sql", query)
+	return ctx, nil
+}
+
+// OnError will be called if any error happens
+func (h *databaseQueryWrapper) OnError(ctx context.Context, err error, query string, args ...interface{}) error {
+	status := "error"
+	// https://golang.org/pkg/database/sql/driver/#ErrSkip
+	if err == nil || err == driver.ErrSkip {
+		status = "success"
+	}
+
+	begin := ctx.Value(databaseQueryWrapperKey{}).(time.Time)
+	elapsed := time.Since(begin)
+	databaseQueryCounter.WithLabelValues(status).Inc()
+	databaseQueryHistogram.WithLabelValues(status).Observe(elapsed.Seconds())
+	h.log.Debug("query finished", "status", status, "elapsed time", elapsed, "sql", query, "error", err)
+	return err
+}
+
+type databaseQueryWrapperParser struct {
+	dbType string
+}
+
+func (hp *databaseQueryWrapperParser) Parse(string, string) (*core.Uri, error) {
+	return &core.Uri{
+		DbType: core.DbType(hp.dbType),
+	}, nil
+}

pkg/services/sqlstore/migrator/dialect.go

@@ -53,15 +53,21 @@ type Dialect interface {
 	IsDeadlock(err error) bool
 }
 
+type dialectFunc func(*xorm.Engine) Dialect
+
+var supportedDialects = map[string]dialectFunc{
+	MYSQL:                  NewMysqlDialect,
+	SQLITE:                 NewSqlite3Dialect,
+	POSTGRES:               NewPostgresDialect,
+	MYSQL + "WithHooks":    NewMysqlDialect,
+	SQLITE + "WithHooks":   NewSqlite3Dialect,
+	POSTGRES + "WithHooks": NewPostgresDialect,
+}
+
 func NewDialect(engine *xorm.Engine) Dialect {
 	name := engine.DriverName()
-	switch name {
-	case MYSQL:
-		return NewMysqlDialect(engine)
-	case SQLITE:
-		return NewSqlite3Dialect(engine)
-	case POSTGRES:
-		return NewPostgresDialect(engine)
+	if fn, exist := supportedDialects[name]; exist {
+		return fn(engine)
 	}
 
 	panic("Unsupported database type: " + name)

pkg/services/sqlstore/migrator/mysql_dialect.go

@@ -15,7 +15,7 @@ type Mysql struct {
 	BaseDialect
 }
 
-func NewMysqlDialect(engine *xorm.Engine) *Mysql {
+func NewMysqlDialect(engine *xorm.Engine) Dialect {
 	d := Mysql{}
 	d.BaseDialect.dialect = &d
 	d.BaseDialect.engine = engine

pkg/services/sqlstore/migrator/postgres_dialect.go

@@ -15,7 +15,7 @@ type Postgres struct {
 	BaseDialect
 }
 
-func NewPostgresDialect(engine *xorm.Engine) *Postgres {
+func NewPostgresDialect(engine *xorm.Engine) Dialect {
 	d := Postgres{}
 	d.BaseDialect.dialect = &d
 	d.BaseDialect.engine = engine

pkg/services/sqlstore/migrator/sqlite_dialect.go

@@ -12,7 +12,7 @@ type Sqlite3 struct {
 	BaseDialect
 }
 
-func NewSqlite3Dialect(engine *xorm.Engine) *Sqlite3 {
+func NewSqlite3Dialect(engine *xorm.Engine) Dialect {
 	d := Sqlite3{}
 	d.BaseDialect.dialect = &d
 	d.BaseDialect.engine = engine

pkg/services/sqlstore/sqlstore.go

@@ -238,6 +238,10 @@ func (ss *SqlStore) getEngine() (*xorm.Engine, error) {
 		return nil, err
 	}
 
+	if ss.Cfg.IsDatabaseMetricsEnabled() {
+		ss.dbCfg.Type = WrapDatabaseDriverWithHooks(ss.dbCfg.Type)
+	}
+
 	sqlog.Info("Connecting to DB", "dbtype", ss.dbCfg.Type)
 	if ss.dbCfg.Type == migrator.SQLITE && strings.HasPrefix(connectionString, "file:") {
 		exists, err := fs.Exists(ss.dbCfg.Path)

pkg/setting/setting.go

@@ -340,6 +340,10 @@ func (c Cfg) IsNgAlertEnabled() bool {
 	return c.FeatureToggles["ngalert"]
 }
 
+func (c Cfg) IsDatabaseMetricsEnabled() bool {
+	return c.FeatureToggles["database_metrics"]
+}
+
 type CommandLineArgs struct {
 	Config   string
 	HomePath string