Add specs for changes surrounding "next" desired sharding key

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