Issue an explicit rollback if a migration statement returns any error

Author: bclarkx2

database/pgx/pgx.go

@@ -285,7 +285,7 @@ func (p *Postgres) runStatement(statement []byte) error {
 		return nil
 	}
 	if _, err := p.conn.ExecContext(ctx, query); err != nil {
-
+		var migrationErr error
 		if pgErr, ok := err.(*pgconn.PgError); ok {
 			var line uint
 			var col uint
@@ -298,9 +298,21 @@ func (p *Postgres) runStatement(statement []byte) error {
 			if pgErr.Detail != "" {
 				message = fmt.Sprintf("%s, %s", message, pgErr.Detail)
 			}
-			return database.Error{OrigErr: err, Err: message, Query: statement, Line: line}
+			migrationErr = database.Error{OrigErr: err, Err: message, Query: statement, Line: line}
+		} else {
+			migrationErr = database.Error{OrigErr: err, Err: "migration failed", Query: statement}
+		}
+
+		// For safety, always issue a rollback on error. In multi-statement
+		// mode, this is necessary to make sure that the connection is not left
+		// in an aborted state. In single-statement mode, this will be a no-op
+		// outside of the implicit transaction block that was already rolled
+		// back.
+		if _, rollbackErr := p.conn.ExecContext(ctx, "ROLLBACK"); rollbackErr != nil {
+			rollbackErr = fmt.Errorf("failed to rollback migration tx: %w", rollbackErr)
+			return multierror.Append(migrationErr, rollbackErr)
 		}
-		return database.Error{OrigErr: err, Err: "migration failed", Query: statement}
+		return migrationErr
 	}
 	return nil
 }

database/pgx/pgx_test.go

@@ -167,6 +167,48 @@ func TestMultipleStatements(t *testing.T) {
 	})
 }
 
+func TestMultipleStatementsError(t *testing.T) {
+	dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
+		ip, port, err := c.FirstPort()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		addr := pgConnectionString(ip, port)
+		p := &Postgres{}
+		d, err := p.Open(addr)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer func() {
+			if err := d.Close(); err != nil {
+				t.Error(err)
+			}
+		}()
+
+		// Run a migration with explicit transaction that we expect to fail
+		migrationErr := d.Run(strings.NewReader("BEGIN; SELECT 1/0; COMMIT;"))
+
+		// Migration should return expected error
+		var e *database.Error
+		if !errors.As(migrationErr, &e) || err == nil {
+			t.Fatalf("Unexpected error, want migration error. Got: %s", err)
+		}
+		if !strings.Contains(e.OrigErr.Error(), "division by zero") {
+			t.Fatalf("Migration error missing expected message. Got: %s", err)
+		}
+
+		// Connection should still be usable after failed migration
+		var result int
+		if err := d.(*Postgres).conn.QueryRowContext(context.Background(), "SELECT 123").Scan(&result); err != nil {
+			t.Fatalf("Unexpected error, want connection to be usable. Got: %s", err)
+		}
+		if result != 123 {
+			t.Fatalf("Unexpected result, want 123. Got: %d", result)
+		}
+	})
+}
+
 func TestMultipleStatementsInMultiStatementMode(t *testing.T) {
 	dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
 		ip, port, err := c.FirstPort()

database/pgx/v5/pgx.go

@@ -283,7 +283,7 @@ func (p *Postgres) runStatement(statement []byte) error {
 		return nil
 	}
 	if _, err := p.conn.ExecContext(ctx, query); err != nil {
-
+		var migrationErr error
 		if pgErr, ok := err.(*pgconn.PgError); ok {
 			var line uint
 			var col uint
@@ -296,9 +296,21 @@ func (p *Postgres) runStatement(statement []byte) error {
 			if pgErr.Detail != "" {
 				message = fmt.Sprintf("%s, %s", message, pgErr.Detail)
 			}
-			return database.Error{OrigErr: err, Err: message, Query: statement, Line: line}
+			migrationErr = database.Error{OrigErr: err, Err: message, Query: statement, Line: line}
+		} else {
+			migrationErr = database.Error{OrigErr: err, Err: "migration failed", Query: statement}
+		}
+
+		// For safety, always issue a rollback on error. In multi-statement
+		// mode, this is necessary to make sure that the connection is not left
+		// in an aborted state. In single-statement mode, this will be a no-op
+		// outside of the implicit transaction block that was already rolled
+		// back.
+		if _, rollbackErr := p.conn.ExecContext(ctx, "ROLLBACK"); rollbackErr != nil {
+			rollbackErr = fmt.Errorf("failed to rollback migration tx: %w", rollbackErr)
+			return multierror.Append(migrationErr, rollbackErr)
 		}
-		return database.Error{OrigErr: err, Err: "migration failed", Query: statement}
+		return migrationErr
 	}
 	return nil
 }

database/postgres/TUTORIAL.md

@@ -68,6 +68,12 @@ Make sure to check if your database changed as expected in this case as well.
 
 ## Database transactions
 
+By default, all the statements in the migration file will be run inside one
+implicit transaction. However, you can also use `BEGIN` and `COMMIT` statements
+to explicitly control what transactions your migration uses (for example, if you
+want to break your migration into multiple transactions to avoid holding onto a
+lock for too long).
+
 To show database transactions usage, let's create another set of migrations by running:
 ```
 migrate create -ext sql -dir db/migrations -seq add_mood_to_users
@@ -76,7 +82,6 @@ Again, it should create for us two migrations files:
 - 000002_add_mood_to_users.down.sql
 - 000002_add_mood_to_users.up.sql
 
-In Postgres, when we want our queries to be done in a transaction, we need to wrap it with `BEGIN` and `COMMIT` commands.
 In our example, we are going to add a column to our database that can only accept enumerable values or NULL.
 Migration up:
 ```

database/postgres/postgres.go

@@ -296,6 +296,7 @@ func (p *Postgres) runStatement(statement []byte) error {
 		return nil
 	}
 	if _, err := p.conn.ExecContext(ctx, query); err != nil {
+		var migrationErr error
 		if pgErr, ok := err.(*pq.Error); ok {
 			var line uint
 			var col uint
@@ -312,9 +313,21 @@ func (p *Postgres) runStatement(statement []byte) error {
 			if pgErr.Detail != "" {
 				message = fmt.Sprintf("%s, %s", message, pgErr.Detail)
 			}
-			return database.Error{OrigErr: err, Err: message, Query: statement, Line: line}
+			migrationErr = database.Error{OrigErr: err, Err: message, Query: statement, Line: line}
+		} else {
+			migrationErr = database.Error{OrigErr: err, Err: "migration failed", Query: statement}
 		}
-		return database.Error{OrigErr: err, Err: "migration failed", Query: statement}
+
+		// For safety, always issue a rollback on error. In multi-statement
+		// mode, this is necessary to make sure that the connection is not left
+		// in an aborted state. In single-statement mode, this will be a no-op
+		// outside of the implicit transaction block that was already rolled
+		// back.
+		if _, rollbackErr := p.conn.ExecContext(ctx, "ROLLBACK"); rollbackErr != nil {
+			rollbackErr = fmt.Errorf("failed to rollback migration tx: %w", rollbackErr)
+			return multierror.Append(migrationErr, rollbackErr)
+		}
+		return migrationErr
 	}
 	return nil
 }