Merge pull request #1212 from ecordell/spannerclose

spanner: close row iterators when done
authzed · Mar 20, 2023 · 03896f6 · 03896f6
2 parents cf1004c + d4e47b9
commit 03896f6
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 34 deletions.
diff --git a/internal/datastore/spanner/migrations/driver.go b/internal/datastore/spanner/migrations/driver.go
@@ -58,26 +58,22 @@ func NewSpannerDriver(database, credentialsFilePath, emulatorHost string) (*Span
 }
 
 func (smd *SpannerMigrationDriver) Version(ctx context.Context) (string, error) {
-	rows := smd.client.Single().Read(
+	var schemaRevision string
+	if err := smd.client.Single().Read(
 		ctx,
 		tableSchemaVersion,
 		spanner.AllKeys(),
 		[]string{colVersionNum},
-	)
-	row, err := rows.Next()
-	if err != nil {
+	).Do(func(r *spanner.Row) error {
+		return r.Columns(&schemaRevision)
+	}); err != nil {
 		if spanner.ErrCode(err) == codes.NotFound {
 			// There is no schema table, empty database
 			return "", nil
 		}
 		return "", err
 	}
 
-	var schemaRevision string
-	if err := row.Columns(&schemaRevision); err != nil {
-		return "", err
-	}
-
 	return schemaRevision, nil
 }
 

diff --git a/internal/datastore/spanner/revisions.go b/internal/datastore/spanner/revisions.go
@@ -29,16 +29,10 @@ func (sd spannerDatastore) now(ctx context.Context) (time.Time, error) {
 	ctx, span := tracer.Start(ctx, "now")
 	defer span.End()
 
-	iter := sd.client.Single().Query(ctx, spanner.NewStatement("SELECT CURRENT_TIMESTAMP()"))
-	defer iter.Stop()
-
-	row, err := iter.Next()
-	if err != nil {
-		return time.Time{}, err
-	}
-
 	var timestamp time.Time
-	if err := row.Columns(&timestamp); err != nil {
+	if err := sd.client.Single().Query(ctx, spanner.NewStatement("SELECT CURRENT_TIMESTAMP()")).Do(func(r *spanner.Row) error {
+		return r.Columns(&timestamp)
+	}); err != nil {
 		return time.Time{}, err
 	}
 

diff --git a/internal/datastore/spanner/stats.go b/internal/datastore/spanner/stats.go
@@ -15,19 +15,15 @@ import (
 var queryRelationshipEstimate = fmt.Sprintf("SELECT SUM(%s) FROM %s", colCount, tableCounters)
 
 func (sd spannerDatastore) Statistics(ctx context.Context) (datastore.Stats, error) {
-	idRows := sd.client.Single().Read(
+	var uniqueID string
+	if err := sd.client.Single().Read(
 		context.Background(),
 		tableMetadata,
 		spanner.AllKeys(),
 		[]string{colUniqueID},
-	)
-	idRow, err := idRows.Next()
-	if err != nil {
-		return datastore.Stats{}, fmt.Errorf("unable to read metadata table: %w", err)
-	}
-
-	var uniqueID string
-	if err := idRow.Columns(&uniqueID); err != nil {
+	).Do(func(r *spanner.Row) error {
+		return r.Columns(&uniqueID)
+	}); err != nil {
 		return datastore.Stats{}, fmt.Errorf("unable to read unique ID: %w", err)
 	}
 
@@ -44,14 +40,10 @@ func (sd spannerDatastore) Statistics(ctx context.Context) (datastore.Stats, err
 	}
 
 	var estimate spanner.NullInt64
-	countRows := sd.client.Single().Query(ctx, spanner.Statement{SQL: queryRelationshipEstimate})
-	countRow, err := countRows.Next()
-	if err != nil && spanner.ErrCode(err) != codes.NotFound {
+	if err := sd.client.Single().Query(ctx, spanner.Statement{SQL: queryRelationshipEstimate}).Do(func(r *spanner.Row) error {
+		return r.Columns(&estimate)
+	}); err != nil {
 		return datastore.Stats{}, fmt.Errorf("unable to read row counts: %w", err)
-	} else if err == nil {
-		if err := countRow.Columns(&estimate); err != nil {
-			return datastore.Stats{}, fmt.Errorf("unable to decode row count: %w", err)
-		}
 	}
 
 	return datastore.Stats{