From 8db050af6b1ad43244ec60a815255ff019e9f0e1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Peter=20=C5=A0tibran=C3=BD?= Date: Thu, 20 Mar 2025 16:50:50 +0100 Subject: [PATCH] Use in-memory sequence number generator when running integration tests against Spanner emulator. (#102522) --- .../sqlstore/migrator/spanner_dialect.go | 29 +------------- pkg/util/xorm/dialect_spanner.go | 23 +++++++++++ pkg/util/xorm/engine.go | 5 +-- pkg/util/xorm/sequence_inmem.go | 39 +++++++++++++++++++ pkg/util/xorm/xorm.go | 23 +++++++++-- 5 files changed, 85 insertions(+), 34 deletions(-) create mode 100644 pkg/util/xorm/sequence_inmem.go diff --git a/pkg/services/sqlstore/migrator/spanner_dialect.go b/pkg/services/sqlstore/migrator/spanner_dialect.go index c66032367ff..950aabf2f2a 100644 --- a/pkg/services/sqlstore/migrator/spanner_dialect.go +++ b/pkg/services/sqlstore/migrator/spanner_dialect.go @@ -7,7 +7,6 @@ import ( "encoding/json" "errors" "fmt" - "strconv" "time" "cloud.google.com/go/spanner" @@ -15,12 +14,10 @@ import ( "github.com/googleapis/gax-go/v2" spannerdriver "github.com/googleapis/go-sql-spanner" "github.com/grafana/dskit/concurrency" - "google.golang.org/api/option" - "google.golang.org/grpc" "google.golang.org/grpc/codes" - "google.golang.org/grpc/credentials/insecure" "xorm.io/core" + spannerext "github.com/grafana/grafana/pkg/extensions/spanner" "xorm.io/xorm" _ "embed" @@ -294,7 +291,7 @@ func (s *SpannerDialect) executeDDLStatements(ctx context.Context, engine *xorm. return err } - opts := SpannerConnectorConfigToClientOptions(cfg) + opts := spannerext.SpannerConnectorConfigToClientOptions(cfg) databaseAdminClient, err := database.NewDatabaseAdminClient(ctx, opts...) if err != nil { @@ -319,28 +316,6 @@ func (s *SpannerDialect) executeDDLStatements(ctx context.Context, engine *xorm. return nil } -// SpannerConnectorConfigToClientOptions is adapted from https://github.com/googleapis/go-sql-spanner/blob/main/driver.go#L341-L477, from version 1.11.1. -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 (s *SpannerDialect) UnionDistinct() string { return "UNION DISTINCT" } diff --git a/pkg/util/xorm/dialect_spanner.go b/pkg/util/xorm/dialect_spanner.go index 4269107cbb6..df4b0e3965e 100644 --- a/pkg/util/xorm/dialect_spanner.go +++ b/pkg/util/xorm/dialect_spanner.go @@ -9,7 +9,10 @@ import ( "strings" _ "github.com/googleapis/go-sql-spanner" + spannerdriver "github.com/googleapis/go-sql-spanner" "xorm.io/core" + + spannerext "github.com/grafana/grafana/pkg/extensions/spanner" ) func init() { @@ -370,3 +373,23 @@ func (s *spanner) GetIndexes(tableName string) (map[string]*core.Index, error) { } return indexes, res.Err() } + +func (s *spanner) CreateSequenceGenerator(db *sql.DB) (SequenceGenerator, error) { + dsn := s.DataSourceName() + connectorConfig, err := spannerdriver.ExtractConnectorConfig(dsn) + if err != nil { + return nil, err + } + + if spannerext.UsePlainText(connectorConfig) { + // Plain-text means we're either using spannertest or Spanner emulator. + // Switch to fake in-memory sequence number generator in that case. + // + // Using database-based sequence generator doesn't work with emulator, as emulator + // only supports single transaction. If there is already another transaction started + // generating new ID via database-based sequence generator would always fail. + return newInMemSequenceGenerator(), nil + } + + return newSequenceGenerator(db), nil +} diff --git a/pkg/util/xorm/engine.go b/pkg/util/xorm/engine.go index 8db50041ae2..4d29facbd77 100644 --- a/pkg/util/xorm/engine.go +++ b/pkg/util/xorm/engine.go @@ -44,7 +44,7 @@ type Engine struct { tagHandlers map[string]tagHandler defaultContext context.Context - sequenceGenerator *sequenceGenerator // If not nil, this generator is used to generate auto-increment values for inserts. + sequenceGenerator SequenceGenerator // If not nil, this generator is used to generate auto-increment values for inserts. } // CondDeleted returns the conditions whether a record is soft deleted. @@ -239,9 +239,6 @@ func (engine *Engine) NewSession() *Session { // Close the engine func (engine *Engine) Close() error { - if engine.sequenceGenerator != nil { - engine.sequenceGenerator.close() - } return engine.db.Close() } diff --git a/pkg/util/xorm/sequence_inmem.go b/pkg/util/xorm/sequence_inmem.go new file mode 100644 index 00000000000..e412e687664 --- /dev/null +++ b/pkg/util/xorm/sequence_inmem.go @@ -0,0 +1,39 @@ +package xorm + +import ( + "context" + "fmt" + "math/rand/v2" + "sync" +) + +type inMemSequenceGenerator struct { + sequencesMu sync.Mutex + nextValues map[string]int +} + +func newInMemSequenceGenerator() *inMemSequenceGenerator { + return &inMemSequenceGenerator{ + nextValues: make(map[string]int), + } +} + +func (g *inMemSequenceGenerator) Next(_ context.Context, table, column string) (int64, error) { + if table == "migration_log" { + // Don't use sequential IDs for migration log entries, as we don't clean up migration_log table between tests, + // so restarting the sequence can lead to conflicting IDs. + return rand.Int64(), nil + } + + key := fmt.Sprintf("%s:%s", table, column) + + g.sequencesMu.Lock() + defer g.sequencesMu.Unlock() + + seq, ok := g.nextValues[key] + if !ok { + seq = 1 + } + g.nextValues[key] = seq + 1 + return int64(seq), nil +} diff --git a/pkg/util/xorm/xorm.go b/pkg/util/xorm/xorm.go index 25ce8bd952f..109a0bca16a 100644 --- a/pkg/util/xorm/xorm.go +++ b/pkg/util/xorm/xorm.go @@ -9,6 +9,7 @@ package xorm import ( "context" + "database/sql" "fmt" "os" "reflect" @@ -22,6 +23,8 @@ import ( const ( // Version show the xorm's version Version string = "0.8.0.1015" + + Spanner = "spanner" ) func regDrvsNDialects() bool { @@ -97,7 +100,7 @@ func NewEngine(driverName string, dataSourceName string) (*Engine, error) { switch uri.DbType { case core.SQLITE: engine.DatabaseTZ = time.UTC - case "spanner": + case Spanner: engine.DatabaseTZ = time.UTC // We need to specify "Z" to indicate that timestamp is in UTC. // Otherwise Spanner uses default America/Los_Angeles timezone. @@ -114,9 +117,23 @@ func NewEngine(driverName string, dataSourceName string) (*Engine, error) { runtime.SetFinalizer(engine, close) - if dialect.DBType() == "spanner" { - engine.sequenceGenerator = newSequenceGenerator(db.DB) + if ext, ok := dialect.(DialectExt); ok { + engine.sequenceGenerator, err = ext.CreateSequenceGenerator(db.DB) + if err != nil { + return nil, fmt.Errorf("failed to create sequence generator: %w", err) + } } return engine, nil } + +type SequenceGenerator interface { + Next(ctx context.Context, table, column string) (int64, error) +} + +type DialectExt interface { + core.Dialect + + // CreateSequenceGenerator returns optional generator used to create AUTOINCREMENT ids for inserts. + CreateSequenceGenerator(db *sql.DB) (SequenceGenerator, error) +}