From 557c3a9cdda7188337fc3ef55cfebbe7406c1517 Mon Sep 17 00:00:00 2001 From: owensmallwood Date: Tue, 14 May 2024 10:17:38 -0600 Subject: [PATCH] Unified Storage: Adds metrics for database query counts (#87781) --- pkg/services/apiserver/service.go | 2 +- .../apiserver/storage/entity/test/watch_test.go | 2 +- pkg/services/store/entity/db/dbimpl/dbEngine.go | 12 ++++++++---- pkg/services/store/entity/db/dbimpl/dbEngine_test.go | 4 ++-- pkg/services/store/entity/db/dbimpl/dbimpl.go | 9 ++++++--- pkg/services/store/entity/server/service.go | 2 +- .../store/entity/sqlstash/sql_storage_server_test.go | 4 +++- pkg/services/store/entity/tests/common_test.go | 2 +- 8 files changed, 23 insertions(+), 14 deletions(-) diff --git a/pkg/services/apiserver/service.go b/pkg/services/apiserver/service.go index 50ad8387c5b..6dcaa43b433 100644 --- a/pkg/services/apiserver/service.go +++ b/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 } diff --git a/pkg/services/apiserver/storage/entity/test/watch_test.go b/pkg/services/apiserver/storage/entity/test/watch_test.go index 6f05909e413..33bb3af381f 100644 --- a/pkg/services/apiserver/storage/entity/test/watch_test.go +++ b/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() diff --git a/pkg/services/store/entity/db/dbimpl/dbEngine.go b/pkg/services/store/entity/db/dbimpl/dbEngine.go index 889a6ed5f4b..d32d13c3fd6 100644 --- a/pkg/services/store/entity/db/dbimpl/dbEngine.go +++ b/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 } diff --git a/pkg/services/store/entity/db/dbimpl/dbEngine_test.go b/pkg/services/store/entity/db/dbimpl/dbEngine_test.go index 66e29b24874..f67d4d0f3f3 100644 --- a/pkg/services/store/entity/db/dbimpl/dbEngine_test.go +++ b/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) diff --git a/pkg/services/store/entity/db/dbimpl/dbimpl.go b/pkg/services/store/entity/db/dbimpl/dbimpl.go index a2057eaba7d..b064f268a25 100644 --- a/pkg/services/store/entity/db/dbimpl/dbimpl.go +++ b/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 } diff --git a/pkg/services/store/entity/server/service.go b/pkg/services/store/entity/server/service.go index d4879d94db3..89a150ccb39 100644 --- a/pkg/services/store/entity/server/service.go +++ b/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 } diff --git a/pkg/services/store/entity/sqlstash/sql_storage_server_test.go b/pkg/services/store/entity/sqlstash/sql_storage_server_test.go index 41830a540e1..5709c53be4a 100644 --- a/pkg/services/store/entity/sqlstash/sql_storage_server_test.go +++ b/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) diff --git a/pkg/services/store/entity/tests/common_test.go b/pkg/services/store/entity/tests/common_test.go index 58d1444c9bb..ffdaa62a876 100644 --- a/pkg/services/store/entity/tests/common_test.go +++ b/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()