2020-07-28 23:09:34 +05:30
# frozen_string_literal: true
module Gitlab
module Database
module Partitioning
2021-09-30 23:02:18 +05:30
class PartitionManager
2021-10-27 15:23:28 +05:30
UnsafeToDetachPartitionError = Class . new ( StandardError )
2020-07-28 23:09:34 +05:30
def self . register ( model )
raise ArgumentError , " Only models with a # partitioning_strategy can be registered. " unless model . respond_to? ( :partitioning_strategy )
models << model
end
def self . models
@models || = Set . new
end
LEASE_TIMEOUT = 1 . minute
2021-09-30 23:02:18 +05:30
MANAGEMENT_LEASE_KEY = 'database_partition_management_%s'
2021-10-27 15:23:28 +05:30
RETAIN_DETACHED_PARTITIONS_FOR = 1 . week
2020-07-28 23:09:34 +05:30
attr_reader :models
def initialize ( models = self . class . models )
@models = models
end
2021-09-30 23:02:18 +05:30
def sync_partitions
2020-10-24 23:57:45 +05:30
Gitlab :: AppLogger . info ( " Checking state of dynamic postgres partitions " )
2020-07-28 23:09:34 +05:30
models . each do | model |
# Double-checking before getting the lease:
2021-09-30 23:02:18 +05:30
# The prevailing situation is no missing partitions and no extra partitions
next if missing_partitions ( model ) . empty? && extra_partitions ( model ) . empty?
2020-07-28 23:09:34 +05:30
2021-09-30 23:02:18 +05:30
only_with_exclusive_lease ( model , lease_key : MANAGEMENT_LEASE_KEY ) do
2020-07-28 23:09:34 +05:30
partitions_to_create = missing_partitions ( model )
2021-09-30 23:02:18 +05:30
create ( partitions_to_create ) unless partitions_to_create . empty?
2020-07-28 23:09:34 +05:30
2021-10-27 15:23:28 +05:30
if Feature . enabled? ( :partition_pruning , default_enabled : :yaml )
2021-09-30 23:02:18 +05:30
partitions_to_detach = extra_partitions ( model )
detach ( partitions_to_detach ) unless partitions_to_detach . empty?
end
2020-07-28 23:09:34 +05:30
end
2021-06-08 01:23:25 +05:30
rescue StandardError = > e
2021-10-27 15:23:28 +05:30
Gitlab :: AppLogger . error ( message : " Failed to create / detach partition(s) " ,
table_name : model . table_name ,
exception_class : e . class ,
exception_message : e . message )
2020-07-28 23:09:34 +05:30
end
end
private
def missing_partitions ( model )
return [ ] unless connection . table_exists? ( model . table_name )
model . partitioning_strategy . missing_partitions
end
2021-09-30 23:02:18 +05:30
def extra_partitions ( model )
return [ ] unless connection . table_exists? ( model . table_name )
model . partitioning_strategy . extra_partitions
end
def only_with_exclusive_lease ( model , lease_key : )
lease = Gitlab :: ExclusiveLease . new ( lease_key % model . table_name , timeout : LEASE_TIMEOUT )
2020-07-28 23:09:34 +05:30
yield if lease . try_obtain
ensure
lease & . cancel
end
2021-09-30 23:02:18 +05:30
def create ( partitions )
2020-07-28 23:09:34 +05:30
connection . transaction do
with_lock_retries do
partitions . each do | partition |
connection . execute partition . to_sql
2021-10-27 15:23:28 +05:30
Gitlab :: AppLogger . info ( message : " Created partition " ,
partition_name : partition . partition_name ,
table_name : partition . table )
2020-07-28 23:09:34 +05:30
end
end
end
end
2021-09-30 23:02:18 +05:30
def detach ( partitions )
connection . transaction do
with_lock_retries do
partitions . each { | p | detach_one_partition ( p ) }
end
end
end
def detach_one_partition ( partition )
2021-10-27 15:23:28 +05:30
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
2021-09-30 23:02:18 +05:30
end
2020-07-28 23:09:34 +05:30
def with_lock_retries ( & block )
2021-01-03 14:25:43 +05:30
Gitlab :: Database :: WithLockRetries . new (
2020-07-28 23:09:34 +05:30
klass : self . class ,
logger : Gitlab :: AppLogger
2021-01-03 14:25:43 +05:30
) . run ( & block )
2020-07-28 23:09:34 +05:30
end
def connection
ActiveRecord :: Base . connection
end
end
end
end
end