Add group restrictions to Geo::FileDownloadDispatchWorker

app/workers/geo/file_download_dispatch_worker.rb

@@ -9,30 +9,51 @@ module Geo
     end
 
     def load_pending_resources
-      lfs_object_ids = find_lfs_object_ids(db_retrieve_batch_size)
-      objects_ids    = find_object_ids(db_retrieve_batch_size)
+      restricted_project_ids = Gitlab::Geo.current_node.project_ids
+      lfs_object_ids         = find_lfs_object_ids(restricted_project_ids)
+      objects_ids            = find_object_ids(restricted_project_ids)
 
       interleave(lfs_object_ids, objects_ids)
     end
 
-    def find_object_ids(limit)
+    def find_object_ids(restricted_project_ids)
       downloaded_ids = find_downloaded_ids([:attachment, :avatar, :file])
 
-      Upload.where.not(id: downloaded_ids)
-            .order(created_at: :desc)
-            .limit(limit)
-            .pluck(:id, :uploader)
-            .map { |id, uploader| [id, uploader.sub(/Uploader\z/, '').downcase] }
+      relation =
+        if restricted_project_ids
+          uploads_table   = Upload.arel_table
+          group_uploads   = uploads_table[:model_type].eq('Namespace').and(uploads_table[:model_id].in(Gitlab::Geo.current_node.group_ids))
+          project_uploads = uploads_table[:model_type].eq('Project').and(uploads_table[:model_id].in(restricted_project_ids))
+          other_uploads   = uploads_table[:model_type].not_in(%w[Namespace Project])
+
+          Upload.where(group_uploads.or(project_uploads).or(other_uploads))
+
+        else
+          Upload.all
+        end
+
+      relation.where.not(id: downloaded_ids)
+              .order(created_at: :desc)
+              .limit(db_retrieve_batch_size)
+              .pluck(:id, :uploader)
+              .map { |id, uploader| [id, uploader.sub(/Uploader\z/, '').downcase] }
     end
 
-    def find_lfs_object_ids(limit)
+    def find_lfs_object_ids(restricted_project_ids)
       downloaded_ids = find_downloaded_ids([:lfs])
 
-      LfsObject.where.not(id: downloaded_ids)
-               .order(created_at: :desc)
-               .limit(limit)
-               .pluck(:id)
-               .map { |id| [id, :lfs] }
+      relation =
+        if restricted_project_ids
+          LfsObject.joins(:projects).where(projects: { id: restricted_project_ids })
+        else
+          LfsObject.all
+        end
+
+      relation.where.not(id: downloaded_ids)
+              .order(created_at: :desc)
+              .limit(db_retrieve_batch_size)
+              .pluck(:id)
+              .map { |id| [id, :lfs] }
     end
 
     def find_downloaded_ids(file_types)

spec/workers/geo/file_download_dispatch_worker_spec.rb

 require 'spec_helper'
 
 describe Geo::FileDownloadDispatchWorker do
+  let!(:primary)   { create(:geo_node, :primary, host: 'primary-geo-node') }
+  let!(:secondary) { create(:geo_node, :current) }
+
   before do
-    @primary = create(:geo_node, :primary, host: 'primary-geo-node')
-    @secondary = create(:geo_node, :current)
     allow(Gitlab::Geo).to receive(:secondary?).and_return(true)
     allow_any_instance_of(Gitlab::ExclusiveLease)
       .to receive(:try_obtain).and_return(true)
@@ -28,8 +29,8 @@ describe Geo::FileDownloadDispatchWorker do
     it 'does not schedule anything when node is disabled' do
       create(:lfs_object, :with_file)
 
-      @secondary.enabled = false
-      @secondary.save
+      secondary.enabled = false
+      secondary.save
 
       expect(GeoFileDownloadWorker).not_to receive(:perform_async)
 
@@ -73,5 +74,40 @@ describe Geo::FileDownloadDispatchWorker do
         subject.perform
       end
     end
+
+    context 'when node have group restrictions' do
+      let(:group_1)    { create(:group) }
+      let!(:project_1) { create(:empty_project, group: group_1) }
+      let!(:project_2) { create(:empty_project) }
+
+      before do
+        allow(ProjectCacheWorker).to receive(:perform_async).and_return(true)
+        allow_any_instance_of(described_class).to receive(:over_time?).and_return(false)
+
+        secondary.update_attribute(:groups, [group_1])
+      end
+
+      it 'does not perform GeoFileDownloadWorker for LFS object that do not belong to selected groups to replicate' do
+        create(:lfs_objects_project, project: project_1)
+        create(:lfs_objects_project, project: project_2)
+
+        expect(GeoFileDownloadWorker).to receive(:perform_async).once.and_return(spy)
+
+        subject.perform
+      end
+
+      it 'does not perform GeoFileDownloadWorker for upload objects that do not belong to selected groups to replicate' do
+        avatar = fixture_file_upload(Rails.root.join('spec/fixtures/dk.png'))
+        create(:upload, model: group_1, path: avatar)
+        create(:upload, model: create(:group), path: avatar)
+        create(:upload, model: project_1, path: avatar)
+        create(:upload, model: project_2, path: avatar)
+        create(:note, :with_attachment)
+
+        expect(GeoFileDownloadWorker).to receive(:perform_async).exactly(3).times.and_return(spy)
+
+        subject.perform
+      end
+    end
   end
 end