Move `replace_sql` into `Database::MigrationHelpers`

2 files changed, 56 insertions, 0 deletions

diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index 6dabbe0264c..298b1a1f4e6 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -498,6 +498,29 @@ module Gitlab
 
         columns(table).find { |column| column.name == name }
       end
+
+      # This will replace the first occurance of a string in a column with
+      # the replacement
+      # On postgresql we can use `regexp_replace` for that.
+      # On mysql we find the location of the pattern, and overwrite it
+      # with the replacement
+      def replace_sql(column, pattern, replacement)
+        quoted_pattern = Arel::Nodes::Quoted.new(pattern.to_s)
+        quoted_replacement = Arel::Nodes::Quoted.new(replacement.to_s)
+
+        if Database.mysql?
+          locate = Arel::Nodes::NamedFunction.
+            new('locate', [quoted_pattern, column])
+          insert_in_place = Arel::Nodes::NamedFunction.
+            new('insert', [column, locate, pattern.size, quoted_replacement])
+
+          Arel::Nodes::SqlLiteral.new(insert_in_place.to_sql)
+        else
+          replace = Arel::Nodes::NamedFunction.
+            new("regexp_replace", [column, quoted_pattern, quoted_replacement])
+          Arel::Nodes::SqlLiteral.new(replace.to_sql)
+        end
+      end
     end
   end
 end
diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index a044b871730..737fac14f92 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers_spec.rb
@@ -726,4 +726,37 @@ describe Gitlab::Database::MigrationHelpers, lib: true do
       expect(model.column_for(:users, :kittens)).to be_nil
     end
   end
+
+  describe '#replace_sql' do
+    context 'using postgres' do
+      before do
+        allow(Gitlab::Database).to receive(:mysql?).and_return(false)
+      end
+
+      it 'builds the sql with correct functions' do
+        expect(model.replace_sql(Arel::Table.new(:users)[:first_name], "Alice", "Eve").to_s).
+          to include('regexp_replace')
+      end
+    end
+
+    context 'using mysql' do
+      before do
+        allow(Gitlab::Database).to receive(:mysql?).and_return(true)
+      end
+
+      it 'builds the sql with the correct functions' do
+        expect(model.replace_sql(Arel::Table.new(:users)[:first_name], "Alice", "Eve").to_s).
+          to include('locate', 'insert')
+      end
+    end
+
+    describe 'results' do
+      let!(:user) { create(:user, name: 'Kathy Alice Aliceson') }
+
+      it 'replaces the correct part of the string' do
+        model.update_column_in_batches(:users, :name, model.replace_sql(Arel::Table.new(:users)[:name], 'Alice', 'Eve'))
+        expect(user.reload.name).to eq('Kathy Eve Aliceson')
+      end
+    end
+  end
 end