partition_manager.rb 4.51 KB
Newer Older
1 2 3 4 5
# frozen_string_literal: true

module Gitlab
  module Database
    module Partitioning
6
      class PartitionManager
7 8
        UnsafeToDetachPartitionError = Class.new(StandardError)

9
        LEASE_TIMEOUT = 1.minute
10
        MANAGEMENT_LEASE_KEY = 'database_partition_management_%s'
11
        RETAIN_DETACHED_PARTITIONS_FOR = 1.week
12

13 14
        def initialize(model)
          @model = model
15 16
        end

17
        def sync_partitions
18
          Gitlab::AppLogger.info(message: "Checking state of dynamic postgres partitions", table_name: model.table_name)
19

20 21 22
          # Double-checking before getting the lease:
          # The prevailing situation is no missing partitions and no extra partitions
          return if missing_partitions.empty? && extra_partitions.empty?
23

24 25 26
          only_with_exclusive_lease(model, lease_key: MANAGEMENT_LEASE_KEY) do
            partitions_to_create = missing_partitions
            create(partitions_to_create) unless partitions_to_create.empty?
27

28 29 30
            if Feature.enabled?(:partition_pruning, default_enabled: :yaml)
              partitions_to_detach = extra_partitions
              detach(partitions_to_detach) unless partitions_to_detach.empty?
31 32
            end
          end
33 34 35 36 37
        rescue StandardError => e
          Gitlab::AppLogger.error(message: "Failed to create / detach partition(s)",
                                  table_name: model.table_name,
                                  exception_class: e.class,
                                  exception_message: e.message)
38 39
        end

40 41
        private

42 43
        attr_reader :model
        delegate :connection, to: :model
44

45
        def missing_partitions
46 47 48 49 50
          return [] unless connection.table_exists?(model.table_name)

          model.partitioning_strategy.missing_partitions
        end

51
        def extra_partitions
52 53 54 55 56
          return [] unless connection.table_exists?(model.table_name)

          model.partitioning_strategy.extra_partitions
        end

57 58 59 60 61 62
        def only_with_exclusive_lease(model, lease_key:)
          lease = Gitlab::ExclusiveLease.new(lease_key % model.table_name, timeout: LEASE_TIMEOUT)

          yield if lease.try_obtain
        ensure
          lease&.cancel
63 64
        end

65
        def create(partitions)
66 67 68
          # with_lock_retries starts a requires_new transaction most of the time, but not on the last iteration
          with_lock_retries do
            connection.transaction(requires_new: false) do # so we open a transaction here if not already in progress
69 70 71
              partitions.each do |partition|
                connection.execute partition.to_sql

72 73 74
                Gitlab::AppLogger.info(message: "Created partition",
                                       partition_name: partition.partition_name,
                                       table_name: partition.table)
75
              end
76

77
              model.partitioning_strategy.after_adding_partitions
78 79 80 81
            end
          end
        end

82
        def detach(partitions)
83 84 85
          # with_lock_retries starts a requires_new transaction most of the time, but not on the last iteration
          with_lock_retries do
            connection.transaction(requires_new: false) do # so we open a transaction here if not already in progress
86 87 88 89 90 91
              partitions.each { |p| detach_one_partition(p) }
            end
          end
        end

        def detach_one_partition(partition)
92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109
          assert_partition_detachable!(partition)

          connection.execute partition.to_detach_sql

          Postgresql::DetachedPartition.create!(table_name: partition.partition_name,
                                                drop_after: RETAIN_DETACHED_PARTITIONS_FOR.from_now)

          Gitlab::AppLogger.info(message: "Detached Partition",
                                 partition_name: partition.partition_name,
                                 table_name: partition.table)
        end

        def assert_partition_detachable!(partition)
          parent_table_identifier = "#{connection.current_schema}.#{partition.table}"

          if (example_fk = PostgresForeignKey.by_referenced_table_identifier(parent_table_identifier).first)
            raise UnsafeToDetachPartitionError, "Cannot detach #{partition.partition_name}, it would block while checking foreign key #{example_fk.name} on #{example_fk.constrained_table_identifier}"
          end
110 111
        end

112
        def with_lock_retries(&block)
113
          Gitlab::Database::WithLockRetries.new(
114
            klass: self.class,
115 116
            logger: Gitlab::AppLogger,
            connection: connection
117
          ).run(&block)
118 119 120 121 122
        end
      end
    end
  end
end