Add latest changes from gitlab-org/gitlab@13-6-stable-eev13.6.0-rc42

9 files changed, 350 insertions, 30 deletions

diff --git a/lib/gitlab/database/batch_count.rb b/lib/gitlab/database/batch_count.rb
index 11d9881aac2..6f79e965cd5 100644
--- a/lib/gitlab/database/batch_count.rb
+++ b/lib/gitlab/database/batch_count.rb
@@ -128,9 +128,9 @@ module Gitlab
       end
 
       def between_condition(start, finish)
-        return @column.between(start..(finish - 1)) if @column.is_a?(Arel::Attributes::Attribute)
+        return @column.between(start...finish) if @column.is_a?(Arel::Attributes::Attribute)
 
-        { @column => start..(finish - 1) }
+        { @column => start...finish }
       end
 
       def actual_start(start)
diff --git a/lib/gitlab/database/partitioning/monthly_strategy.rb b/lib/gitlab/database/partitioning/monthly_strategy.rb
index ecc05d9654a..82ea1ce26fb 100644
--- a/lib/gitlab/database/partitioning/monthly_strategy.rb
+++ b/lib/gitlab/database/partitioning/monthly_strategy.rb
@@ -17,23 +17,8 @@ module Gitlab
         end
 
         def current_partitions
-          result = connection.select_all(<<~SQL)
-            select
-              pg_class.relname,
-              parent_class.relname as base_table,
-              pg_get_expr(pg_class.relpartbound, inhrelid) as condition
-            from pg_class
-            inner join pg_inherits i on pg_class.oid = inhrelid
-            inner join pg_class parent_class on parent_class.oid = inhparent
-            inner join pg_namespace ON pg_namespace.oid = pg_class.relnamespace
-            where pg_namespace.nspname = #{connection.quote(Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)}
-              and parent_class.relname = #{connection.quote(table_name)}
-              and pg_class.relispartition
-            order by pg_class.relname
-          SQL
-
-          result.map do |record|
-            TimePartition.from_sql(table_name, record['relname'], record['condition'])
+          Gitlab::Database::PostgresPartition.for_parent_table(table_name).map do |partition|
+            TimePartition.from_sql(table_name, partition.name, partition.condition)
           end
         end
 
diff --git a/lib/gitlab/database/partitioning/replace_table.rb b/lib/gitlab/database/partitioning/replace_table.rb
new file mode 100644
index 00000000000..6f6af223fa2
--- /dev/null
+++ b/lib/gitlab/database/partitioning/replace_table.rb
@@ -0,0 +1,114 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Partitioning
+      class ReplaceTable
+        DELIMITER = ";\n\n"
+
+        attr_reader :original_table, :replacement_table, :replaced_table, :primary_key_column,
+          :sequence, :original_primary_key, :replacement_primary_key, :replaced_primary_key
+
+        def initialize(original_table, replacement_table, replaced_table, primary_key_column)
+          @original_table = original_table
+          @replacement_table = replacement_table
+          @replaced_table = replaced_table
+          @primary_key_column = primary_key_column
+
+          @sequence = default_sequence(original_table, primary_key_column)
+          @original_primary_key = default_primary_key(original_table)
+          @replacement_primary_key = default_primary_key(replacement_table)
+          @replaced_primary_key = default_primary_key(replaced_table)
+        end
+
+        def perform
+          yield sql_to_replace_table if block_given?
+
+          execute(sql_to_replace_table)
+        end
+
+        private
+
+        delegate :execute, :quote_table_name, :quote_column_name, to: :connection
+        def connection
+          @connection ||= ActiveRecord::Base.connection
+        end
+
+        def default_sequence(table, column)
+          "#{table}_#{column}_seq"
+        end
+
+        def default_primary_key(table)
+          "#{table}_pkey"
+        end
+
+        def sql_to_replace_table
+          @sql_to_replace_table ||= combined_sql_statements.map(&:chomp).join(DELIMITER)
+        end
+
+        def combined_sql_statements
+          statements = []
+
+          statements << alter_column_default(original_table, primary_key_column, expression: nil)
+          statements << alter_column_default(replacement_table, primary_key_column,
+              expression: "nextval('#{quote_table_name(sequence)}'::regclass)")
+
+          statements << alter_sequence_owned_by(sequence, replacement_table, primary_key_column)
+
+          rename_table_objects(statements, original_table, replaced_table, original_primary_key, replaced_primary_key)
+          rename_table_objects(statements, replacement_table, original_table, replacement_primary_key, original_primary_key)
+
+          statements
+        end
+
+        def rename_table_objects(statements, old_table, new_table, old_primary_key, new_primary_key)
+          statements << rename_table(old_table, new_table)
+          statements << rename_constraint(new_table, old_primary_key, new_primary_key)
+
+          rename_partitions(statements, old_table, new_table)
+        end
+
+        def rename_partitions(statements, old_table_name, new_table_name)
+          Gitlab::Database::PostgresPartition.for_parent_table(old_table_name).each do |partition|
+            new_partition_name = partition.name.sub(/#{old_table_name}/, new_table_name)
+            old_primary_key = default_primary_key(partition.name)
+            new_primary_key = default_primary_key(new_partition_name)
+
+            statements << rename_constraint(partition.identifier, old_primary_key, new_primary_key)
+            statements << rename_table(partition.identifier, new_partition_name)
+          end
+        end
+
+        def alter_column_default(table_name, column_name, expression:)
+          default_clause = expression.nil? ? 'DROP DEFAULT' : "SET DEFAULT #{expression}"
+
+          <<~SQL
+            ALTER TABLE #{quote_table_name(table_name)}
+            ALTER COLUMN #{quote_column_name(column_name)} #{default_clause}
+          SQL
+        end
+
+        def alter_sequence_owned_by(sequence_name, table_name, column_name)
+          <<~SQL
+            ALTER SEQUENCE #{quote_table_name(sequence_name)}
+            OWNED BY #{quote_table_name(table_name)}.#{quote_column_name(column_name)}
+          SQL
+        end
+
+        def rename_table(old_name, new_name)
+          <<~SQL
+            ALTER TABLE #{quote_table_name(old_name)}
+            RENAME TO #{quote_table_name(new_name)}
+          SQL
+        end
+
+        def rename_constraint(table_name, old_name, new_name)
+          <<~SQL
+            ALTER TABLE #{quote_table_name(table_name)}
+            RENAME CONSTRAINT #{quote_column_name(old_name)} TO #{quote_column_name(new_name)}
+          SQL
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/partitioning_migration_helpers.rb b/lib/gitlab/database/partitioning_migration_helpers.rb
index 881177a195e..3196dd20356 100644
--- a/lib/gitlab/database/partitioning_migration_helpers.rb
+++ b/lib/gitlab/database/partitioning_migration_helpers.rb
@@ -5,6 +5,7 @@ module Gitlab
     module PartitioningMigrationHelpers
       include ForeignKeyHelpers
       include TableManagementHelpers
+      include IndexHelpers
     end
   end
 end
diff --git a/lib/gitlab/database/partitioning_migration_helpers/index_helpers.rb b/lib/gitlab/database/partitioning_migration_helpers/index_helpers.rb
new file mode 100644
index 00000000000..f367292f4b0
--- /dev/null
+++ b/lib/gitlab/database/partitioning_migration_helpers/index_helpers.rb
@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module PartitioningMigrationHelpers
+      module IndexHelpers
+        include Gitlab::Database::MigrationHelpers
+        include Gitlab::Database::SchemaHelpers
+
+        # Concurrently creates a new index on a partitioned table. In concept this works similarly to
+        # `add_concurrent_index`, and won't block reads or writes on the table while the index is being built.
+        #
+        # A special helper is required for partitioning because Postgres does not support concurrently building indexes
+        # on partitioned tables. This helper concurrently adds the same index to each partition, and creates the final
+        # index on the parent table once all of the partitions are indexed. This is the recommended safe way to add
+        # indexes to partitioned tables.
+        #
+        # Example:
+        #
+        #     add_concurrent_partitioned_index :users, :some_column
+        #
+        # See Rails' `add_index` for more info on the available arguments.
+        def add_concurrent_partitioned_index(table_name, column_names, options = {})
+          raise ArgumentError, 'A name is required for indexes added to partitioned tables' unless options[:name]
+
+          partitioned_table = find_partitioned_table(table_name)
+
+          if index_name_exists?(table_name, options[:name])
+            Gitlab::AppLogger.warn "Index not created because it already exists (this may be due to an aborted" \
+              " migration or similar): table_name: #{table_name}, index_name: #{options[:name]}"
+
+            return
+          end
+
+          partitioned_table.postgres_partitions.each do |partition|
+            partition_index_name = generated_index_name(partition.identifier, options[:name])
+            partition_options = options.merge(name: partition_index_name)
+
+            add_concurrent_index(partition.identifier, column_names, partition_options)
+          end
+
+          with_lock_retries do
+            add_index(table_name, column_names, options)
+          end
+        end
+
+        # Safely removes an existing index from a partitioned table. The method name is a bit inaccurate as it does not
+        # drop the index concurrently, but it's named as such to maintain consistency with other similar helpers, and
+        # indicate that this should be safe to use in a production environment.
+        #
+        # In current versions of Postgres it's impossible to drop an index concurrently, or drop an index from an
+        # individual partition that exists across the entire partitioned table. As a result this helper drops the index
+        # from the parent table, which automatically cascades to all partitions. While this does require an exclusive
+        # lock, dropping an index is a fast operation that won't block the table for a significant period of time.
+        #
+        # Example:
+        #
+        #     remove_concurrent_partitioned_index_by_name :users, 'index_name_goes_here'
+        def remove_concurrent_partitioned_index_by_name(table_name, index_name)
+          find_partitioned_table(table_name)
+
+          unless index_name_exists?(table_name, index_name)
+            Gitlab::AppLogger.warn "Index not removed because it does not exist (this may be due to an aborted " \
+              "migration or similar): table_name: #{table_name}, index_name: #{index_name}"
+
+            return
+          end
+
+          with_lock_retries do
+            remove_index(table_name, name: index_name)
+          end
+        end
+
+        private
+
+        def find_partitioned_table(table_name)
+          partitioned_table = Gitlab::Database::PostgresPartitionedTable.find_by_name_in_current_schema(table_name)
+
+          raise ArgumentError, "#{table_name} is not a partitioned table" unless partitioned_table
+
+          partitioned_table
+        end
+
+        def generated_index_name(partition_name, index_name)
+          object_name("#{partition_name}_#{index_name}", 'index')
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb b/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb
index f7b0306b769..686dda80207 100644
--- a/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb
+++ b/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb
@@ -66,7 +66,10 @@ module Gitlab
             create_range_partitioned_copy(table_name, partitioned_table_name, partition_column, primary_key)
             create_daterange_partitions(partitioned_table_name, partition_column.name, min_date, max_date)
           end
-          create_trigger_to_sync_tables(table_name, partitioned_table_name, primary_key)
+
+          with_lock_retries do
+            create_trigger_to_sync_tables(table_name, partitioned_table_name, primary_key)
+          end
         end
 
         # Clean up a partitioned copy of an existing table. First, deletes the database function and trigger that were
@@ -81,13 +84,9 @@ module Gitlab
           assert_not_in_transaction_block(scope: ERROR_SCOPE)
 
           with_lock_retries do
-            trigger_name = make_sync_trigger_name(table_name)
-            drop_trigger(table_name, trigger_name)
+            drop_sync_trigger(table_name)
           end
 
-          function_name = make_sync_function_name(table_name)
-          drop_function(function_name)
-
           partitioned_table_name = make_partitioned_table_name(table_name)
           drop_table(partitioned_table_name)
         end
@@ -177,6 +176,53 @@ module Gitlab
           end
         end
 
+        # Replaces a non-partitioned table with its partitioned copy. This is the final step in a partitioning
+        # migration, which makes the partitioned table ready for use by the application. The partitioned copy should be
+        # replaced with the original table in such a way that it appears seamless to any database clients. The replaced
+        # table will be renamed to "#{replaced_table}_archived". Partitions and primary key constraints will also be
+        # renamed to match the naming scheme of the parent table.
+        #
+        # **NOTE** This method should only be used after all other migration steps have completed successfully.
+        # There are several limitations to this method that MUST be handled before, or during, the swap migration:
+        #
+        # - Secondary indexes and foreign keys are not automatically recreated on the partitioned table.
+        # - Some types of constraints (UNIQUE and EXCLUDE) which rely on indexes, will not automatically be recreated
+        #   on the partitioned table, since the underlying index will not be present.
+        # - Foreign keys referencing the original non-partitioned table, would also need to be updated to reference the
+        #   partitioned table, but unfortunately this is not supported in PG11.
+        # - Views referencing the original table will not be automatically updated to reference the partitioned table.
+        #
+        # Example:
+        #
+        #   replace_with_partitioned_table :audit_events
+        #
+        def replace_with_partitioned_table(table_name)
+          assert_table_is_allowed(table_name)
+
+          partitioned_table_name = make_partitioned_table_name(table_name)
+          archived_table_name = make_archived_table_name(table_name)
+          primary_key_name = connection.primary_key(table_name)
+
+          replace_table(table_name, partitioned_table_name, archived_table_name, primary_key_name)
+        end
+
+        # Rolls back a migration that replaced a non-partitioned table with its partitioned copy. This can be used to
+        # restore the original non-partitioned table in the event of an unexpected issue.
+        #
+        # Example:
+        #
+        #   rollback_replace_with_partitioned_table :audit_events
+        #
+        def rollback_replace_with_partitioned_table(table_name)
+          assert_table_is_allowed(table_name)
+
+          partitioned_table_name = make_partitioned_table_name(table_name)
+          archived_table_name = make_archived_table_name(table_name)
+          primary_key_name = connection.primary_key(archived_table_name)
+
+          replace_table(table_name, archived_table_name, partitioned_table_name, primary_key_name)
+        end
+
         private
 
         def assert_table_is_allowed(table_name)
@@ -190,6 +236,10 @@ module Gitlab
           tmp_table_name("#{table}_part")
         end
 
+        def make_archived_table_name(table)
+          "#{table}_archived"
+        end
+
         def make_sync_function_name(table)
           object_name(table, 'table_sync_function')
         end
@@ -270,12 +320,18 @@ module Gitlab
           function_name = make_sync_function_name(source_table_name)
           trigger_name = make_sync_trigger_name(source_table_name)
 
-          with_lock_retries do
-            create_sync_function(function_name, partitioned_table_name, unique_key)
-            create_comment('FUNCTION', function_name, "Partitioning migration: table sync for #{source_table_name} table")
+          create_sync_function(function_name, partitioned_table_name, unique_key)
+          create_comment('FUNCTION', function_name, "Partitioning migration: table sync for #{source_table_name} table")
 
-            create_sync_trigger(source_table_name, trigger_name, function_name)
-          end
+          create_sync_trigger(source_table_name, trigger_name, function_name)
+        end
+
+        def drop_sync_trigger(source_table_name)
+          trigger_name = make_sync_trigger_name(source_table_name)
+          drop_trigger(source_table_name, trigger_name)
+
+          function_name = make_sync_function_name(source_table_name)
+          drop_function(function_name)
         end
 
         def create_sync_function(name, partitioned_table_name, unique_key)
@@ -358,6 +414,21 @@ module Gitlab
             end
           end
         end
+
+        def replace_table(original_table_name, replacement_table_name, replaced_table_name, primary_key_name)
+          replace_table = Gitlab::Database::Partitioning::ReplaceTable.new(original_table_name.to_s,
+              replacement_table_name, replaced_table_name, primary_key_name)
+
+          with_lock_retries do
+            drop_sync_trigger(original_table_name)
+
+            replace_table.perform do |sql|
+              say("replace_table(\"#{sql}\")")
+            end
+
+            create_trigger_to_sync_tables(original_table_name, replaced_table_name, primary_key_name)
+          end
+        end
       end
     end
   end
diff --git a/lib/gitlab/database/postgres_partition.rb b/lib/gitlab/database/postgres_partition.rb
new file mode 100644
index 00000000000..0986372586b
--- /dev/null
+++ b/lib/gitlab/database/postgres_partition.rb
@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    class PostgresPartition < ActiveRecord::Base
+      self.primary_key = :identifier
+
+      belongs_to :postgres_partitioned_table, foreign_key: 'parent_identifier', primary_key: 'identifier'
+
+      scope :by_identifier, ->(identifier) do
+        raise ArgumentError, "Partition name is not fully qualified with a schema: #{identifier}" unless identifier =~ /^\w+\.\w+$/
+
+        find(identifier)
+      end
+
+      scope :for_parent_table, ->(name) { where("parent_identifier = concat(current_schema(), '.', ?)", name).order(:name) }
+
+      def to_s
+        name
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/postgres_partitioned_table.rb b/lib/gitlab/database/postgres_partitioned_table.rb
new file mode 100644
index 00000000000..5d2eaa22ee4
--- /dev/null
+++ b/lib/gitlab/database/postgres_partitioned_table.rb
@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    class PostgresPartitionedTable < ActiveRecord::Base
+      DYNAMIC_PARTITION_STRATEGIES = %w[range list].freeze
+
+      self.primary_key = :identifier
+
+      has_many :postgres_partitions, foreign_key: 'parent_identifier', primary_key: 'identifier'
+
+      scope :by_identifier, ->(identifier) do
+        raise ArgumentError, "Table name is not fully qualified with a schema: #{identifier}" unless identifier =~ /^\w+\.\w+$/
+
+        find(identifier)
+      end
+
+      def self.find_by_name_in_current_schema(name)
+        find_by("identifier = concat(current_schema(), '.', ?)", name)
+      end
+
+      def dynamic?
+        DYNAMIC_PARTITION_STRATEGIES.include?(strategy)
+      end
+
+      def static?
+        !dynamic?
+      end
+
+      def to_s
+        name
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/database/reindexing.rb b/lib/gitlab/database/reindexing.rb
index 074752fe75b..c77e000254f 100644
--- a/lib/gitlab/database/reindexing.rb
+++ b/lib/gitlab/database/reindexing.rb
@@ -10,6 +10,7 @@ module Gitlab
       def self.candidate_indexes
         Gitlab::Database::PostgresIndex
           .regular
+          .where('NOT expression')
           .not_match("^#{ConcurrentReindex::TEMPORARY_INDEX_PREFIX}")
           .not_match("^#{ConcurrentReindex::REPLACED_INDEX_PREFIX}")
       end