Handle NULL migration errors in migration helpers

This ensures that whenever changing the NULL constraint of a column
fails we still drop the column.

2 files changed, 15 insertions, 2 deletions

diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index 978c3f7896d..0f488e968f6 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -126,6 +126,8 @@ module Gitlab
         begin
           transaction do
             update_column_in_batches(table, column, default)
+
+            change_column_null(table, column, false) unless allow_null
           end
         # We want to rescue _all_ exceptions here, even those that don't inherit
         # from StandardError.
@@ -134,8 +136,6 @@ module Gitlab
 
           raise error
         end
-
-        change_column_null(table, column, false) unless allow_null
       end
     end
   end
diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index 83ddabe6b0b..1ec539066a7 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers_spec.rb
@@ -120,6 +120,19 @@ describe Gitlab::Database::MigrationHelpers, lib: true do
           model.add_column_with_default(:projects, :foo, :integer, default: 10)
         end.to raise_error(RuntimeError)
       end
+
+      it 'removes the added column whenever changing a column NULL constraint fails' do
+        expect(model).to receive(:change_column_null).
+          with(:projects, :foo, false).
+          and_raise(RuntimeError)
+
+        expect(model).to receive(:remove_column).
+          with(:projects, :foo)
+
+        expect do
+          model.add_column_with_default(:projects, :foo, :integer, default: 10)
+        end.to raise_error(RuntimeError)
+      end
     end
 
     context 'inside a transaction' do