Instrumentation: re-enable database wrapper feature to expose counter and…

Instrumentation: re-enable database wrapper feature to expose counter and histogram for database queries (#29662)

ref https://github.com/grafana/grafana/issues/29489

pkg/services/sqlstore/database_wrapper.go

@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"database/sql"
 	"database/sql/driver"
 	"database/sql/driver"
 	"errors"
 	"errors"
+	"fmt"
 	"time"
 	"time"
 
 
 	"github.com/gchaincl/sqlhooks"
 	"github.com/gchaincl/sqlhooks"
@@ -18,30 +19,23 @@ import (
 )
 )
 
 
 var (
 var (
-	databaseQueryCounter   *prometheus.CounterVec
-	databaseQueryHistogram prometheus.Histogram
+	databaseQueryHistogram *prometheus.HistogramVec
 )
 )
 
 
 func init() {
 func init() {
-	databaseQueryCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
-		Namespace: "grafana",
-		Name:      "database_queries_total",
-		Help:      "The total amount of Database queries",
-	}, []string{"status"})
-
-	databaseQueryHistogram = prometheus.NewHistogram(prometheus.HistogramOpts{
+	databaseQueryHistogram = prometheus.NewHistogramVec(prometheus.HistogramOpts{
 		Namespace: "grafana",
 		Namespace: "grafana",
 		Name:      "database_queries_duration_seconds",
 		Name:      "database_queries_duration_seconds",
 		Help:      "Database query histogram",
 		Help:      "Database query histogram",
 		Buckets:   prometheus.ExponentialBuckets(0.0001, 4, 9),
 		Buckets:   prometheus.ExponentialBuckets(0.0001, 4, 9),
-	})
+	}, []string{"status"})
 
 
-	prometheus.MustRegister(databaseQueryCounter, databaseQueryHistogram)
+	prometheus.MustRegister(databaseQueryHistogram)
 }
 }
 
 
 // WrapDatabaseDriverWithHooks creates a fake database driver that
 // WrapDatabaseDriverWithHooks creates a fake database driver that
 // executes pre and post functions which we use to gather metrics about
 // executes pre and post functions which we use to gather metrics about
-// database queries.
+// database queries. It also registers the metrics.
 func WrapDatabaseDriverWithHooks(dbType string) string {
 func WrapDatabaseDriverWithHooks(dbType string) string {
 	drivers := map[string]driver.Driver{
 	drivers := map[string]driver.Driver{
 		migrator.SQLite:   &sqlite3.SQLiteDriver{},
 		migrator.SQLite:   &sqlite3.SQLiteDriver{},
@@ -56,7 +50,7 @@ func WrapDatabaseDriverWithHooks(dbType string) string {
 
 
 	driverWithHooks := dbType + "WithHooks"
 	driverWithHooks := dbType + "WithHooks"
 	sql.Register(driverWithHooks, sqlhooks.Wrap(d, &databaseQueryWrapper{log: log.New("sqlstore.metrics")}))
 	sql.Register(driverWithHooks, sqlhooks.Wrap(d, &databaseQueryWrapper{log: log.New("sqlstore.metrics")}))
-	core.RegisterDriver(driverWithHooks, &databaseQueryWrapperParser{dbType: dbType})
+	core.RegisterDriver(driverWithHooks, &databaseQueryWrapperDriver{dbType: dbType})
 	return driverWithHooks
 	return driverWithHooks
 }
 }
 
 
@@ -78,8 +72,7 @@ func (h *databaseQueryWrapper) Before(ctx context.Context, query string, args ..
 func (h *databaseQueryWrapper) After(ctx context.Context, query string, args ...interface{}) (context.Context, error) {
 func (h *databaseQueryWrapper) After(ctx context.Context, query string, args ...interface{}) (context.Context, error) {
 	begin := ctx.Value(databaseQueryWrapperKey{}).(time.Time)
 	begin := ctx.Value(databaseQueryWrapperKey{}).(time.Time)
 	elapsed := time.Since(begin)
 	elapsed := time.Since(begin)
-	databaseQueryCounter.WithLabelValues("success").Inc()
-	databaseQueryHistogram.Observe(elapsed.Seconds())
+	databaseQueryHistogram.WithLabelValues("success").Observe(elapsed.Seconds())
 	h.log.Debug("query finished", "status", "success", "elapsed time", elapsed, "sql", query)
 	h.log.Debug("query finished", "status", "success", "elapsed time", elapsed, "sql", query)
 	return ctx, nil
 	return ctx, nil
 }
 }
@@ -94,18 +87,20 @@ func (h *databaseQueryWrapper) OnError(ctx context.Context, err error, query str
 
 
 	begin := ctx.Value(databaseQueryWrapperKey{}).(time.Time)
 	begin := ctx.Value(databaseQueryWrapperKey{}).(time.Time)
 	elapsed := time.Since(begin)
 	elapsed := time.Since(begin)
-	databaseQueryCounter.WithLabelValues(status).Inc()
-	databaseQueryHistogram.Observe(elapsed.Seconds())
+	databaseQueryHistogram.WithLabelValues(status).Observe(elapsed.Seconds())
 	h.log.Debug("query finished", "status", status, "elapsed time", elapsed, "sql", query, "error", err)
 	h.log.Debug("query finished", "status", status, "elapsed time", elapsed, "sql", query, "error", err)
 	return err
 	return err
 }
 }
 
 
-type databaseQueryWrapperParser struct {
+// databaseQueryWrapperDriver satisfies the xorm.io/core.Driver interface
+type databaseQueryWrapperDriver struct {
 	dbType string
 	dbType string
 }
 }
 
 
-func (hp *databaseQueryWrapperParser) Parse(string, string) (*core.Uri, error) {
-	return &core.Uri{
-		DbType: core.DbType(hp.dbType),
-	}, nil
+func (hp *databaseQueryWrapperDriver) Parse(driverName, dataSourceName string) (*core.Uri, error) {
+	driver := core.QueryDriver(hp.dbType)
+	if driver == nil {
+		return nil, fmt.Errorf("could not find driver with name %s", hp.dbType)
+	}
+	return driver.Parse(driverName, dataSourceName)
 }
 }

pkg/services/sqlstore/sqlstore.go

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

pkg/setting/setting.go

@@ -348,6 +348,11 @@ func (cfg Cfg) IsNgAlertEnabled() bool {
 	return cfg.FeatureToggles["ngalert"]
 	return cfg.FeatureToggles["ngalert"]
 }
 }
 
 
+// IsDatabaseMetricsEnabled returns whether the database instrumentation feature is enabled.
+func (cfg Cfg) IsDatabaseMetricsEnabled() bool {
+	return cfg.FeatureToggles["database_metrics"]
+}
+
 // IsHTTPRequestHistogramEnabled returns whether the http_request_histogram feature is enabled.
 // IsHTTPRequestHistogramEnabled returns whether the http_request_histogram feature is enabled.
 func (cfg Cfg) IsHTTPRequestHistogramEnabled() bool {
 func (cfg Cfg) IsHTTPRequestHistogramEnabled() bool {
 	return cfg.FeatureToggles["http_request_histogram"]
 	return cfg.FeatureToggles["http_request_histogram"]