diff options
Diffstat (limited to 'lib/generators/gitlab/partitioning/foreign_keys_generator.rb')
-rw-r--r-- | lib/generators/gitlab/partitioning/foreign_keys_generator.rb | 151 |
1 files changed, 151 insertions, 0 deletions
diff --git a/lib/generators/gitlab/partitioning/foreign_keys_generator.rb b/lib/generators/gitlab/partitioning/foreign_keys_generator.rb new file mode 100644 index 00000000000..e013188c753 --- /dev/null +++ b/lib/generators/gitlab/partitioning/foreign_keys_generator.rb @@ -0,0 +1,151 @@ +# frozen_string_literal: true + +require 'rails/generators' +require 'rails/generators/active_record/migration' + +module Gitlab + module Partitioning + class ForeignKeysGenerator < Rails::Generators::Base + include ActiveRecord::Generators::Migration + + desc 'This generator creates the migrations needed for updating the foreign keys when partitioning the tables' + + source_root File.expand_path('templates', __dir__) + + class_option :target, type: :string, required: true, desc: 'Target table name' + class_option :source, type: :string, required: true, desc: 'Source table name' + class_option :partitioning_column, type: :string, default: :partition_id, + desc: 'The column that is used for partitioning' + class_option :database, type: :string, default: :ci, + desc: 'Database name connection' + + def create_fk_index_migration + migration_template( + '../templates/foreign_key_index.rb.template', + fk_index_file_name) + end + + def create_fk_definition_migration + migration_template( + '../templates/foreign_key_definition.rb.template', + fk_definition_file_name) + end + + def create_fk_validation_migration + migration_template( + '../templates/foreign_key_validation.rb.template', + fk_validation_file_name) + end + + def remove_old_fk_migration + migration_template( + '../templates/foreign_key_removal.rb.template', + fk_removal_file_name) + end + + private + + def fk_index_file_name + post_migration_file_path( + "add_fk_index_to_#{source_table_name}_on_#{partitioning_column}_and_#{foreign_key_column}.rb") + end + + def fk_definition_file_name + post_migration_file_path( + "add_fk_to_#{source_table_name}_on_#{partitioning_column}_and_#{foreign_key_column}.rb") + end + + def fk_validation_file_name + post_migration_file_path( + "validate_fk_on_#{source_table_name}_#{partitioning_column}_and_#{foreign_key_column}.rb") + end + + def fk_removal_file_name + post_migration_file_path( + "remove_fk_to_#{target_table_name}_#{source_table_name}_on_#{foreign_key_column}.rb") + end + + def post_migration_file_path(name) + File.join(db_migrate_path, name) + end + + def db_migrate_path + super.sub('migrate', 'post_migrate') + end + + def source_table_name + options[:source] + end + + def target_table_name + options[:target] + end + + def partitioning_column + options[:partitioning_column] + end + + def foreign_keys_candidates + connection + .foreign_keys(source_table_name) + .select { |fk| fk.to_table == target_table_name } + .reject { |fk| fk.name.end_with?('_p') } + end + + def fk_candidate + @fk_candidate ||= select_foreign_key + end + + def foreign_key_name + fk_candidate.name + end + + def partitioned_foreign_key_name + "#{foreign_key_name}_p" + end + + def foreign_key_column + fk_candidate.column + end + + def fk_on_delete_option + fk_candidate.on_delete + end + + def fk_target_column + fk_candidate.primary_key + end + + def connection + Gitlab::Database + .database_base_models + .fetch(options[:database]) + .connection + end + + def select_foreign_key + case foreign_keys_candidates.size + when 0 + raise Thor::InvocationError, "No FK found between #{source_table_name} and #{target_table_name}" + when 1 + foreign_keys_candidates.first + else + select_fk_from_user_input + end + end + + def select_fk_from_user_input + options = (0...foreign_keys_candidates.size).to_a.map(&:to_s) + + say "There are multiple FKs between #{source_table_name} and #{target_table_name}:" + foreign_keys_candidates.each.with_index do |fk, index| + say "\t#{index} : #{fk}" + end + + input = ask("Please select one:", limited_to: options, default: '0') + + foreign_keys_candidates.fetch(input.to_i) + end + end + end +end |