Support multiple columns in column copy triggers

lib/gitlab/database/migration_helpers.rb

@@ -577,17 +577,7 @@ module Gitlab
       # old_column - The name of the old column.
       # new_column - The name of the new column.
       def install_rename_triggers(table, old_column, new_column)
-        trigger_name = rename_trigger_name(table, old_column, new_column)
-        quoted_table = quote_table_name(table)
-        quoted_old = quote_column_name(old_column)
-        quoted_new = quote_column_name(new_column)
-
-        install_rename_triggers_for_postgresql(
-          trigger_name,
-          quoted_table,
-          quoted_old,
-          quoted_new
-        )
+        install_rename_triggers_for_postgresql(table, old_column, new_column)
       end
 
       # Changes the type of a column concurrently.
@@ -1054,43 +1044,18 @@ module Gitlab
       end
 
       # Performs a concurrent column rename when using PostgreSQL.
-      def install_rename_triggers_for_postgresql(trigger, table, old, new)
-        execute <<-EOF.strip_heredoc
-        CREATE OR REPLACE FUNCTION #{trigger}()
-        RETURNS trigger AS
-        $BODY$
-        BEGIN
-          NEW.#{new} := NEW.#{old};
-          RETURN NEW;
-        END;
-        $BODY$
-        LANGUAGE 'plpgsql'
-        VOLATILE
-        EOF
-
-        execute <<-EOF.strip_heredoc
-        DROP TRIGGER IF EXISTS #{trigger}
-        ON #{table}
-        EOF
-
-        execute <<-EOF.strip_heredoc
-        CREATE TRIGGER #{trigger}
-        BEFORE INSERT OR UPDATE
-        ON #{table}
-        FOR EACH ROW
-        EXECUTE FUNCTION #{trigger}()
-        EOF
+      def install_rename_triggers_for_postgresql(table, old, new, trigger_name: nil)
+        Gitlab::Database::UnidirectionalCopyTrigger.on_table(table).create(old, new, trigger_name: trigger_name)
       end
 
       # Removes the triggers used for renaming a PostgreSQL column concurrently.
       def remove_rename_triggers_for_postgresql(table, trigger)
-        execute("DROP TRIGGER IF EXISTS #{trigger} ON #{table}")
-        execute("DROP FUNCTION IF EXISTS #{trigger}()")
+        Gitlab::Database::UnidirectionalCopyTrigger.on_table(table).drop(trigger)
       end
 
       # Returns the (base) name to use for triggers when renaming columns.
       def rename_trigger_name(table, old, new)
-        'trigger_' + Digest::SHA256.hexdigest("#{table}_#{old}_#{new}").first(12)
+        Gitlab::Database::UnidirectionalCopyTrigger.on_table(table).name(old, new)
       end
 
       # Returns an Array containing the indexes for the given column

lib/gitlab/database/unidirectional_copy_trigger.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    class UnidirectionalCopyTrigger
+      def self.on_table(table_name, connection: ActiveRecord::Base.connection)
+        new(table_name, connection)
+      end
+
+      def name(from_column_names, to_column_names)
+        from_column_names, to_column_names = check_column_names!(from_column_names, to_column_names)
+
+        unchecked_name(from_column_names, to_column_names)
+      end
+
+      def create(from_column_names, to_column_names, trigger_name: nil)
+        from_column_names, to_column_names = check_column_names!(from_column_names, to_column_names)
+        trigger_name ||= unchecked_name(from_column_names, to_column_names)
+
+        assignment_clauses = assignment_clauses_for_columns(from_column_names, to_column_names)
+
+        connection.execute(<<~SQL)
+          CREATE OR REPLACE FUNCTION #{trigger_name}()
+          RETURNS trigger AS
+          $BODY$
+          BEGIN
+            #{assignment_clauses};
+            RETURN NEW;
+          END;
+          $BODY$
+          LANGUAGE 'plpgsql'
+          VOLATILE
+        SQL
+
+        connection.execute(<<~SQL)
+          DROP TRIGGER IF EXISTS #{trigger_name}
+          ON #{quoted_table_name}
+        SQL
+
+        connection.execute(<<~SQL)
+          CREATE TRIGGER #{trigger_name}
+          BEFORE INSERT OR UPDATE
+          ON #{quoted_table_name}
+          FOR EACH ROW
+          EXECUTE FUNCTION #{trigger_name}()
+        SQL
+      end
+
+      def drop(trigger_name)
+        connection.execute("DROP TRIGGER IF EXISTS #{trigger_name} ON #{quoted_table_name}")
+        connection.execute("DROP FUNCTION IF EXISTS #{trigger_name}()")
+      end
+
+      private
+
+      attr_reader :table_name, :connection
+
+      def initialize(table_name, connection)
+        @table_name = table_name
+        @connection = connection
+      end
+
+      def quoted_table_name
+        @quoted_table_name ||= connection.quote_table_name(table_name)
+      end
+
+      def check_column_names!(from_column_names, to_column_names)
+        from_column_names = Array.wrap(from_column_names)
+        to_column_names = Array.wrap(to_column_names)
+
+        unless from_column_names.size == to_column_names.size
+          raise ArgumentError, 'number of source and destination columns must match'
+        end
+
+        [from_column_names, to_column_names]
+      end
+
+      def unchecked_name(from_column_names, to_column_names)
+        joined_column_names = from_column_names.zip(to_column_names).flatten.join('_')
+        'trigger_' + Digest::SHA256.hexdigest("#{table_name}_#{joined_column_names}").first(12)
+      end
+
+      def assignment_clauses_for_columns(from_column_names, to_column_names)
+        combined_column_names = to_column_names.zip(from_column_names)
+
+        assignment_clauses = combined_column_names.map do |(new_name, old_name)|
+          new_name = connection.quote_column_name(new_name)
+          old_name = connection.quote_column_name(old_name)
+
+          "NEW.#{new_name} := NEW.#{old_name}"
+        end
+
+        assignment_clauses.join(";\n  ")
+      end
+    end
+  end
+end

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -835,7 +835,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           expect(model).to receive(:check_trigger_permissions!).with(:users)
 
           expect(model).to receive(:install_rename_triggers_for_postgresql)
-            .with(trigger_name, '"users"', '"old"', '"new"')
+            .with(:users, :old, :new)
 
           expect(model).to receive(:add_column)
             .with(:users, :new, :integer,
@@ -860,14 +860,18 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
         context 'with existing records and type casting' do
           let(:trigger_name) { model.rename_trigger_name(:users, :id, :new) }
           let(:user) { create(:user) }
+          let(:copy_trigger) { double('copy trigger') }
+
+          before do
+            expect(Gitlab::Database::UnidirectionalCopyTrigger).to receive(:on_table)
+              .with(:users).and_return(copy_trigger)
+          end
 
           it 'copies the value to the new column using the type_cast_function', :aggregate_failures do
             expect(model).to receive(:copy_indexes).with(:users, :id, :new)
             expect(model).to receive(:add_not_null_constraint).with(:users, :new)
             expect(model).to receive(:execute).with("UPDATE \"users\" SET \"new\" = cast_to_jsonb_with_default(\"users\".\"id\") WHERE \"users\".\"id\" >= #{user.id}")
-            expect(model).to receive(:execute).with("DROP TRIGGER IF EXISTS #{trigger_name}\nON \"users\"\n")
-            expect(model).to receive(:execute).with("CREATE TRIGGER #{trigger_name}\nBEFORE INSERT OR UPDATE\nON \"users\"\nFOR EACH ROW\nEXECUTE FUNCTION #{trigger_name}()\n")
-            expect(model).to receive(:execute).with("CREATE OR REPLACE FUNCTION #{trigger_name}()\nRETURNS trigger AS\n$BODY$\nBEGIN\n  NEW.\"new\" := NEW.\"id\";\n  RETURN NEW;\nEND;\n$BODY$\nLANGUAGE 'plpgsql'\nVOLATILE\n")
+            expect(copy_trigger).to receive(:create).with(:id, :new, trigger_name: nil)
 
             model.rename_column_concurrently(:users, :id, :new, type_cast_function: 'cast_to_jsonb_with_default')
           end
@@ -996,7 +1000,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
         expect(model).to receive(:check_trigger_permissions!).with(:users)
 
         expect(model).to receive(:install_rename_triggers_for_postgresql)
-          .with(trigger_name, '"users"', '"old"', '"new"')
+          .with(:users, :old, :new)
 
         expect(model).to receive(:add_column)
           .with(:users, :old, :integer,
@@ -1156,7 +1160,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           .with(:users, temp_undo_cleanup_column, :old)
 
         expect(model).to receive(:install_rename_triggers_for_postgresql)
-          .with(trigger_name, '"users"', '"old"', '"old_for_type_change"')
+          .with(:users, :old, 'old_for_type_change')
 
         model.undo_cleanup_concurrent_column_type_change(:users, :old, :string)
       end
@@ -1182,7 +1186,7 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
           .with(:users, temp_undo_cleanup_column, :old)
 
         expect(model).to receive(:install_rename_triggers_for_postgresql)
-          .with(trigger_name, '"users"', '"old"', '"old_for_type_change"')
+          .with(:users, :old, 'old_for_type_change')
 
         model.undo_cleanup_concurrent_column_type_change(
           :users,
@@ -1204,28 +1208,25 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
   describe '#install_rename_triggers_for_postgresql' do
     it 'installs the triggers for PostgreSQL' do
-      expect(model).to receive(:execute)
-        .with(/CREATE OR REPLACE FUNCTION foo()/m)
+      copy_trigger = double('copy trigger')
 
-      expect(model).to receive(:execute)
-        .with(/DROP TRIGGER IF EXISTS foo/m)
+      expect(Gitlab::Database::UnidirectionalCopyTrigger).to receive(:on_table)
+        .with(:users).and_return(copy_trigger)
 
-      expect(model).to receive(:execute)
-        .with(/CREATE TRIGGER foo/m)
+      expect(copy_trigger).to receive(:create).with(:old, :new, trigger_name: 'foo')
 
-      model.install_rename_triggers_for_postgresql('foo', :users, :old, :new)
-    end
-
-    it 'does not fail if trigger already exists' do
-      model.install_rename_triggers_for_postgresql('foo', :users, :old, :new)
-      model.install_rename_triggers_for_postgresql('foo', :users, :old, :new)
+      model.install_rename_triggers_for_postgresql(:users, :old, :new, trigger_name: 'foo')
     end
   end
 
   describe '#remove_rename_triggers_for_postgresql' do
     it 'removes the function and trigger' do
-      expect(model).to receive(:execute).with('DROP TRIGGER IF EXISTS foo ON bar')
-      expect(model).to receive(:execute).with('DROP FUNCTION IF EXISTS foo()')
+      copy_trigger = double('copy trigger')
+
+      expect(Gitlab::Database::UnidirectionalCopyTrigger).to receive(:on_table)
+        .with('bar').and_return(copy_trigger)
+
+      expect(copy_trigger).to receive(:drop).with('foo')
 
       model.remove_rename_triggers_for_postgresql('bar', 'foo')
     end

spec/lib/gitlab/database/unidirectional_copy_trigger_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::UnidirectionalCopyTrigger do
+  include Database::TriggerHelpers
+
+  let(:table_name) { '_test_table' }
+  let(:connection) { ActiveRecord::Base.connection }
+  let(:copy_trigger) { described_class.on_table(table_name) }
+
+  describe '#name' do
+    context 'when a single column name is given' do
+      subject(:trigger_name) { copy_trigger.name('id', 'other_id') }
+
+      it 'returns the trigger name' do
+        expect(trigger_name).to eq('trigger_cfce7a56a9d6')
+      end
+    end
+
+    context 'when multiple column names are given' do
+      subject(:trigger_name) { copy_trigger.name(%w[id fk_id], %w[other_id other_fk_id]) }
+
+      it 'returns the trigger name' do
+        expect(trigger_name).to eq('trigger_166626e51481')
+      end
+    end
+
+    context 'when a different number of new and old column names are given' do
+      it 'raises an error' do
+        expect do
+          copy_trigger.name(%w[id fk_id], %w[other_id])
+        end.to raise_error(ArgumentError, 'number of source and destination columns must match')
+      end
+    end
+  end
+
+  describe '#create' do
+    let(:model) { Class.new(ActiveRecord::Base) }
+
+    before do
+      connection.execute(<<~SQL)
+        CREATE TABLE #{table_name} (
+          id serial NOT NULL PRIMARY KEY,
+          other_id integer,
+          fk_id bigint,
+          other_fk_id bigint);
+      SQL
+
+      model.table_name = table_name
+    end
+
+    context 'when a single column name is given' do
+      let(:trigger_name) { 'trigger_cfce7a56a9d6' }
+
+      it 'creates the trigger and function' do
+        expect_function_not_to_exist(trigger_name)
+        expect_trigger_not_to_exist(table_name, trigger_name)
+
+        copy_trigger.create('id', 'other_id')
+
+        expect_function_to_exist(trigger_name)
+        expect_valid_function_trigger(table_name, trigger_name, trigger_name, before: %w[insert update])
+      end
+
+      it 'properly copies the column data using the trigger function' do
+        copy_trigger.create('id', 'other_id')
+
+        record = model.create!(id: 10)
+        expect(record.reload).to have_attributes(other_id: 10)
+
+        record.update!({ id: 20 })
+        expect(record.reload).to have_attributes(other_id: 20)
+      end
+    end
+
+    context 'when multiple column names are given' do
+      let(:trigger_name) { 'trigger_166626e51481' }
+
+      it 'creates the trigger and function to set all the columns' do
+        expect_function_not_to_exist(trigger_name)
+        expect_trigger_not_to_exist(table_name, trigger_name)
+
+        copy_trigger.create(%w[id fk_id], %w[other_id other_fk_id])
+
+        expect_function_to_exist(trigger_name)
+        expect_valid_function_trigger(table_name, trigger_name, trigger_name, before: %w[insert update])
+      end
+
+      it 'properly copies the columns using the trigger function' do
+        copy_trigger.create(%w[id fk_id], %w[other_id other_fk_id])
+
+        record = model.create!(id: 10, fk_id: 20)
+        expect(record.reload).to have_attributes(other_id: 10, other_fk_id: 20)
+
+        record.update!(id: 30, fk_id: 50)
+        expect(record.reload).to have_attributes(other_id: 30, other_fk_id: 50)
+      end
+    end
+
+    context 'when a custom trigger name is given' do
+      let(:trigger_name) { '_test_trigger' }
+
+      it 'creates the trigger and function with the custom name' do
+        expect_function_not_to_exist(trigger_name)
+        expect_trigger_not_to_exist(table_name, trigger_name)
+
+        copy_trigger.create('id', 'other_id', trigger_name: trigger_name)
+
+        expect_function_to_exist(trigger_name)
+        expect_valid_function_trigger(table_name, trigger_name, trigger_name, before: %w[insert update])
+      end
+    end
+
+    context 'when the trigger function already exists' do
+      let(:trigger_name) { 'trigger_cfce7a56a9d6' }
+
+      it 'does not raise an error' do
+        expect_function_not_to_exist(trigger_name)
+        expect_trigger_not_to_exist(table_name, trigger_name)
+
+        copy_trigger.create('id', 'other_id')
+
+        expect_function_to_exist(trigger_name)
+        expect_valid_function_trigger(table_name, trigger_name, trigger_name, before: %w[insert update])
+
+        copy_trigger.create('id', 'other_id')
+
+        expect_function_to_exist(trigger_name)
+        expect_valid_function_trigger(table_name, trigger_name, trigger_name, before: %w[insert update])
+      end
+    end
+
+    context 'when a different number of new and old column names are given' do
+      it 'raises an error' do
+        expect do
+          copy_trigger.create(%w[id fk_id], %w[other_id])
+        end.to raise_error(ArgumentError, 'number of source and destination columns must match')
+      end
+    end
+  end
+
+  describe '#drop' do
+    let(:trigger_name) { '_test_trigger' }
+
+    before do
+      connection.execute(<<~SQL)
+        CREATE TABLE #{table_name} (
+          id serial NOT NULL PRIMARY KEY,
+          other_id integer NOT NULL);
+
+        CREATE FUNCTION #{trigger_name}()
+        RETURNS trigger
+        LANGUAGE plpgsql AS
+        $$
+        BEGIN
+          RAISE NOTICE 'hello';
+          RETURN NEW;
+        END
+        $$;
+
+        CREATE TRIGGER #{trigger_name}
+        BEFORE INSERT OR UPDATE
+        ON #{table_name}
+        FOR EACH ROW
+        EXECUTE FUNCTION #{trigger_name}();
+      SQL
+    end
+
+    it 'drops the trigger and function for the given arguments' do
+      expect_function_to_exist(trigger_name)
+      expect_valid_function_trigger(table_name, trigger_name, trigger_name, before: %w[insert update])
+
+      copy_trigger.drop(trigger_name)
+
+      expect_trigger_not_to_exist(table_name, trigger_name)
+      expect_function_not_to_exist(trigger_name)
+    end
+
+    context 'when the trigger does not exist' do
+      it 'does not raise an error' do
+        copy_trigger.drop(trigger_name)
+
+        expect_trigger_not_to_exist(table_name, trigger_name)
+        expect_function_not_to_exist(trigger_name)
+
+        copy_trigger.drop(trigger_name)
+      end
+    end
+  end
+end