diff --git a/ee/spec/lib/gitlab/database/desired_sharding_key_spec.rb b/ee/spec/lib/gitlab/database/desired_sharding_key_spec.rb index 1c90e8f0e4844722cd69eba4cf4dbd8665b350af..a8e7d0ca29f92b75e0712a30b89f654d755ad5aa 100644 --- a/ee/spec/lib/gitlab/database/desired_sharding_key_spec.rb +++ b/ee/spec/lib/gitlab/database/desired_sharding_key_spec.rb @@ -15,43 +15,67 @@ end end - it 'must be possible to backfill it via backfill_via' do - desired_sharding_key_entries.each do |entry| - entry.desired_sharding_key.each do |desired_column, details| - table = entry.table_name - connection = Gitlab::Database.schemas_to_base_models[entry.gitlab_schema].first.connection - sharding_key = desired_column - parent = details['backfill_via']['parent'] - foreign_key = parent['foreign_key'] - parent_table = parent['table'] - parent_sharding_key = parent['sharding_key'] - - connection.execute("ALTER TABLE #{table} ADD COLUMN IF NOT EXISTS #{sharding_key} bigint") - - # Confirming it at least produces a valid query - connection.execute <<~SQL - UPDATE #{table} - SET #{sharding_key} = #{parent_table}.#{parent_sharding_key} - FROM #{parent_table} - WHERE #{table}.#{foreign_key} = #{parent_table}.id - SQL + context 'for tables that already have a backfilled, non-nullable sharding key on their parent' do + it 'must be possible to backfill it via backfill_via' do + desired_sharding_key_entries_not_awaiting_backfill_on_parent.each do |entry| + entry.desired_sharding_key.each do |desired_column, details| + table = entry.table_name + connection = Gitlab::Database.schemas_to_base_models[entry.gitlab_schema].first.connection + sharding_key = desired_column + parent = details['backfill_via']['parent'] + foreign_key = parent['foreign_key'] + parent_table = parent['table'] + parent_sharding_key = parent['sharding_key'] + + connection.execute("ALTER TABLE #{table} ADD COLUMN IF NOT EXISTS #{sharding_key} bigint") + + # Confirming it at least produces a valid query + connection.execute <<~SQL + EXPLAIN UPDATE #{table} + SET #{sharding_key} = #{parent_table}.#{parent_sharding_key} + FROM #{parent_table} + WHERE #{table}.#{foreign_key} = #{parent_table}.id + SQL + end + end + end + + it 'the parent.belongs_to must be a model with the parent.sharding_key column' do + desired_sharding_key_entries_not_awaiting_backfill_on_parent.each do |entry| + model = entry.classes.first.constantize + entry.desired_sharding_key.each do |_column, details| + parent = details['backfill_via']['parent'] + parent_sharding_key = parent['sharding_key'] + belongs_to = parent['belongs_to'] + parent_association = model.reflect_on_association(belongs_to) + expect(parent_association).not_to be_nil, + "Invalid backfill_via.parent.belongs_to: #{belongs_to} in db/docs for #{entry.table_name}" + parent_columns = parent_association.klass.columns.map(&:name) + + expect(parent_columns).to include(parent_sharding_key) + end end end end - it 'the parent.belongs_to must be a model with the parent.sharding_key column' do - desired_sharding_key_entries.each do |entry| - model = entry.classes.first.constantize - entry.desired_sharding_key.each do |_column, details| - parent = details['backfill_via']['parent'] - parent_sharding_key = parent['sharding_key'] - belongs_to = parent['belongs_to'] - parent_association = model.reflect_on_association(belongs_to) - expect(parent_association).not_to be_nil, - "Invalid backfil_via.parent.belongs_to: #{belongs_to} in db/docs for #{entry.table_name}" - parent_columns = parent_association.klass.columns.map(&:name) - - expect(parent_columns).to include(parent_sharding_key) + context 'for tables that do not already have a backfilled, non-nullable sharding key on their parent' \ + 'but only has a desired sharding key on their parent' do + it 'the parent.belongs_to must be a model with a desired_sharding_key' do + desired_sharding_key_entries_awaiting_backfill_on_parent.each do |entry| + model = entry.classes.first.constantize + entry.desired_sharding_key.each do |column, details| + parent = details['backfill_via']['parent'] + parent_table = parent['table'] + belongs_to = parent['belongs_to'] + sharding_key_of_parent = parent['sharding_key'] + + parent_association = model.reflect_on_association(belongs_to) + expect(parent_association).not_to be_nil, + "Invalid backfill_via.parent.belongs_to: #{belongs_to} in db/docs for #{entry.table_name}" + + expect(desired_sharding_keys_of(parent_table).keys).to include(column) + expect(desired_sharding_keys_of(parent_table).keys).to include(sharding_key_of_parent) + end end end end @@ -72,4 +96,22 @@ def desired_sharding_key_entries entry.desired_sharding_key.present? end end + + def desired_sharding_key_entries_not_awaiting_backfill_on_parent + ::Gitlab::Database::Dictionary.entries.select do |entry| + entry.desired_sharding_key.present? && + entry.desired_sharding_key.all? { |_, details| details['awaiting_backfill_on_parent'].blank? } + end + end + + def desired_sharding_key_entries_awaiting_backfill_on_parent + ::Gitlab::Database::Dictionary.entries.select do |entry| + entry.desired_sharding_key.present? && + entry.desired_sharding_key.all? { |_, details| details['awaiting_backfill_on_parent'] } + end + end + + def desired_sharding_keys_of(table_name) + Gitlab::Database::Dictionary.entries.find_by_table_name(table_name).desired_sharding_key + end end