Issue an explicit rollback if a migration statement returns any error

nicheinc · Sep 20, 2023 · 723584e · 723584e
1 parent 856ea12
commit 723584e
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 8 deletions.
diff --git a/database/pgx/pgx.go b/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
 }

diff --git a/database/pgx/v5/pgx.go b/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
 }

diff --git a/database/postgres/postgres.go b/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
 }