SQLStore: Introduce new test infra (#101676)

This creates the new testinfra. It is not used anywhere yet.

pkg/services/sqlstore/sqlstore.go

@@ -61,7 +61,7 @@ func ProvideService(cfg *setting.Cfg,
 	// by that mimic the functionality of how it was functioning before
 	// xorm's changes above.
 	xorm.DefaultPostgresSchema = ""
-	s, err := newSQLStore(cfg, nil, features, migrations, bus, tracer)
+	s, err := newStore(cfg, nil, features, migrations, bus, tracer, false)
 	if err != nil {
 		return nil, err
 	}
@@ -87,25 +87,21 @@ func ProvideServiceForTests(t sqlutil.ITestDB, cfg *setting.Cfg, features featur
 func NewSQLStoreWithoutSideEffects(cfg *setting.Cfg,
 	features featuremgmt.FeatureToggles,
 	bus bus.Bus, tracer tracing.Tracer) (*SQLStore, error) {
-	return newSQLStore(cfg, nil, features, nil, bus, tracer)
+	return newStore(cfg, nil, features, nil, bus, tracer, true)
 }
 
-func newSQLStore(cfg *setting.Cfg, engine *xorm.Engine, features featuremgmt.FeatureToggles,
-	migrations registry.DatabaseMigrator, bus bus.Bus, tracer tracing.Tracer, opts ...InitTestDBOpt) (*SQLStore, error) {
+func newStore(cfg *setting.Cfg, engine *xorm.Engine, features featuremgmt.FeatureToggles,
+	migrations registry.DatabaseMigrator, bus bus.Bus, tracer tracing.Tracer,
+	skipEnsureDefaultOrgAndUser bool) (*SQLStore, error) {
 	ss := &SQLStore{
 		cfg:                         cfg,
 		log:                         log.New("sqlstore"),
-		skipEnsureDefaultOrgAndUser: false,
+		skipEnsureDefaultOrgAndUser: skipEnsureDefaultOrgAndUser,
 		migrations:                  migrations,
 		bus:                         bus,
 		tracer:                      tracer,
 		features:                    features,
 	}
-	for _, opt := range opts {
-		if !opt.EnsureDefaultOrgAndUser {
-			ss.skipEnsureDefaultOrgAndUser = true
-		}
-	}
 
 	if err := ss.initEngine(engine); err != nil {
 		return nil, fmt.Errorf("%v: %w", "failed to connect to database", err)
@@ -600,9 +596,17 @@ func TestMain(m *testing.M) {
 		engine.DatabaseTZ = time.UTC
 		engine.TZLocation = time.UTC
 
+		skipEnsureDefaultOrgAndUser := false
+		for _, opt := range opts {
+			if !opt.EnsureDefaultOrgAndUser {
+				skipEnsureDefaultOrgAndUser = true
+				break
+			}
+		}
+
 		tracer := tracing.InitializeTracerForTest()
 		bus := bus.ProvideBus(tracer)
-		testSQLStore, err = newSQLStore(cfg, engine, features, migration, bus, tracer, opts...)
+		testSQLStore, err = newStore(cfg, engine, features, migration, bus, tracer, skipEnsureDefaultOrgAndUser)
 		if err != nil {
 			return nil, err
 		}

pkg/services/sqlstore/sqlstore_testinfra.go

@@ -0,0 +1,369 @@
+// This file sets up the test environment for the sqlstore.
+// Its intent is to create a database for use in tests. This database should be entirely isolated and possible to use in parallel tests.
+
+package sqlstore
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/grafana/grafana/pkg/bus"
+	"github.com/grafana/grafana/pkg/infra/tracing"
+	"github.com/grafana/grafana/pkg/registry"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
+	"github.com/grafana/grafana/pkg/services/sqlstore/migrations"
+	"github.com/grafana/grafana/pkg/setting"
+	"xorm.io/xorm"
+)
+
+// testingTB is an interface that is implemented by *testing.T and *testing.B. Similar to testing.TB.
+type TestingTB interface {
+	// Helper marks the calling function as a test helper function. See also (*testing.T).Helper.
+	Helper()
+	// Cleanup registers a new function the testing suite will run after the test completes. See also (*testing.T).Cleanup.
+	Cleanup(func())
+	// Fatalf logs a message and marks the test as failed. The syntax is similar to that of fmt.Printf. See also (*testing.T).Fatalf.
+	Fatalf(format string, args ...any)
+}
+
+type testOptions struct {
+	FeatureFlags     map[string]bool
+	MigratorFactory  func(featuremgmt.FeatureToggles) registry.DatabaseMigrator
+	Tracer           tracing.Tracer
+	Bus              bus.Bus
+	NoDefaultUserOrg bool
+	Cfg              *setting.Cfg
+	Truncate         bool
+}
+
+type TestOption func(*testOptions)
+
+// WithFeatureFlags adds the feature flags to the other flags already set with a value of true.
+func WithFeatureFlags(flags ...string) TestOption {
+	return func(o *testOptions) {
+		for _, flag := range flags {
+			o.FeatureFlags[flag] = true
+		}
+	}
+}
+
+// WithoutFeatureFlags adds the feature flags to the other flags already set with a value of true.
+func WithoutFeatureFlags(flags ...string) TestOption {
+	return func(o *testOptions) {
+		for _, flag := range flags {
+			o.FeatureFlags[flag] = false
+		}
+	}
+}
+
+// WithFeatureFlag sets the flag to the specified value.
+func WithFeatureFlag(flag string, val bool) TestOption {
+	return func(o *testOptions) {
+		o.FeatureFlags[flag] = val
+	}
+}
+
+// WithOSSMigrations sets the migrator to the OSS migrations.
+// This effectively works _after_ all other options are passed, including WithMigrator.
+func WithOSSMigrations() TestOption {
+	return func(o *testOptions) {
+		o.MigratorFactory = func(ft featuremgmt.FeatureToggles) registry.DatabaseMigrator {
+			return migrations.ProvideOSSMigrations(ft) // the return type isn't exactly registry.DatabaseMigrator, hence the wrapper.
+		}
+	}
+}
+
+func WithMigrator(migrator registry.DatabaseMigrator) TestOption {
+	return func(o *testOptions) {
+		o.MigratorFactory = func(_ featuremgmt.FeatureToggles) registry.DatabaseMigrator {
+			return migrator
+		}
+	}
+}
+
+// WithoutMigrator explicitly opts out of migrations.
+func WithoutMigrator() TestOption {
+	return WithMigrator(nil)
+}
+
+func WithTracer(tracer tracing.Tracer, bus bus.Bus) TestOption {
+	return func(o *testOptions) {
+		o.Tracer = tracer
+		o.Bus = bus
+	}
+}
+
+func WithoutDefaultOrgAndUser() TestOption {
+	return func(o *testOptions) {
+		o.NoDefaultUserOrg = true
+	}
+}
+
+// WithCfg configures a *setting.Cfg to base the configuration upon.
+// Note that if this is set, we will modify the configuration object's [database] section.
+func WithCfg(cfg *setting.Cfg) TestOption {
+	return func(o *testOptions) {
+		o.Cfg = cfg
+	}
+}
+
+// WithTruncation enables truncating the entire database's tables after setup.
+// This is similar to the old infrastructure's behaviour.
+//
+// Most tests should just run with the data the migrations create, as they should assume a position very close to a customer's database, and customers are not going to truncate their database before updating.
+func WithTruncation() TestOption {
+	return func(o *testOptions) {
+		o.Truncate = true
+	}
+}
+
+// NewTestStore creates a new SQLStore with a test database. It is useful in parallel tests.
+// All cleanup is scheduled via the passed TestingTB; the caller does not need to do anything about it.
+// Temporary, clean databases are created for each test, and are destroyed when the test finishes.
+// When using subtests, create a new store for each subtest instead of sharing one across the entire test.
+// By default, OSS migrations are run. Enterprise migrations need to be opted into manually. Migrations can also be opted out of entirely.
+//
+// The opts are called in order. That means that a destructive option should be added last if you want it to be truly destructive.
+func NewTestStore(tb TestingTB, opts ...TestOption) *SQLStore {
+	tb.Helper()
+
+	tracer := tracing.InitializeTracerForTest()
+	options := &testOptions{
+		FeatureFlags:     make(map[string]bool),
+		Tracer:           tracer,
+		Bus:              bus.ProvideBus(tracer),
+		NoDefaultUserOrg: true,
+	}
+	WithOSSMigrations()(options) // Assign some default migrations
+	for _, opt := range opts {
+		opt(options)
+	}
+
+	features := newFeatureToggles(options.FeatureFlags)
+	testDB, err := createTemporaryDatabase(tb)
+	if err != nil {
+		tb.Fatalf("failed to create a temporary database: %v", err)
+		panic("unreachable")
+	}
+
+	cfg, err := newTestCfg(options.Cfg, features, testDB)
+	if err != nil {
+		tb.Fatalf("failed to create a test cfg: %v", err)
+		panic("unreachable")
+	}
+
+	engine, err := xorm.NewEngine(testDB.Driver, testDB.Conn)
+	if err != nil {
+		tb.Fatalf("failed to connect to temporary database: %v", err)
+		panic("unreachable")
+	}
+	tb.Cleanup(func() {
+		_ = engine.Close()
+	})
+	engine.DatabaseTZ = time.UTC
+	engine.TZLocation = time.UTC
+
+	store, err := newStore(cfg, engine, features, options.MigratorFactory(features),
+		options.Bus, options.Tracer, options.NoDefaultUserOrg || options.Truncate)
+	if err != nil {
+		tb.Fatalf("failed to create a new SQLStore: %v", err)
+		panic("unreachable")
+	}
+
+	if err := store.Migrate(false); err != nil {
+		tb.Fatalf("failed to migrate database: %v", err)
+		panic("unreachable")
+	}
+
+	if options.Truncate {
+		if err := store.dialect.TruncateDBTables(store.GetEngine()); err != nil {
+			tb.Fatalf("failed to truncate DB tables after migrations: %v", err)
+			panic("unreachable")
+		}
+	}
+
+	return store
+}
+
+func getTestDBType() string {
+	dbType := "sqlite3"
+
+	if db, present := os.LookupEnv("GRAFANA_TEST_DB"); present {
+		dbType = db
+	}
+	return dbType
+}
+
+func newFeatureToggles(toggles map[string]bool) featuremgmt.FeatureToggles {
+	spec := make([]any, 0, len(toggles)*2)
+	for flag, val := range toggles {
+		spec = append(spec, flag, val)
+	}
+	return featuremgmt.WithFeatures(spec...)
+}
+
+func newTestCfg(
+	cfg *setting.Cfg,
+	features featuremgmt.FeatureToggles,
+	testDB *testDB,
+) (*setting.Cfg, error) {
+	if cfg == nil {
+		cfg = setting.NewCfg()
+	}
+	cfg.IsFeatureToggleEnabled = features.IsEnabledGlobally
+
+	sec, err := cfg.Raw.NewSection("database")
+	if err != nil {
+		return nil, fmt.Errorf("failed to create database section in config: %w", err)
+	}
+
+	if _, err := sec.NewKey("type", getTestDBType()); err != nil {
+		return nil, fmt.Errorf("failed to set database.type: %w", err)
+	}
+	if _, err := sec.NewKey("connection_string", testDB.Conn); err != nil {
+		return nil, fmt.Errorf("failed to set database.connection_string: %w", err)
+	}
+	if _, err := sec.NewKey("path", testDB.Path); err != nil {
+		return nil, fmt.Errorf("failed to set database.path: %w", err)
+	}
+
+	return cfg, nil
+}
+
+type testDB struct {
+	Driver string
+	Conn   string
+	Path   string
+}
+
+// createTemporaryDatabase returns a connection string to a temporary database.
+// The database is created by us, and destroyed by the TestingTB cleanup function.
+// This means every database is entirely empty and isolated. Migrations are not run here.
+// If cleanup fails, the database and its data may be partially or entirely left behind.
+//
+// We assume the database credentials we are given in environment variables are those of a super user who can create databases.
+func createTemporaryDatabase(tb TestingTB) (*testDB, error) {
+	dbType := getTestDBType()
+	if dbType == "sqlite3" {
+		// SQLite doesn't have a concept of a database server, so we always create a new file with no connections required.
+		return newSQLite3DB(tb)
+	}
+
+	// On the remaining databases, we first connect to the configured credentials, create a new database, then return this new database's info as a connection string.
+	// We use databases rather than schemas as MySQL has no concept of schemas, so this aligns them more closely.
+	var driver, connString string
+	switch dbType {
+	case "sqlite3":
+		panic("unreachable; handled above")
+	case "mysql":
+		driver, connString = newMySQLConnString(env("MYSQL_DB", "grafana_tests"))
+	case "postgres":
+		driver, connString = newPostgresConnString(env("POSTGRES_DB", "grafanatest"))
+	default:
+		return nil, fmt.Errorf("unknown test db type: %s", dbType)
+	}
+
+	// We don't need the ORM here, but it's handy to connect with as we implicitly assert our driver names are correct.
+	engine, err := xorm.NewEngine(driver, connString)
+	if err != nil {
+		return nil, fmt.Errorf("failed to connect to database: %w", err)
+	}
+	defer func() {
+		// If the engine closing isn't possible to do cleanly, we don't mind.
+		_ = engine.Close()
+	}()
+
+	// The database name has to be unique amongst all tests. It is highly unlikely we will have a collision here.
+	// The database name has to be <= 64 chars long on MySQL, and <= 31 chars on Postgres.
+	id := "grafana_test_" + randomLowerHex(18)
+	_, err = engine.Exec("CREATE DATABASE " + id)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create a new database %s: %w", id, err)
+	}
+	tb.Cleanup(func() {
+		engine, err := xorm.NewEngine(driver, connString)
+		if err == nil {
+			// Clean up after ourselves at the end as well.
+			_, _ = engine.Exec("DROP DATABASE " + id)
+			_ = engine.Close()
+		}
+	})
+
+	db := &testDB{}
+	switch dbType {
+	case "mysql":
+		db.Driver, db.Conn = newMySQLConnString(id)
+	case "postgres":
+		db.Driver, db.Conn = newPostgresConnString(id)
+	default:
+		panic("unreachable; handled in the switch statement above")
+	}
+	return db, nil
+}
+
+func env(name, fallback string) string {
+	if v := os.Getenv(name); v != "" {
+		return v
+	}
+	return fallback
+}
+
+func newPostgresConnString(dbname string) (driver, connString string) {
+	return "postgres", fmt.Sprintf("user=%s password=%s host=%s port=%s dbname=%s sslmode=%s",
+		env("POSTGRES_USER", "grafanatest"),
+		env("POSTGRES_PASSWORD", "grafanatest"),
+		env("POSTGRES_HOST", "localhost"),
+		env("POSTGRES_PORT", "5432"),
+		dbname,
+		env("POSTGRES_SSL", "disable"),
+	)
+}
+
+func newMySQLConnString(dbname string) (driver, connString string) {
+	// The parseTime=true parameter is required for MySQL to parse time.Time values correctly.
+	// It converts the timezone of the time.Time to the configured timezone of the connection.
+	return "mysql", fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?collation=utf8mb4_unicode_ci&sql_mode='ANSI_QUOTES'&parseTime=true",
+		env("MYSQL_USER", "root"),
+		env("MYSQL_PASSWORD", "rootpass"),
+		env("MYSQL_HOST", "localhost"),
+		env("MYSQL_PORT", "3306"),
+		dbname,
+	)
+}
+
+func newSQLite3DB(tb TestingTB) (*testDB, error) {
+	if os.Getenv("SQLITE_INMEMORY") == "true" {
+		return &testDB{Driver: "sqlite3", Conn: "file::memory:"}, nil
+	}
+
+	tmp, err := os.CreateTemp("", "grafana-test-sqlite-*.db")
+	if err != nil {
+		return nil, fmt.Errorf("failed to create temp file: %w", err)
+	}
+	tb.Cleanup(func() {
+		// Do best efforts at cleaning up after ourselves.
+		_ = tmp.Close()
+		_ = os.Remove(tmp.Name())
+	})
+
+	// For tests, set sync=OFF for faster commits. Reference: https://www.sqlite.org/pragma.html#pragma_synchronous
+	// Sync is used in more production-y environments to avoid the database becoming corrupted. Test databases are fine to break.
+	return &testDB{
+		Driver: "sqlite3",
+		Path:   tmp.Name(),
+		Conn:   fmt.Sprintf("file:%s?cache=private&mode=rwc&_journal_mode=WAL&_synchronous=OFF", tmp.Name()),
+	}, nil
+}
+
+func randomLowerHex(length int) string {
+	buf := make([]byte, length)
+	_, err := rand.Read(buf)
+	if err != nil {
+		panic("invariant: failed to read random bytes -- crypto/rand's documentation says this cannot happen")
+	}
+
+	return hex.EncodeToString(buf)[:length]
+}

pkg/services/sqlstore/sqlstore_testinfra_test.go

@@ -0,0 +1,139 @@
+package sqlstore_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/grafana/grafana/pkg/registry"
+	"github.com/grafana/grafana/pkg/services/org"
+	"github.com/grafana/grafana/pkg/services/sqlstore"
+	"github.com/grafana/grafana/pkg/services/sqlstore/migrator"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"xorm.io/xorm"
+)
+
+// Ensure that we can get any connection at all.
+// If this test fails, it may be sensible to ignore a lot of other test failures as they may be rooted in this.
+func TestIntegrationTempDatabaseConnect(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping integration test")
+	}
+
+	sqlStore := sqlstore.NewTestStore(t, sqlstore.WithoutMigrator())
+	err := sqlStore.WithDbSession(context.Background(), func(sess *sqlstore.DBSession) error {
+		_, err := sess.Exec("SELECT 1")
+		return err
+	})
+	require.NoError(t, err, "failed to execute a SELECT 1")
+}
+
+// Ensure that migrations work on the database.
+// If this test fails, it may be sensible to ignore a lot of other test failures as they may be rooted in this.
+// This only applies OSS migrations, with no feature flags.
+func TestIntegrationTempDatabaseOSSMigrate(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping integration test")
+	}
+
+	_ = sqlstore.NewTestStore(t, sqlstore.WithOSSMigrations())
+}
+
+func TestIntegrationUniqueConstraintViolation(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping integration test")
+	}
+
+	testCases := []struct {
+		desc string
+		f    func(t *testing.T, sess *sqlstore.DBSession, dialect migrator.Dialect) error
+	}{
+		{
+			desc: "successfully detect primary key violations",
+			f: func(t *testing.T, sess *sqlstore.DBSession, dialect migrator.Dialect) error {
+				// Attempt to insert org with provided ID (primary key) twice
+				now := time.Now()
+				org := org.Org{Name: "test org primary key violation", Created: now, Updated: now, ID: 42}
+				err := sess.InsertId(&org, dialect)
+				require.NoError(t, err)
+
+				// Provide a different name to avoid unique constraint violation
+				org.Name = "test org 2"
+				return sess.InsertId(&org, dialect)
+			},
+		},
+		{
+			desc: "successfully detect unique constrain violations",
+			f: func(t *testing.T, sess *sqlstore.DBSession, dialect migrator.Dialect) error {
+				// Attempt to insert org with reserved name
+				now := time.Now()
+				org := org.Org{Name: "test org unique constrain violation", Created: now, Updated: now, ID: 43}
+				err := sess.InsertId(&org, dialect)
+				require.NoError(t, err)
+
+				// Provide a different ID to avoid primary key violation
+				org.ID = 44
+				return sess.InsertId(&org, dialect)
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			store := sqlstore.NewTestStore(t)
+			err := store.WithDbSession(context.Background(), func(sess *sqlstore.DBSession) error {
+				return tc.f(t, sess, store.GetDialect())
+			})
+			require.Error(t, err)
+			assert.True(t, store.GetDialect().IsUniqueConstraintViolation(err))
+		})
+	}
+}
+
+func TestIntegrationTruncateDatabase(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping integration test")
+	}
+
+	migrator := &truncateDatabaseSetup{}
+	store := sqlstore.NewTestStore(t, sqlstore.WithMigrator(migrator), sqlstore.WithTruncation())
+
+	var beans []*truncateBean
+	err := store.WithDbSession(context.Background(), func(sess *sqlstore.DBSession) error {
+		return sess.Find(&beans)
+	})
+	require.NoError(t, err, "could not find truncateBeans")
+
+	require.Empty(t, beans, "database should have no truncateBeans")
+}
+
+var (
+	_ registry.DatabaseMigrator = (*truncateDatabaseSetup)(nil)
+	_ migrator.CodeMigration    = (*truncateDatabaseSetup)(nil)
+)
+
+type truncateDatabaseSetup struct {
+	migrator.MigrationBase
+}
+
+func (t *truncateDatabaseSetup) AddMigration(mg *migrator.Migrator) {
+	mg.AddCreateMigration()
+	mg.AddMigration("add_to_truncate_table", t)
+}
+
+func (*truncateDatabaseSetup) SQL(dialect migrator.Dialect) string {
+	return "code migration"
+}
+
+func (t *truncateDatabaseSetup) Exec(sess *xorm.Session, migrator *migrator.Migrator) error {
+	if err := sess.CreateTable(&truncateBean{}); err != nil {
+		return err
+	}
+	_, err := sess.InsertOne(&truncateBean{1234})
+	return err
+}
+
+type truncateBean struct {
+	Value int
+}