1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
|
# frozen_string_literal: true
module Gitlab
module Database
module PartitioningMigrationHelpers
module TableManagementHelpers
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`.
#
# A copy of the original table is required as PG currently does not support partitioning existing tables.
#
# Example:
#
# partition_table_by_date :audit_events, :created_at, min_date: Date.new(2020, 1), max_date: Date.new(2020, 6)
#
# Required options are:
# :min_date - a date specifying the lower bounds of the partition range
# :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)
raise "primary key not defined for #{table_name}" if primary_key.nil?
partition_column = find_column_definition(table_name, column_name)
raise "partition column #{column_name} does not exist on #{table_name}" if partition_column.nil?
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.
#
# Example:
#
# drop_partitioned_table_for :audit_events
#
def drop_partitioned_table_for(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)
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
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)
if column_of_type?(table_name, primary_key, :integer)
# Default to int8 primary keys to prevent overflow
change_column(table_name, primary_key, :bigint)
end
end
end
def column_of_type?(table_name, column, type)
find_column_definition(table_name, column).type == type
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_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')}"
next_date = min_date.next_month
lower_bound = to_sql_date_literal(min_date)
upper_bound = to_sql_date_literal(next_date)
create_range_partition_safely(partition_name, table_name, lower_bound, upper_bound)
min_date = next_date
end
end
def to_sql_date_literal(date)
connection.quote(date.strftime('%Y-%m-%d'))
end
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
end
end
|
