1 files changed, 68 insertions, 0 deletions
diff --git a/lib/tasks/gitlab/db/decomposition/rollback/bump_ci_sequences.rake b/lib/tasks/gitlab/db/decomposition/rollback/bump_ci_sequences.rake
new file mode 100644
index 00000000000..4d78acb3011
--- /dev/null
+++ b/lib/tasks/gitlab/db/decomposition/rollback/bump_ci_sequences.rake
@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+namespace :gitlab do
+  namespace :db do
+    namespace :decomposition do
+      namespace :rollback do
+        SEQUENCE_NAME_MATCHER = /nextval\('([a-z_]+)'::regclass\)/.freeze
+
+        desc 'Bump all the CI tables sequences on the Main Database'
+        task :bump_ci_sequences, [:increase_by] => :environment do |_t, args|
+          increase_by = args.increase_by.to_i
+          if increase_by < 1
+            puts 'Please specify a positive integer `increase_by` value'.color(:red)
+            puts 'Example: rake gitlab:db:decomposition:rollback:bump_ci_sequences[100000]'.color(:green)
+            exit 1
+          end
+
+          sequences_by_gitlab_schema(ApplicationRecord, :gitlab_ci).each do |sequence_name|
+            increment_sequence_by(ApplicationRecord.connection, sequence_name, increase_by)
+          end
+        end
+      end
+    end
+  end
+end
+
+# base_model is to choose which connection to use to query the tables
+# gitlab_schema, can be 'gitlab_main', 'gitlab_ci', 'gitlab_shared'
+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
+
+  models = tables.map do |table|
+    model = Class.new(base_model)
+    model.table_name = table
+    model
+  end
+
+  sequences = []
+  models.each do |model|
+    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