Add latest changes from gitlab-org/gitlab@master

2 files changed, 82 insertions, 36 deletions

diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index c3a7b384172..82adb0dd369 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -292,23 +292,34 @@ module Gitlab
       #                      order of the ALTER TABLE. This can be useful in situations where the foreign
       #                      key creation could deadlock with another process.
       #
-      # rubocop: disable Metrics/ParameterLists
-      def add_concurrent_foreign_key(source, target, column:, on_delete: :cascade, on_update: nil, target_column: :id, name: nil, validate: true, reverse_lock_order: false)
+      def add_concurrent_foreign_key(source, target, column:, **options)
+        options.reverse_merge!({
+          on_delete: :cascade,
+          on_update: nil,
+          target_column: :id,
+          validate: true,
+          reverse_lock_order: false,
+          allow_partitioned: false,
+          column: column
+        })
+
         # Transactions would result in ALTER TABLE locks being held for the
         # duration of the transaction, defeating the purpose of this method.
         if transaction_open?
           raise 'add_concurrent_foreign_key can not be run inside a transaction'
         end
 
-        options = {
-          column: column,
-          on_delete: on_delete,
-          on_update: on_update,
-          name: name.presence || concurrent_foreign_key_name(source, column),
-          primary_key: target_column
-        }
+        if !options.delete(:allow_partitioned) && table_partitioned?(source)
+          raise ArgumentError, 'add_concurrent_foreign_key can not be used on a partitioned ' \
+            'table. Please use add_concurrent_partitioned_foreign_key on the partitioned table ' \
+            'as we need to create foreign keys on each partition and a FK on the parent table'
+        end
 
-        if foreign_key_exists?(source, target, **options)
+        options[:name] ||= concurrent_foreign_key_name(source, column)
+        options[:primary_key] = options[:target_column]
+        check_options = options.slice(:column, :on_delete, :on_update, :name, :primary_key)
+
+        if foreign_key_exists?(source, target, **check_options)
           warning_message = "Foreign key not created because it exists already " \
             "(this may be due to an aborted migration or similar): " \
             "source: #{source}, target: #{target}, column: #{options[:column]}, "\
@@ -321,17 +332,18 @@ module Gitlab
           # validating it. This means we keep the ALTER TABLE lock only for a
           # short period of time. The key _is_ enforced for any newly created
           # data.
+          not_valid_option = 'NOT VALID' unless table_partitioned?(source)
 
           with_lock_retries do
-            execute("LOCK TABLE #{target}, #{source} IN SHARE ROW EXCLUSIVE MODE") if reverse_lock_order
+            execute("LOCK TABLE #{target}, #{source} IN SHARE ROW EXCLUSIVE MODE") if options[:reverse_lock_order]
             execute <<-EOF.strip_heredoc
             ALTER TABLE #{source}
             ADD CONSTRAINT #{options[:name]}
             FOREIGN KEY (#{multiple_columns(options[:column])})
-            REFERENCES #{target} (#{multiple_columns(target_column)})
+            REFERENCES #{target} (#{multiple_columns(options[:target_column])})
             #{on_update_statement(options[:on_update])}
             #{on_delete_statement(options[:on_delete])}
-            NOT VALID;
+            #{not_valid_option};
             EOF
           end
         end
@@ -345,13 +357,12 @@ module Gitlab
         #
         # Note this is a no-op in case the constraint is VALID already
 
-        if validate
+        if options[:validate]
           disable_statement_timeout do
             execute("ALTER TABLE #{source} VALIDATE CONSTRAINT #{options[:name]};")
           end
         end
       end
-      # rubocop: enable Metrics/ParameterLists
 
       def validate_foreign_key(source, column, name: nil)
         fk_name = name || concurrent_foreign_key_name(source, column)
@@ -1239,6 +1250,12 @@ into similar problems in the future (e.g. when new tables are created).
         end
       end
 
+      def table_partitioned?(table_name)
+        Gitlab::Database::PostgresPartitionedTable
+          .find_by_name_in_current_schema(table_name)
+          .present?
+      end
+
       private
 
       def multiple_columns(columns, separator: ', ')
diff --git a/lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers.rb b/lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers.rb
index 8849191f356..7d9c12d776e 100644
--- a/lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers.rb
+++ b/lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers.rb
@@ -32,46 +32,75 @@ module Gitlab
         # column - The name of the column to create the foreign key on.
         # on_delete - The action to perform when associated data is removed,
         #             defaults to "CASCADE".
+        # on_update - The action to perform when associated data is updated,
+        #             no default value is set.
         # name - The name of the foreign key.
+        # validate - Flag that controls whether the new foreign key will be
+        #            validated after creation and if it will be added on the parent table.
+        #            If the flag is not set, the constraint will only be enforced for new data
+        #            in the existing partitions. The helper will need to be called again
+        #            with the flag set to `true` to add the foreign key on the parent table
+        #            after validating it on all partitions.
+        #            `validate: false` should be paired with `prepare_partitioned_async_foreign_key_validation`
+        # reverse_lock_order - Flag that controls whether we should attempt to acquire locks in the reverse
+        #                      order of the ALTER TABLE. This can be useful in situations where the foreign
+        #                      key creation could deadlock with another process.
         #
-        def add_concurrent_partitioned_foreign_key(source, target, column:, on_delete: :cascade, name: nil)
+        def add_concurrent_partitioned_foreign_key(source, target, column:, **options)
           assert_not_in_transaction_block(scope: ERROR_SCOPE)
 
-          partition_options = {
-            column: column,
-            on_delete: on_delete,
+          options.reverse_merge!({
+            target_column: :id,
+            on_delete: :cascade,
+            on_update: nil,
+            name: nil,
+            validate: true,
+            reverse_lock_order: false,
+            column: column
+          })
 
-            # We'll use the same FK name for all partitions and match it to
-            # the name used for the partitioned table to follow the convention
-            # used by PostgreSQL when adding FKs to new partitions
-            name: name.presence || concurrent_partitioned_foreign_key_name(source, column),
+          # We'll use the same FK name for all partitions and match it to
+          # the name used for the partitioned table to follow the convention
+          # used by PostgreSQL when adding FKs to new partitions
+          options[:name] ||= concurrent_partitioned_foreign_key_name(source, column)
+          check_options = options.slice(:column, :on_delete, :on_update, :name)
+          check_options[:primary_key] = options[:target_column]
 
-            # Force the FK validation to true for partitions (and the partitioned table)
-            validate: true
-          }
-
-          if foreign_key_exists?(source, target, **partition_options)
+          if foreign_key_exists?(source, target, **check_options)
             warning_message = "Foreign key not created because it exists already " \
               "(this may be due to an aborted migration or similar): " \
-              "source: #{source}, target: #{target}, column: #{partition_options[:column]}, "\
-              "name: #{partition_options[:name]}, on_delete: #{partition_options[:on_delete]}"
+              "source: #{source}, target: #{target}, column: #{options[:column]}, "\
+              "name: #{options[:name]}, on_delete: #{options[:on_delete]}, "\
+              "on_update: #{options[:on_update]}"
 
             Gitlab::AppLogger.warn warning_message
 
             return
           end
 
-          partitioned_table = find_partitioned_table(source)
-
-          partitioned_table.postgres_partitions.order(:name).each do |partition|
-            add_concurrent_foreign_key(partition.identifier, target, **partition_options)
+          Gitlab::Database::PostgresPartitionedTable.each_partition(source) do |partition|
+            add_concurrent_foreign_key(partition.identifier, target, **options)
           end
 
-          with_lock_retries do
-            add_foreign_key(source, target, **partition_options)
+          # If we are to add the FK on the parent table now, it will trigger
+          # the validation on all partitions. The helper must be called one
+          # more time with `validate: true` after the FK is valid on all partitions.
+          return unless options[:validate]
+
+          options[:allow_partitioned] = true
+          add_concurrent_foreign_key(source, target, **options)
+        end
+
+        def validate_partitioned_foreign_key(source, column, name: nil)
+          assert_not_in_transaction_block(scope: ERROR_SCOPE)
+
+          Gitlab::Database::PostgresPartitionedTable.each_partition(source) do |partition|
+            validate_foreign_key(partition.identifier, column, name: name)
           end
         end
 
+        private
+
         # Returns the name for a concurrent partitioned foreign key.
         #
         # Similar to concurrent_foreign_key_name (Gitlab::Database::MigrationHelpers)