From 68c6e88b0045e263252b1fc0014cf765893cc605 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Peter=20=C5=A0tibran=C3=BD?= Date: Mon, 7 Apr 2025 12:15:16 +0200 Subject: [PATCH] Support creating temporary Spanner databases for integration tests. (#103347) * Support creating temporary Spanner databases for integration tests. --- pkg/services/sqlstore/sqlstore_testinfra.go | 146 +++++++++++++++++- .../sqlstore/sqlstore_testinfra_test.go | 7 +- 2 files changed, 146 insertions(+), 7 deletions(-) diff --git a/pkg/services/sqlstore/sqlstore_testinfra.go b/pkg/services/sqlstore/sqlstore_testinfra.go index 0a5c4b5db4f..97ea47fb07e 100644 --- a/pkg/services/sqlstore/sqlstore_testinfra.go +++ b/pkg/services/sqlstore/sqlstore_testinfra.go @@ -4,12 +4,23 @@ package sqlstore import ( + "context" "crypto/rand" "encoding/hex" "fmt" "os" + "strconv" + "testing" "time" + database "cloud.google.com/go/spanner/admin/database/apiv1" + "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" + "cloud.google.com/go/spanner/spannertest" + spannerdriver "github.com/googleapis/go-sql-spanner" + "google.golang.org/api/option" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + "github.com/grafana/grafana/pkg/bus" "github.com/grafana/grafana/pkg/infra/tracing" "github.com/grafana/grafana/pkg/registry" @@ -19,7 +30,7 @@ import ( "xorm.io/xorm" ) -// testingTB is an interface that is implemented by *testing.T and *testing.B. Similar to testing.TB. +// 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() @@ -27,8 +38,14 @@ type TestingTB interface { 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) + // Logf formats and logs its arguments. See also (*testing.T).Logf. + Logf(format string, args ...any) + // Context returns a context that is canceled just before Cleanup-registered functions are called. See also (*testing.T).Context. + Context() context.Context } +var _ TestingTB = (testing.TB)(nil) + type testOptions struct { FeatureFlags map[string]bool MigratorFactory func(featuremgmt.FeatureToggles) registry.DatabaseMigrator @@ -252,6 +269,9 @@ func createTemporaryDatabase(tb TestingTB) (*testDB, error) { // SQLite doesn't have a concept of a database server, so we always create a new file with no connections required. return newSQLite3DB(tb) } + if dbType == "spanner" { + return newSpannerDB(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. @@ -277,9 +297,7 @@ func createTemporaryDatabase(tb TestingTB) (*testDB, error) { _ = 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) + id := generateDatabaseName() _, err = engine.Exec("CREATE DATABASE " + id) if err != nil { return nil, fmt.Errorf("failed to create a new database %s: %w", id, err) @@ -305,6 +323,126 @@ func createTemporaryDatabase(tb TestingTB) (*testDB, error) { return db, nil } +func generateDatabaseName() string { + // 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. + // Database ID length on Spanner must be between 2 and 30 characters. (https://cloud.google.com/spanner/quotas#database-limits) + return "grafana_test_" + randomLowerHex(17) +} + +func newSpannerDB(tb TestingTB) (*testDB, error) { + // See https://github.com/googleapis/go-sql-spanner/blob/main/driver.go#L56-L81 for connection string options. + spannerDB := env("SPANNER_DB", "emulator") + if spannerDB == "spannertest" { + // Start new in-memory spannertest instance. This is mostly useless for our tests + // (spannertest doesn't support many things that we use), but added for completion. + // Each spannertest instance is a separate db. + srv, err := spannertest.NewServer("localhost:0") + if err != nil { + return nil, err + } + tb.Cleanup(srv.Close) + + return &testDB{ + Driver: "spanner", + Conn: fmt.Sprintf("%s/projects/grafanatest/instances/grafanatest/databases/grafanatest;usePlainText=true", srv.Addr), + }, nil + } + + conn := spannerDB + if spannerDB == "emulator" { + host := env("SPANNER_EMULATOR_HOST", "localhost:9010") + conn = fmt.Sprintf("%s/projects/grafanatest/instances/grafanatest/databases/grafanatest;usePlainText=true", host) + } + + cfg, err := spannerdriver.ExtractConnectorConfig(conn) + if err != nil { + return nil, err + } + + clientOptions := spannerConnectorConfigToClientOptions(cfg) + + dbname := generateDatabaseName() + fullDbName := fmt.Sprintf("projects/%s/instances/%s/databases/%s", cfg.Project, cfg.Instance, dbname) + dbCreated := false + + databaseAdminClient, err := database.NewDatabaseAdminClient(tb.Context(), clientOptions...) + if err != nil { + return nil, fmt.Errorf("failed to create database admin client: %v", err) + } + tb.Cleanup(func() { + if dbCreated { + // Drop database in the cleanup. + // Can't use tb.Context() here, since that is canceled before calling Cleanup functions. + err := databaseAdminClient.DropDatabase(context.Background(), &databasepb.DropDatabaseRequest{ + Database: fullDbName, + }) + if err != nil { + tb.Logf("Failed to drop Spanner database %s due to error %v", fullDbName, err) + } else { + tb.Logf("Dropped temporary Spanner database %s", fullDbName) + } + } + + _ = databaseAdminClient.Close() + }) + + op, err := databaseAdminClient.CreateDatabase(tb.Context(), &databasepb.CreateDatabaseRequest{ + Parent: fmt.Sprintf("projects/%s/instances/%s", cfg.Project, cfg.Instance), + CreateStatement: fmt.Sprintf("CREATE DATABASE `%s`", dbname), + DatabaseDialect: databasepb.DatabaseDialect_GOOGLE_STANDARD_SQL, + }) + if err != nil { + return nil, fmt.Errorf("failed to create database: %v", err) + } + _, err = op.Wait(tb.Context()) + if err != nil { + return nil, fmt.Errorf("failed to create database: %v", err) + } + tb.Logf("Created temporary Spanner database %s", fullDbName) + + dbCreated = true + + // Rebuild connection string, but change database to ID of just-created database. + // Example: `localhost:9010/projects/test-project/instances/test-instance/databases/test-database;usePlainText=true;disableRouteToLeader=true;enableEndToEndTracing=true` + connString := "" + if cfg.Host != "" { + connString = fmt.Sprintf("%s/", cfg.Host) + } + // Use new DB name instead of cfg.Database. + connString = connString + fmt.Sprintf("projects/%s/instances/%s/databases/%s", cfg.Project, cfg.Instance, dbname) + for k, v := range cfg.Params { + connString = connString + fmt.Sprintf(";%s=%s", k, v) + } + + return &testDB{ + Driver: "spanner", + Conn: connString, + }, nil +} + +// This is same code as xorm.SpannerConnectorConfigToClientOptions, but we cannot use that because it's under "enterprise" build tag. +func spannerConnectorConfigToClientOptions(connectorConfig spannerdriver.ConnectorConfig) []option.ClientOption { + var opts []option.ClientOption + if connectorConfig.Host != "" { + opts = append(opts, option.WithEndpoint(connectorConfig.Host)) + } + if strval, ok := connectorConfig.Params["credentials"]; ok { + opts = append(opts, option.WithCredentialsFile(strval)) + } + if strval, ok := connectorConfig.Params["credentialsjson"]; ok { + opts = append(opts, option.WithCredentialsJSON([]byte(strval))) + } + if strval, ok := connectorConfig.Params["useplaintext"]; ok { + if val, err := strconv.ParseBool(strval); err == nil && val { + opts = append(opts, + option.WithGRPCDialOption(grpc.WithTransportCredentials(insecure.NewCredentials())), + option.WithoutAuthentication()) + } + } + return opts +} + func env(name, fallback string) string { if v := os.Getenv(name); v != "" { return v diff --git a/pkg/services/sqlstore/sqlstore_testinfra_test.go b/pkg/services/sqlstore/sqlstore_testinfra_test.go index 2bd3b55e89e..623865ed903 100644 --- a/pkg/services/sqlstore/sqlstore_testinfra_test.go +++ b/pkg/services/sqlstore/sqlstore_testinfra_test.go @@ -5,12 +5,13 @@ import ( "testing" "time" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "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" ) @@ -23,7 +24,7 @@ func TestIntegrationTempDatabaseConnect(t *testing.T) { sqlStore := sqlstore.NewTestStore(t, sqlstore.WithoutMigrator()) err := sqlStore.WithDbSession(context.Background(), func(sess *sqlstore.DBSession) error { - _, err := sess.Exec("SELECT 1") + _, err := sess.Query("SELECT 1") return err }) require.NoError(t, err, "failed to execute a SELECT 1")