Merge branch 'backfill-trigger-helpers' into 'master'

Add database helper for creating triggers to assign sharding keys

See merge request !147401



Merged-by: Krasimir Angelov <kangelov@gitlab.com>
Approved-by: Ian Baum <ibaum@gitlab.com>
Approved-by: Krasimir Angelov <kangelov@gitlab.com>
Reviewed-by: Ian Baum <ibaum@gitlab.com>
Reviewed-by: Krasimir Angelov <kangelov@gitlab.com>
Co-authored-by: Ian Baum <ibaum@gitlab.com>
Co-authored-by: Tiger <twatson@gitlab.com>

lib/gitlab/database/migration_helpers.rb

@@ -541,6 +541,28 @@ def rename_trigger_name(table, old, new)
         Gitlab::Database::UnidirectionalCopyTrigger.on_table(table, connection: connection).name(old, new)
       end
 
+      # Installs a trigger in a table that assigns a sharding key from an associated table.
+      #
+      # table: The table to install the trigger in.
+      # sharding_key: The column to be assigned on `table`.
+      # parent_table: The associated table with the sharding key to be copied.
+      # parent_sharding_key: The sharding key on the parent table that will be copied to `sharding_key` on `table`.
+      # foreign_key: The column used to fetch the relevant record from `parent_table`.
+      def install_sharding_key_assignment_trigger(**args)
+        Gitlab::Database::Triggers::AssignDesiredShardingKey.new(**args.merge(connection: connection)).create
+      end
+
+      # Removes trigger used for assigning sharding keys.
+      #
+      # table: The table to install the trigger in.
+      # sharding_key: The column to be assigned on `table`.
+      # parent_table: The associated table with the sharding key to be copied.
+      # parent_sharding_key: The sharding key on the parent table that will be copied to `sharding_key` on `table`.
+      # foreign_key: The column used to fetch the relevant record from `parent_table`.
+      def remove_sharding_key_assignment_trigger(**args)
+        Gitlab::Database::Triggers::AssignDesiredShardingKey.new(**args.merge(connection: connection)).drop
+      end
+
       # Changes the type of a column concurrently.
       #
       # table - The table containing the column.

lib/gitlab/database/triggers/assign_desired_sharding_key.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Triggers
+      class AssignDesiredShardingKey
+        include Gitlab::Database::SchemaHelpers
+
+        attr_reader :name
+
+        delegate :execute, :quote_table_name, :quote_column_name, to: :connection, private: true
+
+        def initialize(
+          table:, sharding_key:, parent_table:, parent_sharding_key:,
+          foreign_key:, connection:, trigger_name: nil
+        )
+          @table = table
+          @sharding_key = sharding_key
+          @parent_table = parent_table
+          @parent_sharding_key = parent_sharding_key
+          @foreign_key = foreign_key
+          @name = trigger_name || generated_name
+          @connection = connection
+        end
+
+        def create
+          quoted_table_name = quote_table_name(table)
+          quoted_parent_table = quote_table_name(parent_table)
+          quoted_sharding_key = quote_column_name(sharding_key)
+          quoted_parent_sharding_key = quote_column_name(parent_sharding_key)
+          quoted_primary_key = quote_column_name('id')
+          quoted_foreign_key = quote_column_name(foreign_key)
+
+          create_trigger_function(name) do
+            <<~SQL
+              IF NEW.#{quoted_sharding_key} IS NULL THEN
+                SELECT #{quoted_parent_sharding_key}
+                INTO NEW.#{quoted_sharding_key}
+                FROM #{quoted_parent_table}
+                WHERE #{quoted_parent_table}.#{quoted_primary_key} = NEW.#{quoted_foreign_key};
+              END IF;
+
+              RETURN NEW;
+            SQL
+          end
+
+          # Postgres 14 adds the `OR REPLACE` option to trigger creation, so
+          # this line can be removed and `OR REPLACE` added to `#create_trigger`
+          # when the minimum supported version is updated to 14 (milestone 17.0).
+          drop_trigger(quoted_table_name, name)
+
+          create_trigger(quoted_table_name, name, name, fires: 'BEFORE INSERT OR UPDATE')
+        end
+
+        def drop
+          drop_trigger(quote_table_name(table), name)
+          drop_function(name)
+        end
+
+        private
+
+        attr_reader :table, :sharding_key, :parent_table, :parent_sharding_key, :foreign_key, :connection
+
+        def generated_name
+          identifier = "#{table}_assign_#{sharding_key}"
+
+          "trigger_#{Digest::SHA256.hexdigest(identifier).first(12)}"
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -1756,6 +1756,42 @@
     end
   end
 
+  describe '#install_sharding_key_assignment_trigger' do
+    let(:trigger) { double }
+    let(:connection) { ActiveRecord::Base.connection }
+
+    it do
+      expect(Gitlab::Database::Triggers::AssignDesiredShardingKey).to receive(:new)
+        .with(table: :test_table, sharding_key: :project_id, parent_table: :parent_table,
+          parent_sharding_key: :parent_project_id, foreign_key: :foreign_key, connection: connection,
+          trigger_name: 'trigger_name').and_return(trigger)
+
+      expect(trigger).to receive(:create)
+
+      model.install_sharding_key_assignment_trigger(table: :test_table, sharding_key: :project_id,
+        parent_table: :parent_table, parent_sharding_key: :parent_project_id, foreign_key: :foreign_key,
+        trigger_name: 'trigger_name')
+    end
+  end
+
+  describe '#remove_sharding_key_assignment_trigger' do
+    let(:trigger) { double }
+    let(:connection) { ActiveRecord::Base.connection }
+
+    it do
+      expect(Gitlab::Database::Triggers::AssignDesiredShardingKey).to receive(:new)
+        .with(table: :test_table, sharding_key: :project_id, parent_table: :parent_table,
+          parent_sharding_key: :parent_project_id, foreign_key: :foreign_key, connection: connection,
+          trigger_name: 'trigger_name').and_return(trigger)
+
+      expect(trigger).to receive(:drop)
+
+      model.remove_sharding_key_assignment_trigger(table: :test_table, sharding_key: :project_id,
+        parent_table: :parent_table, parent_sharding_key: :parent_project_id, foreign_key: :foreign_key,
+        trigger_name: 'trigger_name')
+    end
+  end
+
   describe '#indexes_for' do
     it 'returns the indexes for a column' do
       idx1 = double(:idx, columns: %w[project_id])

spec/lib/gitlab/database/triggers/assign_desired_sharding_key_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Triggers::AssignDesiredShardingKey, feature_category: :database do
+  include Database::TriggerHelpers
+
+  let(:table_name) { :_test_table }
+  let(:connection) { ActiveRecord::Base.connection }
+  let(:trigger) { described_class.new(**attributes) }
+  let(:trigger_name) { trigger.name }
+
+  let(:attributes) do
+    {
+      table: table_name,
+      sharding_key: :project_id,
+      parent_table: :_test_project_parent,
+      parent_sharding_key: :parent_project_id,
+      foreign_key: :project_fk_id,
+      connection: connection
+    }
+  end
+
+  describe '#create' do
+    let(:model) { Class.new(ActiveRecord::Base) }
+
+    let(:valid_project_parent_id) { 10 }
+    let(:valid_project_parent_sharding_key) { 20 }
+    let(:invalid_project_parent_id) { 60 }
+
+    subject(:create_trigger) { trigger.create } # rubocop: disable Rails/SaveBang -- Not an ActiveRecord model
+
+    before do
+      connection.execute(<<~SQL)
+        CREATE TABLE _test_project_parent (
+          id serial NOT NULL PRIMARY KEY,
+          parent_project_id bigint);
+
+        CREATE TABLE #{table_name} (
+          id serial NOT NULL PRIMARY KEY,
+          project_fk_id bigint,
+          project_id bigint);
+
+        INSERT INTO _test_project_parent (id, parent_project_id) VALUES
+          (#{valid_project_parent_id}, #{valid_project_parent_sharding_key}),
+          (#{invalid_project_parent_id}, NULL);
+      SQL
+
+      model.table_name = table_name
+    end
+
+    it 'creates the trigger and function' do
+      expect_function_not_to_exist(trigger_name)
+      expect_trigger_not_to_exist(table_name, trigger_name)
+
+      create_trigger
+
+      expect_function_to_exist(trigger_name)
+      expect_valid_function_trigger(table_name, trigger_name, trigger_name, before: %w[insert update])
+    end
+
+    it 'assigns the sharding key using the trigger function' do
+      create_trigger
+
+      record = model.create!(project_fk_id: valid_project_parent_id)
+      expect(record.reload).to have_attributes(project_id: valid_project_parent_sharding_key)
+    end
+
+    context 'when the sharding key is already set' do
+      it 'does not change the sharding key' do
+        create_trigger
+
+        record = model.create!(project_fk_id: valid_project_parent_id, project_id: 99)
+        expect(record.reload).to have_attributes(project_id: 99)
+      end
+    end
+
+    context 'when no matching record is found' do
+      it 'does not set the sharding key' do
+        create_trigger
+
+        record = model.create!(project_fk_id: non_existing_record_id)
+        expect(record.reload).to have_attributes(project_id: nil)
+      end
+    end
+
+    context 'when a matching record is found but the sharding key is missing' do
+      it 'does not set the sharding key' do
+        create_trigger
+
+        record = model.create!(project_fk_id: invalid_project_parent_id)
+        expect(record.reload).to have_attributes(project_id: nil)
+      end
+    end
+
+    context 'when a custom trigger name is supplied' do
+      let(:trigger) { described_class.new(**attributes.merge(trigger_name: trigger_name)) }
+      let(:trigger_name) { 'trigger_with_custom_name' }
+
+      it 'creates the trigger and function using the custom name' do
+        expect_function_not_to_exist(trigger_name)
+        expect_trigger_not_to_exist(table_name, trigger_name)
+
+        create_trigger
+
+        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 already exists' do
+      before do
+        connection.execute(<<~SQL)
+          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 'does not raise an error' do
+        expect_function_to_exist(trigger_name)
+
+        create_trigger
+
+        expect_function_to_exist(trigger_name)
+        expect_valid_function_trigger(table_name, trigger_name, trigger_name, before: %w[insert update])
+      end
+    end
+  end
+
+  describe '#drop' do
+    subject(:drop_trigger) { trigger.drop }
+
+    before do
+      connection.execute(<<~SQL)
+        CREATE TABLE #{table_name} (
+          id serial NOT NULL PRIMARY KEY,
+          project_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])
+
+      drop_trigger
+
+      expect_trigger_not_to_exist(table_name, trigger_name)
+      expect_function_not_to_exist(trigger_name)
+    end
+
+    context 'when the trigger has a custom name' do
+      let(:trigger) { described_class.new(**attributes.merge(trigger_name: trigger_name)) }
+      let(:trigger_name) { 'trigger_with_custom_name' }
+
+      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])
+
+        drop_trigger
+
+        expect_trigger_not_to_exist(table_name, trigger_name)
+        expect_function_not_to_exist(trigger_name)
+      end
+    end
+
+    context 'when the trigger does not exist' do
+      it 'does not raise an error' do
+        drop_trigger
+
+        expect_trigger_not_to_exist(table_name, trigger_name)
+        expect_function_not_to_exist(trigger_name)
+
+        drop_trigger
+      end
+    end
+  end
+end