2 files changed, 92 insertions, 9 deletions

diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index 2ae807697bc..c77fdf55789 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -446,10 +446,22 @@ module Gitlab
       # This method can also take a block which is passed directly to the
       # `update_column_in_batches` method.
       def add_column_with_default(table, column, type, default:, limit: nil, allow_null: false, &block)
-        if transaction_open?
-          raise 'add_column_with_default can not be run inside a transaction, ' \
-            'you can disable transactions by calling disable_ddl_transaction! ' \
-            'in the body of your migration class'
+        has_rows = has_rows?(table)
+        if has_rows && transaction_open?
+          raise(<<~MSG.chomp)
+            `add_column_with_default` cannot be run inside a transaction,
+            unless the table is empty!
+
+            You can disable the "single transaction" mode by calling `disable_ddl_transaction!`
+            in the body of your migration class
+          MSG
+        end
+
+        # It is safe to add the column atomically here, since this table is empty
+        unless has_rows
+          opts = { default: default, null: allow_null, limit: limit }
+
+          return add_column(table, column, type, opts.reject { |_k, v| v.nil? })
         end
 
         disable_statement_timeout do
@@ -480,6 +492,13 @@ module Gitlab
         end
       end
 
+      # Tests whether a table contains any data, without knowing anything about
+      # what kind of data that table is meant to hold.
+      def has_rows?(table_name)
+        q = Arel::Table.new(table_name).project(1).take(1)
+        ActiveRecord::Base.connection.exec_query(q.to_sql).present?
+      end
+
       # Renames a column without requiring downtime.
       #
       # Concurrent renames work by using database triggers to ensure both the
diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index dd0033bbc14..94e7b0d863d 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers_spec.rb
@@ -576,6 +576,12 @@ describe Gitlab::Database::MigrationHelpers do
   end
 
   describe '#add_column_with_default' do
+    let(:table_has_rows) { true }
+
+    before do
+      allow(model).to receive(:has_rows?).and_return(table_has_rows)
+    end
+
     context 'outside of a transaction' do
       context 'when a column limit is not set' do
         before do
@@ -675,14 +681,72 @@ describe Gitlab::Database::MigrationHelpers do
     end
 
     context 'inside a transaction' do
-      it 'raises RuntimeError' do
-        expect(model).to receive(:transaction_open?).and_return(true)
+      subject do
+        -> { model.add_column_with_default(:projects, :foo, :integer, default: 10) }
+      end
 
-        expect do
-          model.add_column_with_default(:projects, :foo, :integer, default: 10)
-        end.to raise_error(RuntimeError)
+      context 'the table has rows' do
+        before do
+          expect(model).to receive(:transaction_open?).and_return(true)
+        end
+
+        it { is_expected.to raise_error(RuntimeError) }
+      end
+
+      context 'the table has no rows' do
+        let(:table_has_rows) { false }
+
+        it { is_expected.not_to raise_error(RuntimeError) }
+      end
+    end
+  end
+
+  describe '#has_rows?' do
+    let(:table_name) { :spec_has_rows_test }
+
+    before do
+      model.create_table(table_name) do |t|
+        t.integer :foo
+        t.text :bar
+      end
+    end
+
+    after do
+      model.drop_table table_name
+    end
+
+    subject { model.has_rows? table_name }
+
+    context 'there are no rows' do
+      it { is_expected.to be_falsey }
+    end
+
+    def insert_rows(*rows)
+      table = Arel::Table.new(table_name)
+      rows.each do |foo:, bar:|
+        stm = table.create_insert.insert([
+          [table[:foo], foo],
+          [table[:bar], bar]
+        ])
+        ActiveRecord::Base.connection.insert(stm.to_sql)
       end
     end
+
+    context 'there is one row' do
+      before do
+        insert_rows(foo: 1, bar: 'wibble')
+      end
+
+      it { is_expected.to be_truthy }
+    end
+
+    context 'there are several rows' do
+      before do
+        insert_rows(*(1..3).map { |i| { foo: i, bar: "wibble-#{i}" } })
+      end
+
+      it { is_expected.to be_truthy }
+    end
   end
 
   describe '#rename_column_concurrently' do