Wrap postgres up migrations into a single transaction (#887)

internal/db/schema/postgres/postgres.go

@@ -57,6 +57,8 @@ type Postgres struct {
 	// Locking and unlocking need to use the same connection
 	conn *sql.Conn
 	db   *sql.DB
+
+	tx *sql.Tx
 }
 
 // New returns a postgres pointer with the provided db verified as
@@ -154,10 +156,41 @@ func (p *Postgres) UnlockShared(ctx context.Context) error {
 	return nil
 }
 
+// StartRun starts a transaction that all subsequent calls to Run will use.
+func (p *Postgres) StartRun(ctx context.Context) error {
+	tx, err := p.conn.BeginTx(ctx, nil)
+	if err != nil {
+		return err
+	}
+	p.tx = tx
+	return nil
+}
+
+// CommitRun commits the pending transaction if there is one
+func (p *Postgres) CommitRun() error {
+	const op = "postgres.(Postgres).CommitRun"
+	defer func() {
+		p.tx = nil
+	}()
+	if p.tx == nil {
+		return errors.New(errors.MigrationIntegrity, op, "no pending transaction")
+	}
+	if err := p.tx.Commit(); err != nil {
+		if errRollback := p.tx.Rollback(); errRollback != nil {
+			return multierror.Append(err, errRollback)
+		}
+	}
+	return nil
+}
+
+type execContexter interface {
+	ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error)
+}
+
 // Executes the sql provided in the passed in io.Reader.  The contents of the reader must
 // fit in memory as the full content is read into a string before being passed to the
 // backing database.
-func (p *Postgres) Run(ctx context.Context, migration io.Reader) error {
+func (p *Postgres) Run(ctx context.Context, migration io.Reader, version int) error {
 	const op = "postgres.(Postgres).Run"
 	migr, err := ioutil.ReadAll(migration)
 	if err != nil {
@@ -165,7 +198,21 @@ func (p *Postgres) Run(ctx context.Context, migration io.Reader) error {
 	}
 	// Run migration
 	query := string(migr)
-	if _, err := p.conn.ExecContext(ctx, query); err != nil {
+
+	var extr execContexter = p.conn
+	rollback := func() error { return nil }
+	if p.tx != nil {
+		extr = p.tx
+		rollback = func() error {
+			defer func() { p.tx = nil }()
+			return p.tx.Rollback()
+		}
+	}
+
+	if _, err := extr.ExecContext(ctx, query); err != nil {
+		if rollbackErr := rollback(); rollbackErr != nil {
+			err = multierror.Append(err, rollbackErr)
+		}
 		if pgErr, ok := err.(*pq.Error); ok {
 			var line uint
 			var col uint
@@ -187,6 +234,9 @@ func (p *Postgres) Run(ctx context.Context, migration io.Reader) error {
 		}
 		return errors.Wrap(err, op, errors.WithMsg(fmt.Sprintf("migration failed: %s", migr)))
 	}
+	if err := p.setVersion(ctx, version, false); err != nil {
+		return errors.Wrap(err, op)
+	}
 
 	return nil
 }
@@ -226,18 +276,26 @@ func runesLastIndex(input []rune, target rune) int {
 	return -1
 }
 
-// SetVersion sets the version number, and whether the database is in a dirty state.
+// setVersion sets the version number, and whether the database is in a dirty state.
 // A version value of -1 indicates no version is set.
-func (p *Postgres) SetVersion(ctx context.Context, version int, dirty bool) error {
+func (p *Postgres) setVersion(ctx context.Context, version int, dirty bool) error {
 	const op = "postgres.(Postgres).SetVersion"
-	tx, err := p.conn.BeginTx(ctx, &sql.TxOptions{})
-	if err != nil {
-		return err
+	tx := p.tx
+	var err error
+	if tx == nil {
+		tx, err = p.conn.BeginTx(ctx, nil)
+		if err != nil {
+			return err
+		}
+	}
+	rollback := func() error {
+		defer func() { p.tx = nil }()
+		return tx.Rollback()
 	}
 
 	query := `TRUNCATE ` + pq.QuoteIdentifier(defaultMigrationsTable)
 	if _, err := tx.ExecContext(ctx, query); err != nil {
-		if errRollback := tx.Rollback(); errRollback != nil {
+		if errRollback := rollback(); errRollback != nil {
 			err = multierror.Append(err, errRollback)
 		}
 		return errors.Wrap(err, op)
@@ -250,24 +308,25 @@ func (p *Postgres) SetVersion(ctx context.Context, version int, dirty bool) erro
 		query = `INSERT INTO ` + pq.QuoteIdentifier(defaultMigrationsTable) +
 			` (Version, dirty) VALUES ($1, $2)`
 		if _, err := tx.ExecContext(ctx, query, version, dirty); err != nil {
-			if errRollback := tx.Rollback(); errRollback != nil {
+			if errRollback := rollback(); errRollback != nil {
 				err = multierror.Append(err, errRollback)
 			}
 			return errors.Wrap(err, op)
 		}
 	}
 
-	if err := tx.Commit(); err != nil {
-		return errors.Wrap(err, op)
+	if p.tx == nil {
+		if err := tx.Commit(); err != nil {
+			return errors.Wrap(err, op)
+		}
 	}
-
 	return nil
 }
 
 // CurrentState returns the version, if the database is currently in a dirty state, and any error.
 // A version value of -1 indicates no version is set.
 func (p *Postgres) CurrentState(ctx context.Context) (version int, dirty bool, err error) {
-	const op = "postgres.(Postgres).Version"
+	const op = "postgres.(Postgres).CurrentState"
 	query := `SELECT Version, dirty FROM ` + pq.QuoteIdentifier(defaultMigrationsTable) + ` LIMIT 1`
 	err = p.conn.QueryRowContext(ctx, query).Scan(&version, &dirty)
 	switch {

internal/db/schema/postgres/postgres_test.go

@@ -163,7 +163,7 @@ func TestMultiStatement(t *testing.T) {
 				t.Error(err)
 			}
 		}()
-		if err := d.Run(ctx, strings.NewReader("CREATE TABLE foo (foo text); CREATE TABLE bar (bar text);")); err != nil {
+		if err := d.Run(ctx, strings.NewReader("CREATE TABLE foo (foo text); CREATE TABLE bar (bar text);"), 2); err != nil {
 			t.Fatalf("expected err to be nil, got %v", err)
 		}
 
@@ -178,6 +178,45 @@ func TestMultiStatement(t *testing.T) {
 	})
 }
 
+func TestTransaction(t *testing.T) {
+	dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
+		ctx := context.Background()
+		ip, port, err := c.FirstPort()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		addr := pgConnectionString(ip, port)
+		d, err := open(t, ctx, addr)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer func() {
+			if err := d.close(t); err != nil {
+				t.Error(err)
+			}
+		}()
+
+		assert.NoError(t, d.StartRun(ctx))
+		assert.NoError(t, d.Run(ctx, strings.NewReader("CREATE TABLE foo (foo text);"), 2))
+		assert.NoError(t, d.Run(ctx, strings.NewReader("SELECT 1"), 3))
+		assert.NoError(t, d.CommitRun())
+		v, dirty, err := d.CurrentState(ctx)
+		assert.NoError(t, err)
+		assert.False(t, dirty)
+		assert.Equal(t, 3, v)
+
+		assert.NoError(t, d.StartRun(ctx))
+		assert.NoError(t, d.Run(ctx, strings.NewReader("CREATE TABLE bar (bar text);"), 20))
+		assert.Error(t, d.Run(ctx, strings.NewReader("SELECT 1 FROM NonExistingTable"), 30))
+		assert.Error(t, d.CommitRun())
+		v, dirty, err = d.CurrentState(ctx)
+		assert.NoError(t, err)
+		assert.False(t, dirty)
+		assert.Equal(t, 3, v)
+	})
+}
+
 func TestWithSchema(t *testing.T) {
 	dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
 		ctx := context.Background()
@@ -198,10 +237,7 @@ func TestWithSchema(t *testing.T) {
 		}()
 
 		// create foobar schema
-		if err := d.Run(ctx, strings.NewReader("CREATE SCHEMA foobar AUTHORIZATION postgres")); err != nil {
-			t.Fatal(err)
-		}
-		if err := d.SetVersion(ctx, 1, false); err != nil {
+		if err := d.Run(ctx, strings.NewReader("CREATE SCHEMA foobar AUTHORIZATION postgres"), 1); err != nil {
 			t.Fatal(err)
 		}
 
@@ -226,7 +262,7 @@ func TestWithSchema(t *testing.T) {
 		}
 
 		// now update CurrentState and compare
-		if err := d2.SetVersion(ctx, 2, false); err != nil {
+		if err := d2.setVersion(ctx, 2, false); err != nil {
 			t.Fatal(err)
 		}
 		version, _, err = d2.CurrentState(ctx)
@@ -303,7 +339,7 @@ func TestRun_Error(t *testing.T) {
 			require.NoError(t, p.close(t))
 		})
 
-		err = p.Run(ctx, bytes.NewReader([]byte("SELECT *\nFROM foo")))
+		err = p.Run(ctx, bytes.NewReader([]byte("SELECT *\nFROM foo")), 2)
 		assert.Error(t, err)
 	})
 }

internal/db/schema/postgres/testing.go

@@ -98,7 +98,7 @@ func testRun(t *testing.T, d *Postgres, migration io.Reader) {
 		t.Fatal("migration can't be nil")
 	}
 
-	if err := d.Run(ctx, migration); err != nil {
+	if err := d.Run(ctx, migration, 1); err != nil {
 		t.Fatal(err)
 	}
 }
@@ -132,7 +132,7 @@ func testSetVersion(t *testing.T, d *Postgres) {
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			err := d.SetVersion(ctx, tc.version, tc.dirty)
+			err := d.setVersion(ctx, tc.version, tc.dirty)
 			if err != tc.expectedErr {
 				t.Fatal("Got unexpected error:", err, "!=", tc.expectedErr)
 			}