Unified Storage: Adds metrics for database query counts (#87781)

pkg/services/apiserver/service.go

@@ -249,7 +249,7 @@ func (s *service) start(ctx context.Context) error {
 			return fmt.Errorf("unified storage requires the unifiedStorage feature flag")
 		}
 
-		eDB, err := dbimpl.ProvideEntityDB(s.db, s.cfg, s.features)
+		eDB, err := dbimpl.ProvideEntityDB(s.db, s.cfg, s.features, s.tracing)
 		if err != nil {
 			return err
 		}

pkg/services/apiserver/storage/entity/test/watch_test.go

@@ -66,7 +66,7 @@ func createTestContext(t *testing.T) (entityStore.EntityStoreClient, factory.Des
 	db := sqlstore.InitTestDBWithMigration(t, nil, sqlstore.InitTestDBOpt{EnsureDefaultOrgAndUser: false})
 	require.NoError(t, err)
 
-	eDB, err := dbimpl.ProvideEntityDB(db, cfg, featureToggles)
+	eDB, err := dbimpl.ProvideEntityDB(db, cfg, featureToggles, nil)
 	require.NoError(t, err)
 
 	err = eDB.Init()

pkg/services/store/entity/db/dbimpl/dbEngine.go

@@ -5,12 +5,14 @@ import (
 	"strings"
 	"time"
 
+	"github.com/grafana/grafana/pkg/infra/tracing"
+	"github.com/grafana/grafana/pkg/services/sqlstore"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/util"
 	"xorm.io/xorm"
 )
 
-func getEngineMySQL(cfgSection *setting.DynamicSection) (*xorm.Engine, error) {
+func getEngineMySQL(cfgSection *setting.DynamicSection, tracer tracing.Tracer) (*xorm.Engine, error) {
 	dbHost := cfgSection.Key("db_host").MustString("")
 	dbName := cfgSection.Key("db_name").MustString("")
 	dbUser := cfgSection.Key("db_user").MustString("")
@@ -24,7 +26,8 @@ func getEngineMySQL(cfgSection *setting.DynamicSection) (*xorm.Engine, error) {
 
 	connectionString := connectionStringMySQL(dbUser, dbPass, protocol, dbHost, dbName)
 
-	engine, err := xorm.NewEngine("mysql", connectionString)
+	driverName := sqlstore.WrapDatabaseDriverWithHooks("mysql", tracer)
+	engine, err := xorm.NewEngine(driverName, connectionString)
 	if err != nil {
 		return nil, err
 	}
@@ -36,7 +39,7 @@ func getEngineMySQL(cfgSection *setting.DynamicSection) (*xorm.Engine, error) {
 	return engine, nil
 }
 
-func getEnginePostgres(cfgSection *setting.DynamicSection) (*xorm.Engine, error) {
+func getEnginePostgres(cfgSection *setting.DynamicSection, tracer tracing.Tracer) (*xorm.Engine, error) {
 	dbHost := cfgSection.Key("db_host").MustString("")
 	dbName := cfgSection.Key("db_name").MustString("")
 	dbUser := cfgSection.Key("db_user").MustString("")
@@ -52,7 +55,8 @@ func getEnginePostgres(cfgSection *setting.DynamicSection) (*xorm.Engine, error)
 
 	connectionString := connectionStringPostgres(dbUser, dbPass, addr.Host, addr.Port, dbName, dbSslMode)
 
-	engine, err := xorm.NewEngine("postgres", connectionString)
+	driverName := sqlstore.WrapDatabaseDriverWithHooks("postgres", tracer)
+	engine, err := xorm.NewEngine(driverName, connectionString)
 	if err != nil {
 		return nil, err
 	}

pkg/services/store/entity/db/dbimpl/dbEngine_test.go

@@ -19,7 +19,7 @@ func TestGetEnginePostgresFromConfig(t *testing.T) {
 	s.Key("db_user").SetValue("user")
 	s.Key("db_password").SetValue("password")
 
-	engine, err := getEnginePostgres(cfg.SectionWithEnvOverrides("entity_api"))
+	engine, err := getEnginePostgres(cfg.SectionWithEnvOverrides("entity_api"), nil)
 
 	assert.NotNil(t, engine)
 	assert.NoError(t, err)
@@ -36,7 +36,7 @@ func TestGetEngineMySQLFromConfig(t *testing.T) {
 	s.Key("db_user").SetValue("user")
 	s.Key("db_password").SetValue("password")
 
-	engine, err := getEngineMySQL(cfg.SectionWithEnvOverrides("entity_api"))
+	engine, err := getEngineMySQL(cfg.SectionWithEnvOverrides("entity_api"), nil)
 
 	assert.NotNil(t, engine)
 	assert.NoError(t, err)

pkg/services/store/entity/db/dbimpl/dbimpl.go

@@ -6,6 +6,7 @@ import (
 	"github.com/dlmiddlecote/sqlstats"
 	"github.com/grafana/grafana/pkg/infra/db"
 	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/infra/tracing"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/sqlstore/session"
 	entitydb "github.com/grafana/grafana/pkg/services/store/entity/db"
@@ -18,12 +19,13 @@ import (
 
 var _ entitydb.EntityDBInterface = (*EntityDB)(nil)
 
-func ProvideEntityDB(db db.DB, cfg *setting.Cfg, features featuremgmt.FeatureToggles) (*EntityDB, error) {
+func ProvideEntityDB(db db.DB, cfg *setting.Cfg, features featuremgmt.FeatureToggles, tracer tracing.Tracer) (*EntityDB, error) {
 	return &EntityDB{
 		db:       db,
 		cfg:      cfg,
 		features: features,
 		log:      log.New("entity-db"),
+		tracer:   tracer,
 	}, nil
 }
 
@@ -33,6 +35,7 @@ type EntityDB struct {
 	engine   *xorm.Engine
 	cfg      *setting.Cfg
 	log      log.Logger
+	tracer   tracing.Tracer
 }
 
 func (db *EntityDB) Init() error {
@@ -54,7 +57,7 @@ func (db *EntityDB) GetEngine() (*xorm.Engine, error) {
 	// if explicit connection settings are provided, use them
 	if dbType != "" {
 		if dbType == "postgres" {
-			engine, err = getEnginePostgres(cfgSection)
+			engine, err = getEnginePostgres(cfgSection, db.tracer)
 			if err != nil {
 				return nil, err
 			}
@@ -66,7 +69,7 @@ func (db *EntityDB) GetEngine() (*xorm.Engine, error) {
 				// FIXME: return nil, err
 			}
 		} else if dbType == "mysql" {
-			engine, err = getEngineMySQL(cfgSection)
+			engine, err = getEngineMySQL(cfgSection, db.tracer)
 			if err != nil {
 				return nil, err
 			}

pkg/services/store/entity/server/service.go

@@ -106,7 +106,7 @@ func (s *service) start(ctx context.Context) error {
 	// TODO: use wire
 
 	// TODO: support using grafana db connection?
-	eDB, err := dbimpl.ProvideEntityDB(nil, s.cfg, s.features)
+	eDB, err := dbimpl.ProvideEntityDB(nil, s.cfg, s.features, s.tracing)
 	if err != nil {
 		return err
 	}

pkg/services/store/entity/sqlstash/sql_storage_server_test.go

@@ -136,7 +136,9 @@ func setUpTestServer(t *testing.T) entity.EntityStoreServer {
 	entityDB, err := dbimpl.ProvideEntityDB(
 		sqlStore,
 		cfg,
-		featuremgmt.WithFeatures(featuremgmt.FlagUnifiedStorage))
+		featuremgmt.WithFeatures(featuremgmt.FlagUnifiedStorage),
+		nil,
+	)
 	require.NoError(t, err)
 
 	traceConfig, err := tracing.ParseTracingConfig(cfg)

pkg/services/store/entity/tests/common_test.go

@@ -80,7 +80,7 @@ func createTestContext(t *testing.T) testContext {
 
 	authToken, serviceAccountUser := createServiceAccountAdminToken(t, env)
 
-	eDB, err := dbimpl.ProvideEntityDB(env.SQLStore, env.Cfg, env.FeatureToggles)
+	eDB, err := dbimpl.ProvideEntityDB(env.SQLStore, env.Cfg, env.FeatureToggles, nil)
 	require.NoError(t, err)
 
 	err = eDB.Init()