Context DB txn timeout, and interface to avoid connection leak during migration

pkg/dbsql/database.go

@@ -30,6 +30,7 @@ import (
 	"github.com/hyperledger/firefly-common/pkg/i18n"
 	"github.com/hyperledger/firefly-common/pkg/log"
 
+	"github.com/golang-migrate/migrate/v4/database"
 	// Import migrate file source
 	_ "github.com/golang-migrate/migrate/v4/source/file"
 )
@@ -139,8 +140,17 @@ func (s *Database) RunAsGroup(ctx context.Context, fn func(ctx context.Context)
 	return s.CommitTx(ctx, tx, false /* we _are_ the auto-committer */)
 }
 
-func (s *Database) applyDBMigrations(ctx context.Context, config config.Section, provider Provider) error {
-	driver, err := provider.GetMigrationDriver(s.db)
+func (s *Database) applyDBMigrations(ctx context.Context, config config.Section, provider Provider) (err error) {
+	var driver database.Driver
+	providerClosable, isClosable := provider.(ProviderCloseableMigrationDriver)
+	if isClosable {
+		driver, err = providerClosable.GetMigrationDriverClosable(s.db)
+		defer func() {
+			_ = driver.Close()
+		}()
+	} else {
+		driver, err = provider.GetMigrationDriver(s.db)
+	}
 	if err == nil {
 		fileURL := "file://" + config.GetString(SQLConfMigrationsDirectory)
 		log.L(ctx).Infof("Running migrations in: %s", fileURL)
@@ -184,7 +194,7 @@ func (s *Database) BeginOrUseTx(ctx context.Context) (ctx1 context.Context, tx *
 	ctx1 = log.WithLogger(ctx, l)
 	before := time.Now()
 	l.Tracef("SQL-> begin")
-	sqlTX, err := s.db.Begin()
+	sqlTX, err := s.db.BeginTx(ctx /* transaction should auto-rollback on context cancel */, nil /* default options */)
 	if err != nil {
 		return ctx1, nil, false, i18n.WrapError(ctx1, err, i18n.MsgDBBeginFailed)
 	}

pkg/dbsql/mock_provider.go

@@ -40,6 +40,10 @@ type MockProvider struct {
 	mmg    *dbmigratemocks.Driver
 }
 
+type MockProviderConnScoped struct {
+	MockProvider
+}
+
 type MockProviderConfig struct {
 	FakePSQLInsert             bool
 	OpenError                  error
@@ -108,3 +112,7 @@ func (mp *MockProvider) Open(_ string) (*sql.DB, error) {
 func (mp *MockProvider) GetMigrationDriver(_ *sql.DB) (migratedb.Driver, error) {
 	return mp.mmg, mp.GetMigrationDriverError
 }
+
+func (mp *MockProviderConnScoped) GetMigrationDriverConn(_ *sql.DB) (migratedb.Driver, error) {
+	return mp.mmg, mp.GetMigrationDriverError
+}

pkg/dbsql/provider.go

@@ -59,7 +59,7 @@ type Provider interface {
 	// Open creates the DB instances
 	Open(url string) (*sql.DB, error)
 
-	// GetDriver returns the driver implementation
+	// GetDriver returns the driver implementation - preferred if supported to implement GetMigrationDriverConn so the connection can be cleaned up (cleaning this driver up closes the whole DB)
 	GetMigrationDriver(*sql.DB) (migratedb.Driver, error)
 
 	// Features returns database specific configuration switches
@@ -68,3 +68,9 @@ type Provider interface {
 	// ApplyInsertQueryCustomizations updates the INSERT query for returning the Sequence, and returns whether it needs to be run as a query to return the Sequence field
 	ApplyInsertQueryCustomizations(insert sq.InsertBuilder, requestConflictEmptyResult bool) (updatedInsert sq.InsertBuilder, runAsQuery bool)
 }
+
+// Implementing this interface allows cleanup of the connection used during migration
+type ProviderCloseableMigrationDriver interface {
+	// Returns the driver implementation that is safe to close. Specifically this means the WithConnection() use in migration, rather than WithInstance().
+	GetMigrationDriverClosable(*sql.DB) (migratedb.Driver, error)
+}