Run migration and save migration version in single transaction

Signed-off-by: Amelia Downs <adowns@vmware.com>
Signed-off-by: Maria Shaldybin <mariash@vmware.com>

cmd/bbs/fixtures/9999999999_sql_test_migration.go.bs

@@ -23,15 +23,11 @@ func NewSQLTestMigration() migration.Migration {
 
 func (*SQLTestMigration) SetCryptor(cryptor encryption.Cryptor) {}
 
-func (t *SQLTestMigration) SetRawSQLDB(rawSQLDB *sql.DB) {
-	t.rawSQLDB = rawSQLDB
-}
-
 func (*SQLTestMigration) SetClock(clock.Clock) {}
 func (*SQLTestMigration) SetDBFlavor(string)   {}
 
-func (t *SQLTestMigration) Up(logger lager.Logger) error {
-	_, err := t.rawSQLDB.Exec(`CREATE TABLE IF NOT EXISTS sweet_table (
+func (t *SQLTestMigration) Up(tx *sql.Tx, logger lager.Logger) error {
+	_, err := tx.Exec(`CREATE TABLE IF NOT EXISTS sweet_table (
 		something VARCHAR(255) PRIMARY KEY,
 		something_else INT DEFAULT 0
 	);`)

db/migrations/1461790966_init_sql.go

@@ -18,7 +18,6 @@ func init() {
 type InitSQL struct {
 	serializer format.Serializer
 	clock      clock.Clock
-	rawSQLDB   *sql.DB
 	dbFlavor   string
 }
 
@@ -38,61 +37,57 @@ func (e *InitSQL) SetCryptor(cryptor encryption.Cryptor) {
 	e.serializer = format.NewSerializer(cryptor)
 }
 
-func (e *InitSQL) SetRawSQLDB(db *sql.DB) {
-	e.rawSQLDB = db
-}
-
 func (e *InitSQL) SetClock(c clock.Clock)    { e.clock = c }
 func (e *InitSQL) SetDBFlavor(flavor string) { e.dbFlavor = flavor }
 
-func (e *InitSQL) Up(logger lager.Logger) error {
+func (e *InitSQL) Up(tx *sql.Tx, logger lager.Logger) error {
 	logger = logger.Session("init-sql")
 	logger.Info("truncating-tables")
 
 	// Ignore the error as the tables may not exist
-	_ = dropTables(e.rawSQLDB)
+	_ = e.dropTables(tx)
 
-	err := createTables(logger, e.rawSQLDB, e.dbFlavor)
+	err := e.createTables(tx, logger)
 	if err != nil {
 		return err
 	}
 
-	err = createIndices(logger, e.rawSQLDB)
+	err = e.createIndices(tx, logger)
 	if err != nil {
 		return err
 	}
 
 	return nil
 }
 
-func dropTables(db *sql.DB) error {
+func (e *InitSQL) dropTables(tx *sql.Tx) error {
 	tableNames := []string{
 		"domains",
 		"tasks",
 		"desired_lrps",
 		"actual_lrps",
 	}
 	for _, tableName := range tableNames {
-		_, err := db.Exec("DROP TABLE IF EXISTS " + tableName)
+		_, err := tx.Exec("DROP TABLE IF EXISTS " + tableName)
 		if err != nil {
 			return err
 		}
 	}
 	return nil
 }
 
-func createTables(logger lager.Logger, db *sql.DB, flavor string) error {
+func (e *InitSQL) createTables(tx *sql.Tx, logger lager.Logger) error {
 	var createTablesSQL = []string{
-		helpers.RebindForFlavor(createDomainSQL, flavor),
-		helpers.RebindForFlavor(createDesiredLRPsSQL, flavor),
-		helpers.RebindForFlavor(createActualLRPsSQL, flavor),
-		helpers.RebindForFlavor(createTasksSQL, flavor),
+		helpers.RebindForFlavor(createDomainSQL, e.dbFlavor),
+		helpers.RebindForFlavor(createDesiredLRPsSQL, e.dbFlavor),
+		helpers.RebindForFlavor(createActualLRPsSQL, e.dbFlavor),
+		helpers.RebindForFlavor(createTasksSQL, e.dbFlavor),
 	}
 
 	logger.Info("creating-tables")
 	for _, query := range createTablesSQL {
 		logger.Info("creating the table", lager.Data{"query": query})
-		_, err := db.Exec(query)
+		_, err := tx.Exec(query)
 		if err != nil {
 			logger.Error("failed-creating-tables", err)
 			return err
@@ -103,7 +98,7 @@ func createTables(logger lager.Logger, db *sql.DB, flavor string) error {
 	return nil
 }
 
-func createIndices(logger lager.Logger, db *sql.DB) error {
+func (e *InitSQL) createIndices(tx *sql.Tx, logger lager.Logger) error {
 	logger.Info("creating-indices")
 	createIndicesSQL := []string{}
 	createIndicesSQL = append(createIndicesSQL, createDomainsIndices...)
@@ -113,7 +108,7 @@ func createIndices(logger lager.Logger, db *sql.DB) error {
 
 	for _, query := range createIndicesSQL {
 		logger.Info("creating the index", lager.Data{"query": query})
-		_, err := db.Exec(query)
+		_, err := tx.Exec(query)
 		if err != nil {
 			logger.Error("failed-creating-index", err)
 			return err