Propagate context in Version

internal/datastore/crdb/crdb.go

@@ -178,7 +178,7 @@ func (cds *crdbDatastore) IsReady(ctx context.Context) (bool, error) {
 	}
 	defer currentRevision.Close()
 
-	version, err := currentRevision.Version()
+	version, err := currentRevision.Version(ctx)
 	if err != nil {
 		return false, err
 	}

internal/datastore/crdb/migrations/driver.go

@@ -46,10 +46,10 @@ func NewCRDBDriver(url string) (*CRDBDriver, error) {
 
 // Version returns the version of the schema to which the connected database
 // has been migrated.
-func (apd *CRDBDriver) Version() (string, error) {
+func (apd *CRDBDriver) Version(ctx context.Context) (string, error) {
 	var loaded string
 
-	if err := apd.db.QueryRow(context.Background(), queryLoadVersion).Scan(&loaded); err != nil {
+	if err := apd.db.QueryRow(ctx, queryLoadVersion).Scan(&loaded); err != nil {
 		var pgErr *pgconn.PgError
 		if errors.As(err, &pgErr) && pgErr.Code == postgresMissingTableErrorCode {
 			return "", nil

internal/datastore/mysql/datastore.go

@@ -460,7 +460,7 @@ func (mds *Datastore) IsReady(ctx context.Context) (bool, error) {
 		return false, err
 	}
 
-	currentMigrationRevision, err := mds.driver.Version()
+	currentMigrationRevision, err := mds.driver.Version(ctx)
 	if err != nil {
 		return false, err
 	}

internal/datastore/mysql/datastore_test.go

@@ -642,14 +642,14 @@ func TestMySQLMigrations(t *testing.T) {
 	db := datastoreDB(t, false)
 	migrationDriver := migrations.NewMySQLDriverFromDB(db, "")
 
-	version, err := migrationDriver.Version()
+	version, err := migrationDriver.Version(context.Background())
 	req.NoError(err)
 	req.Equal("", version)
 
 	err = migrations.Manager.Run(context.Background(), migrationDriver, migrate.Head, migrate.LiveRun)
 	req.NoError(err)
 
-	version, err = migrationDriver.Version()
+	version, err = migrationDriver.Version(context.Background())
 	req.NoError(err)
 
 	headVersion, err := migrations.Manager.HeadRevision()
@@ -664,14 +664,14 @@ func TestMySQLMigrationsWithPrefix(t *testing.T) {
 	db := datastoreDB(t, false)
 	migrationDriver := migrations.NewMySQLDriverFromDB(db, prefix)
 
-	version, err := migrationDriver.Version()
+	version, err := migrationDriver.Version(context.Background())
 	req.NoError(err)
 	req.Equal("", version)
 
 	err = migrations.Manager.Run(context.Background(), migrationDriver, migrate.Head, migrate.LiveRun)
 	req.NoError(err)
 
-	version, err = migrationDriver.Version()
+	version, err = migrationDriver.Version(context.Background())
 	req.NoError(err)
 
 	headVersion, err := migrations.Manager.HeadRevision()

internal/datastore/mysql/migrations/driver.go

@@ -67,21 +67,21 @@ func columnNameToRevision(columnName string) (string, bool) {
 
 // Version returns the version of the schema to which the connected database
 // has been migrated.
-func (driver *MySQLDriver) Version() (string, error) {
+func (driver *MySQLDriver) Version(ctx context.Context) (string, error) {
 	query, args, err := sb.Select("*").From(driver.migrationVersion()).ToSql()
 	if err != nil {
 		return "", fmt.Errorf("unable to load driver migration revision: %w", err)
 	}
 
-	rows, err := driver.db.Query(query, args...)
+	rows, err := driver.db.QueryContext(ctx, query, args...)
 	if err != nil {
 		var mysqlError *sqlDriver.MySQLError
 		if errors.As(err, &mysqlError) && mysqlError.Number == mysqlMissingTableErrorNumber {
 			return "", nil
 		}
 		return "", fmt.Errorf("unable to load driver migration revision: %w", err)
 	}
-	defer LogOnError(context.Background(), rows.Close)
+	defer LogOnError(ctx, rows.Close)
 	if rows.Err() != nil {
 		return "", fmt.Errorf("unable to load driver migration revision: %w", err)
 	}

internal/datastore/postgres/migrations/driver.go

@@ -38,10 +38,10 @@ func NewAlembicPostgresDriver(url string) (*AlembicPostgresDriver, error) {
 
 // Version returns the version of the schema to which the connected database
 // has been migrated.
-func (apd *AlembicPostgresDriver) Version() (string, error) {
+func (apd *AlembicPostgresDriver) Version(ctx context.Context) (string, error) {
 	var loaded string
 
-	if err := apd.db.QueryRowx("SELECT version_num from alembic_version").Scan(&loaded); err != nil {
+	if err := apd.db.QueryRowxContext(ctx, "SELECT version_num from alembic_version").Scan(&loaded); err != nil {
 		var pqErr *pq.Error
 		if errors.As(err, &pqErr) && pqErr.Code == postgresMissingTableErrorCode {
 			return "", nil

internal/datastore/postgres/postgres.go

@@ -422,7 +422,7 @@ func (pgd *pgDatastore) IsReady(ctx context.Context) (bool, error) {
 	}
 	defer currentRevision.Close()
 
-	version, err := currentRevision.Version()
+	version, err := currentRevision.Version(ctx)
 	if err != nil {
 		return false, err
 	}

internal/datastore/spanner/migrations/driver.go

@@ -48,9 +48,9 @@ func NewSpannerDriver(database, credentialsFilePath, emulatorHost string) (Spann
 	return SpannerMigrationDriver{client, adminClient}, nil
 }
 
-func (smd SpannerMigrationDriver) Version() (string, error) {
+func (smd SpannerMigrationDriver) Version(ctx context.Context) (string, error) {
 	rows := smd.client.Single().Read(
-		context.Background(),
+		ctx,
 		tableSchemaVersion,
 		spanner.AllKeys(),
 		[]string{colVersionNum},

internal/datastore/spanner/spanner.go

@@ -114,7 +114,7 @@ func (sd spannerDatastore) IsReady(ctx context.Context) (bool, error) {
 	}
 	defer currentRevision.Close()
 
-	version, err := currentRevision.Version()
+	version, err := currentRevision.Version(ctx)
 	if err != nil {
 		return false, err
 	}

pkg/migrate/migrate.go

@@ -28,7 +28,7 @@ type Driver interface {
 	// Version returns the current version of the schema in the backing datastore.
 	// If the datastore is brand new, version should return the empty string without
 	// an error.
-	Version() (string, error)
+	Version(ctx context.Context) (string, error)
 
 	// WriteVersion records the newly migrated version to the backing datastore.
 	WriteVersion(ctx context.Context, version, replaced string) error
@@ -83,7 +83,7 @@ func (m *Manager) Register(version, replaces string, up interface{}) error {
 // Run will actually perform the necessary migrations to bring the backing datastore
 // from its current revision to the specified revision.
 func (m *Manager) Run(ctx context.Context, driver Driver, throughRevision string, dryRun RunType) error {
-	starting, err := driver.Version()
+	starting, err := driver.Version(ctx)
 	if err != nil {
 		return fmt.Errorf("unable to compute target revision: %w", err)
 	}
@@ -113,7 +113,7 @@ func (m *Manager) Run(ctx context.Context, driver Driver, throughRevision string
 	if !dryRun {
 		for _, migrationToRun := range toRun {
 			// Double check that the current version reported is the one we expect
-			currentVersion, err := driver.Version()
+			currentVersion, err := driver.Version(ctx)
 			if err != nil {
 				return fmt.Errorf("unable to load version from driver: %w", err)
 			}

pkg/migrate/migrate_test.go

@@ -9,7 +9,7 @@ import (
 
 type fakeDriver struct{}
 
-func (*fakeDriver) Version() (string, error) {
+func (*fakeDriver) Version(ctx context.Context) (string, error) {
 	return "", nil
 }