Require scope by primary key for replicables

No change in behavior. All existing usages of `replicables_for_geo_node`
already apply `primary_key_in` or `id_in` scopes, so this is only a
refactor.

This refactor will allow us to apply the scope to a subquery within
`LfsObject.replicables_for_geo_node` to improve performance.

This refactor will also tend to ensure that `replicables_for_geo_node`
is not accidentally used without restricting the scope.

doc/development/geo/framework.md

@@ -197,7 +197,9 @@ For example, to add support for files referenced by a `Widget` model with a
        file_store == ObjectStorage::Store::LOCAL
      end
 
-     def self.replicables_for_geo_node
+     # @param primary_key_in [Range, Widget] arg to pass to primary_key_in scope
+     # @return [ActiveRecord::Relation<Widget>] everything that should be synced to this node, restricted by primary key
+     def self.replicables_for_geo_node(primary_key_in)
        # Should be implemented. The idea of the method is to restrict
        # the set of synced items depending on synchronization settings
      end

ee/app/models/ee/ci/job_artifact.rb

@@ -82,8 +82,13 @@ module EE
         super
       end
 
-      def replicables_for_geo_node(node = ::Gitlab::Geo.current_node)
-        not_expired.merge(selective_sync_scope(node))
+      # @param primary_key_in [Range, Ci::JobArtifact] arg to pass to primary_key_in scope
+      # @param node [GeoNode] defaults to ::Gitlab::Geo.current_node
+      # @return [ActiveRecord::Relation<Ci::JobArtifact>] everything that should be synced to this node, restricted by primary key
+      def replicables_for_geo_node(primary_key_in, node = ::Gitlab::Geo.current_node)
+        not_expired
+          .primary_key_in(primary_key_in)
+          .merge(selective_sync_scope(node))
           .merge(object_storage_scope(node))
       end
 

ee/app/models/ee/container_repository.rb

@@ -9,8 +9,11 @@ module EE
     end
 
     class_methods do
-      def replicables_for_geo_node(node = ::Gitlab::Geo.current_node)
-        node.container_repositories
+      # @param primary_key_in [Range, ContainerRepository] arg to pass to primary_key_in scope
+      # @param node [GeoNode] defaults to ::Gitlab::Geo.current_node
+      # @return [ActiveRecord::Relation<ContainerRepository>] everything that should be synced to this node, restricted by primary key
+      def replicables_for_geo_node(primary_key_in, node = ::Gitlab::Geo.current_node)
+        node.container_repositories.primary_key_in(primary_key_in)
       end
     end
 

ee/app/models/ee/lfs_object.rb

@@ -17,9 +17,15 @@ module EE
     end
 
     class_methods do
-      def replicables_for_geo_node(node = ::Gitlab::Geo.current_node)
-        local_storage_only = !node&.sync_object_storage
-        local_storage_only ? node.lfs_objects.with_files_stored_locally : node.lfs_objects
+      # @param primary_key_in [Range, LfsObject] arg to pass to primary_key_in scope
+      # @param node [GeoNode] defaults to ::Gitlab::Geo.current_node
+      # @return [ActiveRecord::Relation<LfsObject>] everything that should be synced to this node, restricted by primary key
+      def replicables_for_geo_node(primary_key_in, node = ::Gitlab::Geo.current_node)
+        local_storage_only = !node.sync_object_storage
+
+        scope = node.lfs_objects.primary_key_in(primary_key_in)
+        scope = scope.with_files_stored_locally if local_storage_only
+        scope
       end
     end
 

ee/app/models/ee/merge_request_diff.rb

@@ -29,8 +29,12 @@ module EE
     end
 
     class_methods do
-      def replicables_for_geo_node(node = ::Gitlab::Geo.current_node)
-        has_external_diffs.merge(selective_sync_scope(node))
+      # @param primary_key_in [Range, MergeRequestDiff] arg to pass to primary_key_in scope
+      # @param node [GeoNode] defaults to ::Gitlab::Geo.current_node
+      # @return [ActiveRecord::Relation<MergeRequestDiff>] everything that should be synced to this node, restricted by primary key
+      def replicables_for_geo_node(primary_key_in, node = ::Gitlab::Geo.current_node)
+        has_external_diffs.primary_key_in(primary_key_in)
+                          .merge(selective_sync_scope(node))
                           .merge(object_storage_scope(node))
       end
 

ee/app/models/ee/packages/package_file.rb

@@ -11,9 +11,12 @@ module EE
       end
 
       class_methods do
-        # @return [ActiveRecord::Relation<Packages::PackageFile>] scope of everything that should be synced
-        def replicables_for_geo_node
-          selective_sync_scope.merge(object_storage_scope)
+        # @param primary_key_in [Range, Packages::PackageFile] arg to pass to primary_key_in scope
+        # @return [ActiveRecord::Relation<LfsObject>] everything that should be synced to this node, restricted by primary key
+        def replicables_for_geo_node(primary_key_in)
+          primary_key_in(primary_key_in)
+            .merge(selective_sync_scope)
+            .merge(object_storage_scope)
         end
 
         private

ee/app/models/ee/project.rb

@@ -204,8 +204,11 @@ module EE
     class_methods do
       extend ::Gitlab::Utils::Override
 
-      def replicables_for_geo_node(node = ::Gitlab::Geo.current_node)
-        node.projects
+      # @param primary_key_in [Range, Project] arg to pass to primary_key_in scope
+      # @param node [GeoNode] defaults to ::Gitlab::Geo.current_node
+      # @return [ActiveRecord::Relation<Project>] everything that should be synced to this node, restricted by primary key
+      def replicables_for_geo_node(primary_key_in, node = ::Gitlab::Geo.current_node)
+        node.projects.primary_key_in(primary_key_in)
       end
 
       def search_by_visibility(level)

ee/app/models/ee/snippet_repository.rb

@@ -11,7 +11,10 @@ module EE
     end
 
     class_methods do
-      def replicables_for_geo_node
+      # @param primary_key_in [Range, SnippetRepository] arg to pass to primary_key_in scope
+      # @param node [GeoNode] defaults to ::Gitlab::Geo.current_node
+      # @return [ActiveRecord::Relation<SnippetRepository>] everything that should be synced to this node, restricted by primary key
+      def replicables_for_geo_node(primary_key_in, node = ::Gitlab::Geo.current_node)
         # Not implemented yet. Should be responsible for selective sync
         all
       end

ee/app/models/ee/terraform/state_version.rb

@@ -14,8 +14,13 @@ module EE
       end
 
       class_methods do
-        def replicables_for_geo_node(node = ::Gitlab::Geo.current_node)
-          selective_sync_scope(node).merge(object_storage_scope(node))
+        # @param primary_key_in [Range, Terraform::StateVersion] arg to pass to primary_key_in scope
+        # @param node [GeoNode] defaults to ::Gitlab::Geo.current_node
+        # @return [ActiveRecord::Relation<Terraform::StateVersion>] everything that should be synced to this node, restricted by primary key
+        def replicables_for_geo_node(primary_key_in, node = ::Gitlab::Geo.current_node)
+          primary_key_in(primary_key_in)
+            .merge(selective_sync_scope(node))
+            .merge(object_storage_scope(node))
         end
 
         private

ee/app/models/ee/upload.rb

@@ -18,9 +18,13 @@ module EE
     end
 
     class_methods do
-      # @return [ActiveRecord::Relation<Upload>] scope of everything that should be synced to this node
-      def replicables_for_geo_node(node = ::Gitlab::Geo.current_node)
-        selective_sync_scope(node).merge(object_storage_scope(node))
+      # @param primary_key_in [Range, Upload] arg to pass to primary_key_in scope
+      # @param node [GeoNode] defaults to ::Gitlab::Geo.current_node
+      # @return [ActiveRecord::Relation<Upload>] everything that should be synced to this node, restricted by primary key
+      def replicables_for_geo_node(primary_key_in, node = ::Gitlab::Geo.current_node)
+        primary_key_in(primary_key_in)
+          .merge(selective_sync_scope(node))
+          .merge(object_storage_scope(node))
       end
 
       # Searches for a list of uploads based on the query given in `query`.

ee/app/models/geo/base_registry.rb

@@ -53,8 +53,7 @@ class Geo::BaseRegistry < Geo::TrackingBase
     model_primary_key = self::MODEL_CLASS.primary_key.to_sym
 
     source_ids = self::MODEL_CLASS
-                  .replicables_for_geo_node
-                  .primary_key_in(range)
+                  .replicables_for_geo_node(range)
                   .pluck(self::MODEL_CLASS.arel_table[model_primary_key])
 
     tracked_ids = self.pluck_model_ids_in_range(range)

ee/app/models/geo/upload_registry.rb

@@ -34,8 +34,7 @@ class Geo::UploadRegistry < Geo::BaseRegistry
   #                 For example: [[[1, 'avatar'], [5, 'file']], [[3, 'attachment']]]
   def self.find_registry_differences(range)
     source =
-      self::MODEL_CLASS.replicables_for_geo_node
-          .id_in(range)
+      self::MODEL_CLASS.replicables_for_geo_node(range)
           .pluck(self::MODEL_CLASS.arel_table[:id], self::MODEL_CLASS.arel_table[:uploader])
           .map! { |id, uploader| [id, uploader.sub(/Uploader\z/, '').underscore] }
 

ee/lib/gitlab/geo/replicable_model.rb

@@ -84,7 +84,7 @@ module Gitlab
       end
 
       def in_replicables_for_geo_node?
-        self.class.replicables_for_geo_node.primary_key_in(self).exists?
+        self.class.replicables_for_geo_node(self).exists?
       end
     end
   end

ee/spec/lib/gitlab/geo/replicable_model_spec.rb

@@ -45,7 +45,7 @@ RSpec.describe Gitlab::Geo::ReplicableModel do
 
   describe '#in_replicables_for_geo_node?' do
     it 'reuses replicables_for_geo_node' do
-      expect(DummyModel).to receive(:replicables_for_geo_node).once.and_return(DummyModel.all)
+      expect(DummyModel).to receive(:replicables_for_geo_node).once.with(subject).and_call_original
 
       subject.in_replicables_for_geo_node?
     end

ee/spec/models/ee/ci/job_artifact_spec.rb

@@ -162,7 +162,7 @@ RSpec.describe Ci::JobArtifact do
     end
 
     with_them do
-      subject(:job_artifact_included) { described_class.replicables_for_geo_node.include?(ci_job_artifact) }
+      subject(:job_artifact_included) { described_class.replicables_for_geo_node(ci_job_artifact).exists? }
 
       let(:project) { create(*project_factory) }
       let(:ci_build) { create(:ci_build, project: project) }

ee/spec/models/ee/merge_request_diff_spec.rb

@@ -74,13 +74,13 @@ RSpec.describe MergeRequestDiff do
 
         create(:merge_request, source_project: project)
 
-        expect(described_class.replicables_for_geo_node).to be_empty
+        expect(described_class.replicables_for_geo_node(1..described_class.last.id)).to be_empty
       end
 
       it 'excludes empty diffs' do
         create(:merge_request, source_project: create(:project))
 
-        expect(described_class.replicables_for_geo_node).to be_empty
+        expect(described_class.replicables_for_geo_node(1..described_class.last.id)).to be_empty
       end
     end
 
@@ -119,7 +119,7 @@ RSpec.describe MergeRequestDiff do
       end
 
       it 'returns the proper number of merge request diff states' do
-        expect(described_class.replicables_for_geo_node).to have_attributes(count: synced_states)
+        expect(described_class.replicables_for_geo_node(1..described_class.last.id)).to have_attributes(count: synced_states)
       end
     end
   end

ee/spec/models/ee/terraform/state_version_spec.rb

@@ -60,7 +60,7 @@ RSpec.describe Terraform::StateVersion do
       end
 
       it 'returns the proper number of terraform states' do
-        expect(Terraform::StateVersion.replicables_for_geo_node.count).to eq(synced_states)
+        expect(described_class.replicables_for_geo_node(1..described_class.last.id).count).to eq(synced_states)
       end
     end
   end

ee/spec/models/packages/package_file_spec.rb

@@ -37,7 +37,7 @@ RSpec.describe Packages::PackageFile, type: :model do
   end
 
   describe '.replicables_for_geo_node' do
-    subject { described_class.replicables_for_geo_node }
+    subject { described_class.replicables_for_geo_node(1..described_class.last.id) }
 
     it 'returns a package files scope' do
       secondary = create(:geo_node)

ee/spec/models/upload_spec.rb

@@ -37,7 +37,7 @@ RSpec.describe Upload do
     end
 
     with_them do
-      subject(:upload_included) { described_class.replicables_for_geo_node.include?(upload) }
+      subject(:upload_included) { described_class.replicables_for_geo_node(upload).exists? }
 
       let(:model) { create(*model_factory) }
       let(:node) do

ee/spec/support/helpers/ee/geo_helpers.rb

@@ -83,8 +83,8 @@ module EE
 
         with_replicator Geo::DummyReplicator
 
-        def self.replicables_for_geo_node
-          self.all
+        def self.replicables_for_geo_node(primary_key_in)
+          self.primary_key_in(primary_key_in)
         end
       end
 

ee/spec/support/shared_examples/models/concerns/replicable_model_shared_examples.rb

@@ -36,7 +36,7 @@ RSpec.shared_examples 'a replicable model' do
     end
 
     it 'is implemented' do
-      expect(model_record.class.replicables_for_geo_node).to be_an(ActiveRecord::Relation)
+      expect(model_record.class.replicables_for_geo_node(model_record.id)).to be_an(ActiveRecord::Relation)
     end
   end
 end