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
|
# frozen_string_literal: true
module Gitlab
module Database
class BumpSequences
SEQUENCE_NAME_MATCHER = /nextval\('([a-z_]+)'::regclass\)/
# gitlab_schema: can be 'gitlab_main', 'gitlab_ci', 'gitlab_main_cell', 'gitlab_shared'
# increase_by: positive number, to increase the sequence by
# base_model: is to choose which connection to use to query the tables
def initialize(gitlab_schema, increase_by, base_model = ApplicationRecord)
@base_model = base_model
@gitlab_schema = gitlab_schema
@increase_by = increase_by
end
def execute
sequences_by_gitlab_schema(base_model, gitlab_schema).each do |sequence_name|
increment_sequence_by(base_model.connection, sequence_name, increase_by)
end
end
private
attr_reader :base_model, :gitlab_schema, :increase_by
def sequences_by_gitlab_schema(base_model, gitlab_schema)
tables = Gitlab::Database::GitlabSchema.tables_to_schema.select do |_table_name, schema_name|
schema_name == gitlab_schema
end.keys
sequences = []
tables.each do |table|
model = Class.new(base_model) do
self.table_name = table
end
model.columns.each do |column|
match_result = column.default_function&.match(SEQUENCE_NAME_MATCHER)
next unless match_result
sequences << match_result[1]
end
end
sequences
end
# This method is going to increase the sequence next_value by:
# - increment_by + 1 if the sequence has the attribute is_called = True (which is the common case)
# - increment_by if the sequence has the attribute is_called = False (for example, a newly created sequence)
# It uses ALTER SEQUENCE as a safety mechanism to avoid that no concurrent insertions
# will cause conflicts on the sequence.
# This is because ALTER SEQUENCE blocks concurrent nextval, currval, lastval, and setval calls.
def increment_sequence_by(connection, sequence_name, increment_by)
connection.transaction do
# The first call is to make sure that the sequence's is_called value is set to `true`
# This guarantees that the next call to `nextval` will increase the sequence by `increment_by`
connection.select_value("SELECT nextval($1)", nil, [sequence_name])
connection.execute("ALTER SEQUENCE #{sequence_name} INCREMENT BY #{increment_by}")
connection.select_value("select nextval($1)", nil, [sequence_name])
connection.execute("ALTER SEQUENCE #{sequence_name} INCREMENT BY 1")
end
end
end
end
end
|