Resource server improvements and fixes (#90715)

* cleanup dependencies and improve list method
* Improve Resource Server API, remove unnecessary dependencies
* Reduce the API footprint of ResourceDBInterface and its implementation
* Improve LifecycleHooks to use context
* Improve testing
* reduce API size and improve code
* sqltemplate: add DialectForDriver func and improve naming
* improve lifecycle API
* many small fixes after adding more tests

pkg/storage/unified/sql/db/dbimpl/db.go

@@ -4,11 +4,16 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"strings"
 
-	resourcedb "github.com/grafana/grafana/pkg/storage/unified/sql/db"
+	"github.com/grafana/grafana/pkg/storage/unified/sql/db"
 )
 
-func NewDB(d *sql.DB, driverName string) resourcedb.DB {
+func NewDB(d *sql.DB, driverName string) db.DB {
+	// remove the suffix from the instrumented driver created by the older
+	// Grafana code
+	driverName = strings.TrimSuffix(driverName, "WithHooks")
+
 	return sqldb{
 		DB:         d,
 		driverName: driverName,
@@ -24,7 +29,7 @@ func (d sqldb) DriverName() string {
 	return d.driverName
 }
 
-func (d sqldb) BeginTx(ctx context.Context, opts *sql.TxOptions) (resourcedb.Tx, error) {
+func (d sqldb) BeginTx(ctx context.Context, opts *sql.TxOptions) (db.Tx, error) {
 	t, err := d.DB.BeginTx(ctx, opts)
 	if err != nil {
 		return nil, err
@@ -34,7 +39,7 @@ func (d sqldb) BeginTx(ctx context.Context, opts *sql.TxOptions) (resourcedb.Tx,
 	}, nil
 }
 
-func (d sqldb) WithTx(ctx context.Context, opts *sql.TxOptions, f resourcedb.TxFunc) error {
+func (d sqldb) WithTx(ctx context.Context, opts *sql.TxOptions, f db.TxFunc) error {
 	t, err := d.BeginTx(ctx, opts)
 	if err != nil {
 		return fmt.Errorf("begin tx: %w", err)

pkg/storage/unified/sql/db/dbimpl/dbEngine.go

@@ -7,13 +7,13 @@ import (
 	"time"
 
 	"github.com/go-sql-driver/mysql"
+	"go.opentelemetry.io/otel/trace"
 	"xorm.io/xorm"
 
-	"github.com/grafana/grafana/pkg/infra/tracing"
 	"github.com/grafana/grafana/pkg/services/store/entity/db"
 )
 
-func getEngineMySQL(getter *sectionGetter, _ tracing.Tracer) (*xorm.Engine, error) {
+func getEngineMySQL(getter *sectionGetter, _ trace.Tracer) (*xorm.Engine, error) {
 	config := mysql.NewConfig()
 	config.User = getter.String("db_user")
 	config.Passwd = getter.String("db_pass")
@@ -29,6 +29,11 @@ func getEngineMySQL(getter *sectionGetter, _ tracing.Tracer) (*xorm.Engine, erro
 	config.AllowNativePasswords = true
 	config.ClientFoundRows = true
 
+	// allow executing multiple SQL statements in a single roundtrip, and also
+	// enable executing the CALL statement to run stored procedures that execute
+	// multiple SQL statements.
+	//config.MultiStatements = true
+
 	// TODO: do we want to support these?
 	//	config.ServerPubKey = getter.String("db_server_pub_key")
 	//	config.TLSConfig = getter.String("db_tls_config_name")
@@ -54,7 +59,7 @@ func getEngineMySQL(getter *sectionGetter, _ tracing.Tracer) (*xorm.Engine, erro
 	return engine, nil
 }
 
-func getEnginePostgres(getter *sectionGetter, _ tracing.Tracer) (*xorm.Engine, error) {
+func getEnginePostgres(getter *sectionGetter, _ trace.Tracer) (*xorm.Engine, error) {
 	dsnKV := map[string]string{
 		"user":     getter.String("db_user"),
 		"password": getter.String("db_pass"),

pkg/storage/unified/sql/db/dbimpl/dbEngine_test.go

@@ -6,20 +6,22 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func newValidMySQLGetter() *sectionGetter {
+	return newTestSectionGetter(map[string]string{
+		"db_type":     dbTypeMySQL,
+		"db_host":     "/var/run/mysql.socket",
+		"db_name":     "grafana",
+		"db_user":     "user",
+		"db_password": "password",
+	})
+}
+
 func TestGetEngineMySQLFromConfig(t *testing.T) {
 	t.Parallel()
 
 	t.Run("happy path", func(t *testing.T) {
 		t.Parallel()
-
-		getter := newTestSectionGetter(map[string]string{
-			"db_type":     "mysql",
-			"db_host":     "/var/run/mysql.socket",
-			"db_name":     "grafana",
-			"db_user":     "user",
-			"db_password": "password",
-		})
-		engine, err := getEngineMySQL(getter, nil)
+		engine, err := getEngineMySQL(newValidMySQLGetter(), nil)
 		assert.NotNil(t, engine)
 		assert.NoError(t, err)
 	})
@@ -28,7 +30,7 @@ func TestGetEngineMySQLFromConfig(t *testing.T) {
 		t.Parallel()
 
 		getter := newTestSectionGetter(map[string]string{
-			"db_type":     "mysql",
+			"db_type":     dbTypeMySQL,
 			"db_host":     "/var/run/mysql.socket",
 			"db_name":     string(invalidUTF8ByteSequence),
 			"db_user":     "user",
@@ -37,7 +39,17 @@ func TestGetEngineMySQLFromConfig(t *testing.T) {
 		engine, err := getEngineMySQL(getter, nil)
 		assert.Nil(t, engine)
 		assert.Error(t, err)
-		assert.ErrorIs(t, err, ErrInvalidUTF8Sequence)
+		assert.ErrorIs(t, err, errInvalidUTF8Sequence)
+	})
+}
+
+func newValidPostgresGetter() *sectionGetter {
+	return newTestSectionGetter(map[string]string{
+		"db_type":     dbTypePostgres,
+		"db_host":     "localhost",
+		"db_name":     "grafana",
+		"db_user":     "user",
+		"db_password": "password",
 	})
 }
 
@@ -46,15 +58,7 @@ func TestGetEnginePostgresFromConfig(t *testing.T) {
 
 	t.Run("happy path", func(t *testing.T) {
 		t.Parallel()
-		getter := newTestSectionGetter(map[string]string{
-			"db_type":     "mysql",
-			"db_host":     "localhost",
-			"db_name":     "grafana",
-			"db_user":     "user",
-			"db_password": "password",
-		})
-		engine, err := getEnginePostgres(getter, nil)
-
+		engine, err := getEnginePostgres(newValidPostgresGetter(), nil)
 		assert.NotNil(t, engine)
 		assert.NoError(t, err)
 	})
@@ -62,7 +66,7 @@ func TestGetEnginePostgresFromConfig(t *testing.T) {
 	t.Run("invalid string", func(t *testing.T) {
 		t.Parallel()
 		getter := newTestSectionGetter(map[string]string{
-			"db_type":     "mysql",
+			"db_type":     dbTypePostgres,
 			"db_host":     string(invalidUTF8ByteSequence),
 			"db_name":     "grafana",
 			"db_user":     "user",
@@ -72,13 +76,13 @@ func TestGetEnginePostgresFromConfig(t *testing.T) {
 
 		assert.Nil(t, engine)
 		assert.Error(t, err)
-		assert.ErrorIs(t, err, ErrInvalidUTF8Sequence)
+		assert.ErrorIs(t, err, errInvalidUTF8Sequence)
 	})
 
 	t.Run("invalid hostport", func(t *testing.T) {
 		t.Parallel()
 		getter := newTestSectionGetter(map[string]string{
-			"db_type":     "mysql",
+			"db_type":     dbTypePostgres,
 			"db_host":     "1:1:1",
 			"db_name":     "grafana",
 			"db_user":     "user",

pkg/storage/unified/sql/db/dbimpl/db_test.go

@@ -9,7 +9,7 @@ import (
 	sqlmock "github.com/DATA-DOG/go-sqlmock"
 	"github.com/stretchr/testify/require"
 
-	resourcedb "github.com/grafana/grafana/pkg/storage/unified/sql/db"
+	"github.com/grafana/grafana/pkg/storage/unified/sql/db"
 )
 
 func newCtx(t *testing.T) context.Context {
@@ -71,8 +71,8 @@ func TestDB_BeginTx(t *testing.T) {
 func TestDB_WithTx(t *testing.T) {
 	t.Parallel()
 
-	newTxFunc := func(err error) resourcedb.TxFunc {
-		return func(context.Context, resourcedb.Tx) error {
+	newTxFunc := func(err error) db.TxFunc {
+		return func(context.Context, db.Tx) error {
 			return err
 		}
 	}

pkg/storage/unified/sql/db/dbimpl/dbimpl.go

@@ -1,166 +1,116 @@
 package dbimpl
 
 import (
+	"context"
 	"fmt"
 	"sync"
 
 	"github.com/dlmiddlecote/sqlstats"
-	"github.com/jmoiron/sqlx"
 	"github.com/prometheus/client_golang/prometheus"
+	"go.opentelemetry.io/otel/trace"
 	"xorm.io/xorm"
 
-	"github.com/grafana/grafana/pkg/infra/db"
+	infraDB "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"
 	"github.com/grafana/grafana/pkg/setting"
-	resourcedb "github.com/grafana/grafana/pkg/storage/unified/sql/db"
+	"github.com/grafana/grafana/pkg/storage/unified/sql/db"
 	"github.com/grafana/grafana/pkg/storage/unified/sql/db/migrations"
 )
 
-var _ resourcedb.ResourceDBInterface = (*ResourceDB)(nil)
+const (
+	dbTypeMySQL    = "mysql"
+	dbTypePostgres = "postgres"
+)
 
-func ProvideResourceDB(db db.DB, cfg *setting.Cfg, features featuremgmt.FeatureToggles, tracer tracing.Tracer) (*ResourceDB, error) {
-	return &ResourceDB{
-		db:       db,
-		cfg:      cfg,
-		features: features,
-		log:      log.New("entity-db"),
-		tracer:   tracer,
-	}, nil
-}
-
-type ResourceDB struct {
-	once    sync.Once
-	onceErr error
-
-	db       db.DB
-	features featuremgmt.FeatureToggles
-	engine   *xorm.Engine
-	cfg      *setting.Cfg
-	log      log.Logger
-	tracer   tracing.Tracer
-}
-
-func (db *ResourceDB) Init() error {
-	db.once.Do(func() {
-		db.onceErr = db.init()
-	})
-
-	return db.onceErr
-}
-
-func (db *ResourceDB) GetEngine() (*xorm.Engine, error) {
-	if err := db.Init(); err != nil {
-		return nil, err
+func ProvideResourceDB(grafanaDB infraDB.DB, cfg *setting.Cfg, features featuremgmt.FeatureToggles, tracer trace.Tracer) (db.DBProvider, error) {
+	p, err := newResourceDBProvider(grafanaDB, cfg, features, tracer)
+	if err != nil {
+		return nil, fmt.Errorf("provide Resource DB: %w", err)
 	}
+	var once sync.Once
+	var resourceDB db.DB
 
-	return db.engine, db.onceErr
+	return dbProviderFunc(func(ctx context.Context) (db.DB, error) {
+		once.Do(func() {
+			resourceDB, err = p.init(ctx)
+		})
+		return resourceDB, err
+	}), nil
 }
 
-func (db *ResourceDB) init() error {
-	if db.engine != nil {
-		return nil
-	}
+type dbProviderFunc func(context.Context) (db.DB, error)
 
-	var engine *xorm.Engine
-	var err error
+func (f dbProviderFunc) Init(ctx context.Context) (db.DB, error) {
+	return f(ctx)
+}
 
+type resourceDBProvider struct {
+	engine          *xorm.Engine
+	cfg             *setting.Cfg
+	log             log.Logger
+	migrateFunc     func(context.Context, *xorm.Engine, *setting.Cfg) error
+	registerMetrics bool
+	logQueries      bool
+}
+
+func newResourceDBProvider(grafanaDB infraDB.DB, cfg *setting.Cfg, features featuremgmt.FeatureToggles, tracer trace.Tracer) (p *resourceDBProvider, err error) {
 	// TODO: This should be renamed resource_api
 	getter := &sectionGetter{
-		DynamicSection: db.cfg.SectionWithEnvOverrides("resource_api"),
+		DynamicSection: cfg.SectionWithEnvOverrides("resource_api"),
 	}
 
-	dbType := getter.Key("db_type").MustString("")
-
-	// if explicit connection settings are provided, use them
-	if dbType != "" {
-		if dbType == "postgres" {
-			engine, err = getEnginePostgres(getter, db.tracer)
-			if err != nil {
-				return err
-			}
-
-			// FIXME: this config option is cockroachdb-specific, it's not supported by postgres
-			// FIXME: this only sets this option for the session that we get
-			// from the pool right now. A *sql.DB is a pool of connections,
-			// there is no guarantee that the session where this is run will be
-			// the same where we need to change the type of a column
-			_, err = engine.Exec("SET SESSION enable_experimental_alter_column_type_general=true")
-			if err != nil {
-				db.log.Error("error connecting to postgres", "msg", err.Error())
-				// FIXME: return nil, err
-			}
-		} else if dbType == "mysql" {
-			engine, err = getEngineMySQL(getter, db.tracer)
-			if err != nil {
-				return err
-			}
-
-			if err = engine.Ping(); err != nil {
-				return err
-			}
-		} else {
-			// TODO: sqlite support
-			return fmt.Errorf("invalid db type specified: %s", dbType)
-		}
-
-		// register sql stat metrics
-		if err := prometheus.Register(sqlstats.NewStatsCollector("unified_storage", engine.DB().DB)); err != nil {
-			db.log.Warn("Failed to register unified storage sql stats collector", "error", err)
-		}
-
-		// configure sql logging
-		debugSQL := getter.Key("log_queries").MustBool(false)
-		if !debugSQL {
-			engine.SetLogger(&xorm.DiscardLogger{})
-		} else {
-			// add stack to database calls to be able to see what repository initiated queries. Top 7 items from the stack as they are likely in the xorm library.
-			// engine.SetLogger(sqlstore.NewXormLogger(log.LvlInfo, log.WithSuffix(log.New("sqlstore.xorm"), log.CallerContextKey, log.StackCaller(log.DefaultCallerDepth))))
-			engine.ShowSQL(true)
-			engine.ShowExecTime(true)
-		}
-
-		// otherwise, try to use the grafana db connection
-	} else {
-		if db.db == nil {
-			return fmt.Errorf("no db connection provided")
-		}
-
-		engine = db.db.GetEngine()
+	p = &resourceDBProvider{
+		cfg:        cfg,
+		log:        log.New("entity-db"),
+		logQueries: getter.Key("log_queries").MustBool(false),
+	}
+	if features.IsEnabledGlobally(featuremgmt.FlagUnifiedStorage) {
+		p.migrateFunc = migrations.MigrateResourceStore
 	}
 
-	db.engine = engine
+	switch dbType := getter.Key("db_type").MustString(""); dbType {
+	case dbTypePostgres:
+		p.registerMetrics = true
+		p.engine, err = getEnginePostgres(getter, tracer)
+		return p, err
 
-	if err := migrations.MigrateResourceStore(engine, db.cfg, db.features); err != nil {
-		db.engine = nil
-		return fmt.Errorf("run migrations: %w", err)
+	case dbTypeMySQL:
+		p.registerMetrics = true
+		p.engine, err = getEngineMySQL(getter, tracer)
+		return p, err
+
+	case "":
+		// try to use the grafana db connection
+		if grafanaDB == nil {
+			return p, fmt.Errorf("no db connection provided")
+		}
+		p.engine = grafanaDB.GetEngine()
+		return p, nil
+
+	default:
+		// TODO: sqlite support
+		return p, fmt.Errorf("invalid db type specified: %s", dbType)
 	}
-
-	return nil
 }
 
-func (db *ResourceDB) GetSession() (*session.SessionDB, error) {
-	engine, err := db.GetEngine()
-	if err != nil {
-		return nil, err
+func (p *resourceDBProvider) init(ctx context.Context) (db.DB, error) {
+	if p.registerMetrics {
+		err := prometheus.Register(sqlstats.NewStatsCollector("unified_storage", p.engine.DB().DB))
+		if err != nil {
+			p.log.Warn("Failed to register unified storage sql stats collector", "error", err)
+		}
+	}
+	_ = p.logQueries // TODO: configure SQL logging
+
+	// TODO: change the migrator to use db.DB instead of xorm
+	// Skip migrations if feature flag is not enabled
+	if p.migrateFunc != nil {
+		err := p.migrateFunc(ctx, p.engine, p.cfg)
+		if err != nil {
+			return nil, fmt.Errorf("run migrations: %w", err)
+		}
 	}
 
-	return session.GetSession(sqlx.NewDb(engine.DB().DB, engine.DriverName())), nil
-}
-
-func (db *ResourceDB) GetCfg() *setting.Cfg {
-	return db.cfg
-}
-
-func (db *ResourceDB) GetDB() (resourcedb.DB, error) {
-	engine, err := db.GetEngine()
-	if err != nil {
-		return nil, err
-	}
-
-	ret := NewDB(engine.DB().DB, engine.Dialect().DriverName())
-
-	return ret, nil
+	return NewDB(p.engine.DB().DB, p.engine.Dialect().DriverName()), nil
 }

pkg/storage/unified/sql/db/dbimpl/util.go

@@ -12,9 +12,7 @@ import (
 	"github.com/grafana/grafana/pkg/setting"
 )
 
-var (
-	ErrInvalidUTF8Sequence = errors.New("invalid UTF-8 sequence")
-)
+var errInvalidUTF8Sequence = errors.New("invalid UTF-8 sequence")
 
 type sectionGetter struct {
 	*setting.DynamicSection
@@ -28,7 +26,7 @@ func (g *sectionGetter) Err() error {
 func (g *sectionGetter) String(key string) string {
 	v := g.DynamicSection.Key(key).MustString("")
 	if !utf8.ValidString(v) {
-		g.err = fmt.Errorf("value for key %q: %w", key, ErrInvalidUTF8Sequence)
+		g.err = fmt.Errorf("value for key %q: %w", key, errInvalidUTF8Sequence)
 
 		return ""
 	}
@@ -47,7 +45,7 @@ func MakeDSN(m map[string]string) (string, error) {
 		v := m[k]
 		if !utf8.ValidString(v) {
 			return "", fmt.Errorf("value for DSN key %q: %w", k,
-				ErrInvalidUTF8Sequence)
+				errInvalidUTF8Sequence)
 		}
 		if v == "" {
 			continue

pkg/storage/unified/sql/db/dbimpl/util_test.go

@@ -4,8 +4,9 @@ import (
 	"fmt"
 	"testing"
 
-	"github.com/grafana/grafana/pkg/setting"
 	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/grafana/pkg/setting"
 )
 
 var invalidUTF8ByteSequence = []byte{0xff, 0xfe, 0xfd}
@@ -44,7 +45,7 @@ func TestSectionGetter(t *testing.T) {
 	v = g.String(key)
 	require.Empty(t, v)
 	require.Error(t, g.Err())
-	require.ErrorIs(t, g.Err(), ErrInvalidUTF8Sequence)
+	require.ErrorIs(t, g.Err(), errInvalidUTF8Sequence)
 }
 
 func TestMakeDSN(t *testing.T) {
@@ -55,7 +56,7 @@ func TestMakeDSN(t *testing.T) {
 	})
 	require.Empty(t, s)
 	require.Error(t, err)
-	require.ErrorIs(t, err, ErrInvalidUTF8Sequence)
+	require.ErrorIs(t, err, errInvalidUTF8Sequence)
 
 	s, err = MakeDSN(map[string]string{
 		"skip": "",

pkg/storage/unified/sql/db/migrations/migrator.go

@@ -1,18 +1,16 @@
 package migrations
 
 import (
+	"context"
+
 	"xorm.io/xorm"
 
-	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/sqlstore/migrator"
 	"github.com/grafana/grafana/pkg/setting"
 )
 
-func MigrateResourceStore(engine *xorm.Engine, cfg *setting.Cfg, features featuremgmt.FeatureToggles) error {
-	// Skip if feature flag is not enabled
-	if !features.IsEnabledGlobally(featuremgmt.FlagUnifiedStorage) {
-		return nil
-	}
+func MigrateResourceStore(_ context.Context, engine *xorm.Engine, cfg *setting.Cfg) error {
+	// TODO: use the context.Context
 
 	mg := migrator.NewScopedMigrator(engine, cfg, "resource")
 	mg.AddCreateMigration()

pkg/storage/unified/sql/db/service.go

@@ -3,30 +3,13 @@ package db
 import (
 	"context"
 	"database/sql"
-
-	"xorm.io/xorm"
-
-	"github.com/grafana/grafana/pkg/services/sqlstore/session"
-	"github.com/grafana/grafana/pkg/setting"
 )
 
-const (
-	DriverPostgres = "postgres"
-	DriverMySQL    = "mysql"
-	DriverSQLite   = "sqlite"
-	DriverSQLite3  = "sqlite3"
-)
-
-// ResourceDBInterface provides access to a database capable of supporting the
-// Entity Server.
-type ResourceDBInterface interface {
-	Init() error
-	GetCfg() *setting.Cfg
-	GetDB() (DB, error)
-
-	// TODO: deprecate.
-	GetSession() (*session.SessionDB, error)
-	GetEngine() (*xorm.Engine, error)
+// DBProvider provides access to a SQL Database.
+type DBProvider interface {
+	// Init initializes the SQL Database, running migrations if needed. It is
+	// idempotent and thread-safe.
+	Init(context.Context) (DB, error)
 }
 
 // DB is a thin abstraction on *sql.DB to allow mocking to provide better unit