do not open a new Spanner client for version checks

reuse the existing one in the datastore
authzed · Oct 26, 2023 · 80afb6c · 80afb6c
1 parent c4c6004
commit 80afb6c
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 16 deletions.
diff --git a/internal/datastore/spanner/migrations/driver.go b/internal/datastore/spanner/migrations/driver.go
@@ -33,9 +33,7 @@ type Wrapper struct {
 }
 
 // NewSpannerDriver returns a migration driver for the given Cloud Spanner instance
-func NewSpannerDriver(database, credentialsFilePath, emulatorHost string) (*SpannerMigrationDriver, error) {
-	ctx := context.Background()
-
+func NewSpannerDriver(ctx context.Context, database, credentialsFilePath, emulatorHost string) (*SpannerMigrationDriver, error) {
 	if len(emulatorHost) > 0 {
 		err := os.Setenv(emulatorSettingKey, emulatorHost)
 		if err != nil {
@@ -57,6 +55,16 @@ func NewSpannerDriver(database, credentialsFilePath, emulatorHost string) (*Span
 	return &SpannerMigrationDriver{client, adminClient}, nil
 }
 
+// VersionProvider returns the migration version a specific spanner datastore is running at
+type VersionProvider interface {
+	Version(ctx context.Context) (string, error)
+}
+
+// NewSpannerVersionChecker returns a VersionProvider for the argument spanner.Client
+func NewSpannerVersionChecker(c *spanner.Client) VersionProvider {
+	return &SpannerMigrationDriver{c, nil}
+}
+
 func (smd *SpannerMigrationDriver) Version(ctx context.Context) (string, error) {
 	var schemaRevision string
 	iter := smd.client.Single().Read(

diff --git a/internal/datastore/spanner/spanner.go b/internal/datastore/spanner/spanner.go
@@ -236,17 +236,8 @@ func (sd spannerDatastore) ReadyState(ctx context.Context) (datastore.ReadyState
 		return datastore.ReadyState{}, fmt.Errorf("invalid head migration found for spanner: %w", err)
 	}
 
-	currentRevision, err := migrations.NewSpannerDriver(sd.client.DatabaseName(), sd.config.credentialsFilePath, sd.config.emulatorHost)
-	if err != nil {
-		return datastore.ReadyState{}, err
-	}
-	defer func() {
-		if err := currentRevision.Close(ctx); err != nil {
-			log.Error().Err(err).Msg("failed to close current revision in Datastore.ReadyState")
-		}
-	}()
-
-	version, err := currentRevision.Version(ctx)
+	checker := migrations.NewSpannerVersionChecker(sd.client)
+	version, err := checker.Version(ctx)
 	if err != nil {
 		return datastore.ReadyState{}, err
 	}

diff --git a/internal/testserver/datastore/spanner.go b/internal/testserver/datastore/spanner.go
@@ -138,7 +138,7 @@ func (b *spannerTest) NewDatabase(t testing.TB) string {
 func (b *spannerTest) NewDatastore(t testing.TB, initFunc InitFunc) datastore.Datastore {
 	db := b.NewDatabase(t)
 
-	migrationDriver, err := migrations.NewSpannerDriver(db, "", os.Getenv("SPANNER_EMULATOR_HOST"))
+	migrationDriver, err := migrations.NewSpannerDriver(context.Background(), db, "", os.Getenv("SPANNER_EMULATOR_HOST"))
 	require.NoError(t, err)
 
 	err = migrations.SpannerMigrations.Run(context.Background(), migrationDriver, b.targetMigration, migrate.LiveRun)

diff --git a/pkg/cmd/migrate.go b/pkg/cmd/migrate.go
@@ -74,7 +74,7 @@ func migrateRun(cmd *cobra.Command, args []string) error {
 		if err != nil {
 			log.Ctx(cmd.Context()).Fatal().Err(err).Msg("unable to get spanner emulator host")
 		}
-		migrationDriver, err := spannermigrations.NewSpannerDriver(dbURL, credFile, emulatorHost)
+		migrationDriver, err := spannermigrations.NewSpannerDriver(cmd.Context(), dbURL, credFile, emulatorHost)
 		if err != nil {
 			return fmt.Errorf("unable to create migration driver for %s: %w", datastoreEngine, err)
 		}