From 80afb6c27c3f080ed17053db91f99645fd55b2b3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20Rold=C3=A1n=20Betancort?=
 <vroldanbet@authzed.com>
Date: Wed, 25 Oct 2023 17:20:50 +0100
Subject: [PATCH] do not open a new Spanner client for version checks

reuse the existing one in the datastore
---
 internal/datastore/spanner/migrations/driver.go | 14 +++++++++++---
 internal/datastore/spanner/spanner.go           | 13 ++-----------
 internal/testserver/datastore/spanner.go        |  2 +-
 pkg/cmd/migrate.go                              |  2 +-
 4 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/internal/datastore/spanner/migrations/driver.go b/internal/datastore/spanner/migrations/driver.go
index cab42310da..f3df460aea 100644
--- 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
index 6b6d0abf39..d518000c04 100644
--- 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
index 76191e0192..ca020e1937 100644
--- 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
index 62143c03aa..17c30d0c7e 100644
--- 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)
 		}