Use trigger to sync writes to a partitioned table

When creating a partitioned table using migration helpers, attach a
trigger to the source table that will copy all insert/update/delete
operations to the partitioned copy.

lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers.rb

@@ -79,11 +79,7 @@ module Gitlab
         end
 
         def update_foreign_keys(from_table, to_table, from_column, to_column, cascade_delete = nil)
-          if transaction_open?
-            raise 'partitioned foreign key operations can not be run inside a transaction block, ' \
-              'you can disable transaction blocks by calling disable_ddl_transaction! ' \
-              'in the body of your migration class'
-          end
+          assert_not_in_transaction_block(scope: 'partitioned foreign key')
 
           from_column ||= "#{to_table.to_s.singularize}_id"
           specified_key = fk_from_spec(from_table, to_table, from_column, to_column, cascade_delete)
@@ -103,7 +99,7 @@ module Gitlab
               drop_function(fn_name, if_exists: true)
             else
               create_or_replace_fk_function(fn_name, final_keys)
-              create_function_trigger(trigger_name, fn_name, fires: "AFTER DELETE ON #{to_table}")
+              create_trigger(trigger_name, fn_name, fires: "AFTER DELETE ON #{to_table}")
             end
           end
         end
@@ -116,13 +112,6 @@ module Gitlab
           end
         end
 
-        def with_lock_retries(&block)
-          Gitlab::Database::WithLockRetries.new({
-            klass: self.class,
-            logger: Gitlab::BackgroundMigration::Logger
-          }).run(&block)
-        end
-
         def find_existing_key(keys, key)
           keys.find { |k| k.from_table == key.from_table && k.from_column == key.from_column }
         end

lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb

@@ -4,7 +4,10 @@ module Gitlab
   module Database
     module PartitioningMigrationHelpers
       module TableManagementHelpers
-        include SchemaHelpers
+        include ::Gitlab::Database::SchemaHelpers
+
+        WHITELISTED_TABLES = %w[audit_events].freeze
+        ERROR_SCOPE = 'table partitioning'
 
         # Creates a partitioned copy of an existing table, using a RANGE partitioning strategy on a timestamp column.
         # One partition is created per month between the given `min_date` and `max_date`.
@@ -20,6 +23,9 @@ module Gitlab
         #   :max_date - a date specifying the upper bounds of the partitioning range
         #
         def partition_table_by_date(table_name, column_name, min_date:, max_date:)
+          assert_table_is_whitelisted(table_name)
+          assert_not_in_transaction_block(scope: ERROR_SCOPE)
+
           raise "max_date #{max_date} must be greater than min_date #{min_date}" if min_date >= max_date
 
           primary_key = connection.primary_key(table_name)
@@ -31,6 +37,7 @@ module Gitlab
           new_table_name = partitioned_table_name(table_name)
           create_range_partitioned_copy(new_table_name, table_name, partition_column, primary_key)
           create_daterange_partitions(new_table_name, partition_column.name, min_date, max_date)
+          create_sync_trigger(table_name, new_table_name, primary_key)
         end
 
         # Clean up a partitioned copy of an existing table. This deletes the partitioned table and all partitions.
@@ -40,40 +47,76 @@ module Gitlab
         #   drop_partitioned_table_for :audit_events
         #
         def drop_partitioned_table_for(table_name)
-          drop_table(partitioned_table_name(table_name))
+          assert_table_is_whitelisted(table_name)
+          assert_not_in_transaction_block(scope: ERROR_SCOPE)
+
+          with_lock_retries do
+            trigger_name = sync_trigger_name(table_name)
+            drop_trigger(table_name, trigger_name)
+          end
+
+          function_name = sync_function_name(table_name)
+          drop_function(function_name)
+
+          part_table_name = partitioned_table_name(table_name)
+          drop_table(part_table_name)
         end
 
         private
 
+        def assert_table_is_whitelisted(table_name)
+          return if WHITELISTED_TABLES.include?(table_name.to_s)
+
+          raise "partitioning helpers are in active development, and #{table_name} is not whitelisted for use, " \
+            "for more information please contact the database team"
+        end
+
         def partitioned_table_name(table)
           tmp_table_name("#{table}_part")
         end
 
+        def sync_function_name(table)
+          object_name(table, 'table_sync_function')
+        end
+
+        def sync_trigger_name(table)
+          object_name(table, 'table_sync_trigger')
+        end
+
         def find_column_definition(table, column)
           connection.columns(table).find { |c| c.name == column.to_s }
         end
 
         def create_range_partitioned_copy(table_name, template_table_name, partition_column, primary_key)
-          tmp_column_name = object_name(partition_column.name, 'partition_key')
+          if table_exists?(table_name)
+            # rubocop:disable Gitlab/RailsLogger
+            Rails.logger.warn "Partitioned table not created because it already exists" \
+              " (this may be due to an aborted migration or similar): table_name: #{table_name} "
+            # rubocop:enable Gitlab/RailsLogger
+            return
+          end
 
-          execute(<<~SQL)
-            CREATE TABLE #{table_name} (
-              LIKE #{template_table_name} INCLUDING ALL EXCLUDING INDEXES,
-              #{tmp_column_name} #{partition_column.sql_type} NOT NULL,
-              PRIMARY KEY (#{[primary_key, tmp_column_name].join(", ")})
-            ) PARTITION BY RANGE (#{tmp_column_name})
-          SQL
-
-          remove_column(table_name, partition_column.name)
-          rename_column(table_name, tmp_column_name, partition_column.name)
-          change_column_default(table_name, primary_key, nil)
+          tmp_column_name = object_name(partition_column.name, 'partition_key')
+          transaction do
+            execute(<<~SQL)
+              CREATE TABLE #{table_name} (
+                LIKE #{template_table_name} INCLUDING ALL EXCLUDING INDEXES,
+                #{tmp_column_name} #{partition_column.sql_type} NOT NULL,
+                PRIMARY KEY (#{[primary_key, tmp_column_name].join(", ")})
+              ) PARTITION BY RANGE (#{tmp_column_name})
+            SQL
+
+            remove_column(table_name, partition_column.name)
+            rename_column(table_name, tmp_column_name, partition_column.name)
+            change_column_default(table_name, primary_key, nil)
+          end
         end
 
         def create_daterange_partitions(table_name, column_name, min_date, max_date)
           min_date = min_date.beginning_of_month.to_date
           max_date = max_date.next_month.beginning_of_month.to_date
 
-          create_range_partition("#{table_name}_000000", table_name, 'MINVALUE', to_sql_date_literal(min_date))
+          create_range_partition_safely("#{table_name}_000000", table_name, 'MINVALUE', to_sql_date_literal(min_date))
 
           while min_date < max_date
             partition_name = "#{table_name}_#{min_date.strftime('%Y%m')}"
@@ -81,7 +124,7 @@ module Gitlab
             lower_bound = to_sql_date_literal(min_date)
             upper_bound = to_sql_date_literal(next_date)
 
-            create_range_partition(partition_name, table_name, lower_bound, upper_bound)
+            create_range_partition_safely(partition_name, table_name, lower_bound, upper_bound)
             min_date = next_date
           end
         end
@@ -90,11 +133,55 @@ module Gitlab
           connection.quote(date.strftime('%Y-%m-%d'))
         end
 
-        def create_range_partition(partition_name, table_name, lower_bound, upper_bound)
-          execute(<<~SQL)
-            CREATE TABLE #{partition_name} PARTITION OF #{table_name}
-            FOR VALUES FROM (#{lower_bound}) TO (#{upper_bound})
-          SQL
+        def create_range_partition_safely(partition_name, table_name, lower_bound, upper_bound)
+          if table_exists?(partition_name)
+            # rubocop:disable Gitlab/RailsLogger
+            Rails.logger.warn "Partition not created because it already exists" \
+              " (this may be due to an aborted migration or similar): partition_name: #{partition_name}"
+            # rubocop:enable Gitlab/RailsLogger
+            return
+          end
+
+          create_range_partition(partition_name, table_name, lower_bound, upper_bound)
+        end
+
+        def create_sync_trigger(source_table, target_table, unique_key)
+          function_name = sync_function_name(source_table)
+          trigger_name = sync_trigger_name(source_table)
+
+          with_lock_retries do
+            create_sync_function(function_name, target_table, unique_key)
+            create_comment('FUNCTION', function_name, "Partitioning migration: table sync for #{source_table} table")
+
+            create_trigger(trigger_name, function_name, fires: "AFTER INSERT OR UPDATE OR DELETE ON #{source_table}")
+          end
+        end
+
+        def create_sync_function(name, target_table, unique_key)
+          delimiter = ",\n    "
+          column_names = connection.columns(target_table).map(&:name)
+          set_statements = build_set_statements(column_names, unique_key)
+          insert_values = column_names.map { |name| "NEW.#{name}" }
+
+          create_trigger_function(name, replace: false) do
+            <<~SQL
+              IF (TG_OP = 'DELETE') THEN
+                DELETE FROM #{target_table} where #{unique_key} = OLD.#{unique_key};
+              ELSIF (TG_OP = 'UPDATE') THEN
+                UPDATE #{target_table}
+                SET #{set_statements.join(delimiter)}
+                WHERE #{target_table}.#{unique_key} = NEW.#{unique_key};
+              ELSIF (TG_OP = 'INSERT') THEN
+                INSERT INTO #{target_table} (#{column_names.join(delimiter)})
+                VALUES (#{insert_values.join(delimiter)});
+              END IF;
+              RETURN NULL;
+            SQL
+          end
+        end
+
+        def build_set_statements(column_names, unique_key)
+          column_names.reject { |name| name == unique_key }.map { |column_name| "#{column_name} = NEW.#{column_name}" }
         end
       end
     end

lib/gitlab/database/schema_helpers.rb

@@ -16,12 +16,12 @@ module Gitlab
         SQL
       end
 
-      def create_function_trigger(name, fn_name, fires: nil)
+      def create_trigger(name, function_name, fires: nil)
         execute(<<~SQL)
           CREATE TRIGGER #{name}
           #{fires}
           FOR EACH ROW
-          EXECUTE PROCEDURE #{fn_name}()
+          EXECUTE PROCEDURE #{function_name}()
         SQL
       end
 
@@ -35,6 +35,10 @@ module Gitlab
         execute("DROP TRIGGER #{exists_clause} #{name} ON #{table_name}")
       end
 
+      def create_comment(type, name, text)
+        execute("COMMENT ON #{type} #{name} IS '#{text}'")
+      end
+
       def tmp_table_name(base)
         hashed_base = Digest::SHA256.hexdigest(base).first(10)
 
@@ -48,8 +52,30 @@ module Gitlab
         "#{type}_#{hashed_identifier}"
       end
 
+      def with_lock_retries(&block)
+        Gitlab::Database::WithLockRetries.new({
+          klass: self.class,
+          logger: Gitlab::BackgroundMigration::Logger
+        }).run(&block)
+      end
+
+      def assert_not_in_transaction_block(scope:)
+        return unless transaction_open?
+
+        raise "#{scope} operations can not be run inside a transaction block, " \
+          "you can disable transaction blocks by calling disable_ddl_transaction! " \
+          "in the body of your migration class"
+      end
+
       private
 
+      def create_range_partition(partition_name, table_name, lower_bound, upper_bound)
+        execute(<<~SQL)
+          CREATE TABLE #{partition_name} PARTITION OF #{table_name}
+          FOR VALUES FROM (#{lower_bound}) TO (#{upper_bound})
+        SQL
+      end
+
       def optional_clause(flag, clause)
         flag ? clause : ""
       end

spec/lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers_spec.rb

@@ -3,6 +3,8 @@
 require 'spec_helper'
 
 describe Gitlab::Database::PartitioningMigrationHelpers::ForeignKeyHelpers do
+  include TriggerHelpers
+
   let(:model) do
     ActiveRecord::Migration.new.extend(described_class)
   end
@@ -27,7 +29,7 @@ describe Gitlab::Database::PartitioningMigrationHelpers::ForeignKeyHelpers do
         model.add_partitioned_foreign_key :issue_assignees, referenced_table
 
         expect_function_to_contain(function_name, 'delete from issue_assignees where issue_id = old.id')
-        expect_valid_function_trigger(trigger_name, function_name)
+        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
       end
     end
 
@@ -43,7 +45,7 @@ describe Gitlab::Database::PartitioningMigrationHelpers::ForeignKeyHelpers do
           expect_function_to_contain(function_name,
             'delete from issue_assignees where issue_id = old.id',
             'delete from epic_issues where issue_id = old.id')
-          expect_valid_function_trigger(trigger_name, function_name)
+          expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
         end
       end
 
@@ -59,7 +61,7 @@ describe Gitlab::Database::PartitioningMigrationHelpers::ForeignKeyHelpers do
             expect_function_to_contain(function_name,
               'delete from issues where moved_to_id = old.id',
               'delete from issues where duplicated_to_id = old.id')
-            expect_valid_function_trigger(trigger_name, function_name)
+            expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
           end
         end
 
@@ -68,7 +70,7 @@ describe Gitlab::Database::PartitioningMigrationHelpers::ForeignKeyHelpers do
             model.add_partitioned_foreign_key :issues, referenced_table, column: :moved_to_id
 
             expect_function_to_contain(function_name, 'delete from issues where moved_to_id = old.id')
-            expect_valid_function_trigger(trigger_name, function_name)
+            expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
           end
         end
       end
@@ -79,7 +81,7 @@ describe Gitlab::Database::PartitioningMigrationHelpers::ForeignKeyHelpers do
         model.add_partitioned_foreign_key :issue_assignees, referenced_table, on_delete: :nullify
 
         expect_function_to_contain(function_name, 'update issue_assignees set issue_id = null where issue_id = old.id')
-        expect_valid_function_trigger(trigger_name, function_name)
+        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
       end
     end
 
@@ -88,7 +90,7 @@ describe Gitlab::Database::PartitioningMigrationHelpers::ForeignKeyHelpers do
         model.add_partitioned_foreign_key :issues, referenced_table, column: :duplicated_to_id
 
         expect_function_to_contain(function_name, 'delete from issues where duplicated_to_id = old.id')
-        expect_valid_function_trigger(trigger_name, function_name)
+        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
       end
     end
 
@@ -99,7 +101,7 @@ describe Gitlab::Database::PartitioningMigrationHelpers::ForeignKeyHelpers do
         model.add_partitioned_foreign_key :user_preferences, referenced_table, column: :user_id, primary_key: :user_id
 
         expect_function_to_contain(function_name, 'delete from user_preferences where user_id = old.user_id')
-        expect_valid_function_trigger(trigger_name, function_name)
+        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
       end
     end
 
@@ -137,12 +139,12 @@ describe Gitlab::Database::PartitioningMigrationHelpers::ForeignKeyHelpers do
         expect_function_to_contain(function_name,
           'delete from issue_assignees where issue_id = old.id',
           'delete from epic_issues where issue_id = old.id')
-        expect_valid_function_trigger(trigger_name, function_name)
+        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
 
         model.remove_partitioned_foreign_key :issue_assignees, referenced_table
 
         expect_function_to_contain(function_name, 'delete from epic_issues where issue_id = old.id')
-        expect_valid_function_trigger(trigger_name, function_name)
+        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
       end
     end
 
@@ -153,12 +155,12 @@ describe Gitlab::Database::PartitioningMigrationHelpers::ForeignKeyHelpers do
 
       it 'removes the trigger function altogether' do
         expect_function_to_contain(function_name, 'delete from issue_assignees where issue_id = old.id')
-        expect_valid_function_trigger(trigger_name, function_name)
+        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
 
         model.remove_partitioned_foreign_key :issue_assignees, referenced_table
 
-        expect(find_function_def(function_name)).to be_nil
-        expect(find_trigger_def(trigger_name)).to be_nil
+        expect_function_not_to_exist(function_name)
+        expect_trigger_not_to_exist(referenced_table, trigger_name)
       end
     end
 
@@ -169,12 +171,12 @@ describe Gitlab::Database::PartitioningMigrationHelpers::ForeignKeyHelpers do
 
       it 'ignores the invalid key and properly recreates the trigger function' do
         expect_function_to_contain(function_name, 'delete from issue_assignees where issue_id = old.id')
-        expect_valid_function_trigger(trigger_name, function_name)
+        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
 
         model.remove_partitioned_foreign_key :issues, referenced_table, column: :moved_to_id
 
         expect_function_to_contain(function_name, 'delete from issue_assignees where issue_id = old.id')
-        expect_valid_function_trigger(trigger_name, function_name)
+        expect_valid_function_trigger(referenced_table, trigger_name, function_name, after: 'delete')
       end
     end
 
@@ -188,45 +190,4 @@ describe Gitlab::Database::PartitioningMigrationHelpers::ForeignKeyHelpers do
       end
     end
   end
-
-  def expect_function_to_contain(name, *statements)
-    return_stmt, *body_stmts = parsed_function_statements(name).reverse
-
-    expect(return_stmt).to eq('return old')
-    expect(body_stmts).to contain_exactly(*statements)
-  end
-
-  def expect_valid_function_trigger(name, fn_name)
-    event, activation, definition = cleaned_trigger_def(name)
-
-    expect(event).to eq('delete')
-    expect(activation).to eq('after')
-    expect(definition).to eq("execute procedure #{fn_name}()")
-  end
-
-  def parsed_function_statements(name)
-    cleaned_definition = find_function_def(name)['fn_body'].downcase.gsub(/\s+/, ' ')
-    statements = cleaned_definition.sub(/\A\s*begin\s*(.*)\s*end\s*\Z/, "\\1")
-    statements.split(';').map! { |stmt| stmt.strip.presence }.compact!
-  end
-
-  def find_function_def(name)
-    connection.execute("select prosrc as fn_body from pg_proc where proname = '#{name}';").first
-  end
-
-  def cleaned_trigger_def(name)
-    find_trigger_def(name).values_at('event', 'activation', 'definition').map!(&:downcase)
-  end
-
-  def find_trigger_def(name)
-    connection.execute(<<~SQL).first
-      select
-        string_agg(event_manipulation, ',') as event,
-        action_timing as activation,
-        action_statement as definition
-      from information_schema.triggers
-      where trigger_name = '#{name}'
-      group by 2, 3
-    SQL
-  end
 end

spec/support/helpers/trigger_helpers.rb

+# frozen_string_literal: true
+
+module TriggerHelpers
+  def expect_function_to_exist(name)
+    expect(find_function_def(name)).not_to be_nil
+  end
+
+  def expect_function_not_to_exist(name)
+    expect(find_function_def(name)).to be_nil
+  end
+
+  def expect_function_to_contain(name, *statements)
+    return_stmt, *body_stmts = parsed_function_statements(name).reverse
+
+    expect(return_stmt).to eq('return old')
+    expect(body_stmts).to contain_exactly(*statements)
+  end
+
+  def expect_trigger_not_to_exist(table_name, name)
+    expect(find_trigger_def(table_name, name)).to be_nil
+  end
+
+  def expect_valid_function_trigger(table_name, name, fn_name, fires_on)
+    events, timing, definition = cleaned_trigger_def(table_name, name)
+
+    events = events&.split(',')
+    expected_timing, expected_events = fires_on.first
+    expect(timing).to eq(expected_timing.to_s)
+    expect(events).to match_array(Array.wrap(expected_events))
+    expect(definition).to eq("execute procedure #{fn_name}()")
+  end
+
+  private
+
+  def parsed_function_statements(name)
+    cleaned_definition = find_function_def(name)['body'].downcase.gsub(/\s+/, ' ')
+    statements = cleaned_definition.sub(/\A\s*begin\s*(.*)\s*end\s*\Z/, "\\1")
+    statements.split(';').map! { |stmt| stmt.strip.presence }.compact!
+  end
+
+  def find_function_def(name)
+    connection.select_one(<<~SQL)
+      SELECT prosrc AS body
+      FROM pg_proc
+      WHERE proname = '#{name}'
+    SQL
+  end
+
+  def cleaned_trigger_def(table_name, name)
+    find_trigger_def(table_name, name).values_at('event', 'action_timing', 'action_statement').map!(&:downcase)
+  end
+
+  def find_trigger_def(table_name, name)
+    connection.select_one(<<~SQL)
+      SELECT
+        string_agg(event_manipulation, ',') AS event,
+        action_timing,
+        action_statement
+      FROM information_schema.triggers
+      WHERE event_object_table = '#{table_name}'
+      AND trigger_name = '#{name}'
+      GROUP BY 2, 3
+    SQL
+  end
+end