Use in-memory sequence number generator when running integration tests against Spanner emulator. (#102522)

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"
 }

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
+}

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()
 }
 

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
+}

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)
+}