Add group restrictions to Geo node status

app/models/geo_node_status.rb

@@ -15,7 +15,12 @@ class GeoNodeStatus
   end
 
   def repositories_count
-    @repositories_count ||= Project.count
+    @repositories_count ||=
+      if restricted_project_ids
+        Project.where(id: restricted_project_ids).count
+      else
+        Project.count
+      end
   end
 
   def repositories_count=(value)
@@ -23,7 +28,12 @@ class GeoNodeStatus
   end
 
   def repositories_synced_count
-    @repositories_synced_count ||= Geo::ProjectRegistry.synced.count
+    @repositories_synced_count ||=
+      if restricted_project_ids
+        Geo::ProjectRegistry.synced.where(project_id: restricted_project_ids).count
+      else
+        Geo::ProjectRegistry.synced.count
+      end
   end
 
   def repositories_synced_count=(value)
@@ -35,7 +45,12 @@ class GeoNodeStatus
   end
 
   def repositories_failed_count
-    @repositories_failed_count ||= Geo::ProjectRegistry.failed.count
+    @repositories_failed_count ||=
+      if restricted_project_ids
+        Geo::ProjectRegistry.failed.where(project_id: restricted_project_ids).count
+      else
+        Geo::ProjectRegistry.failed.count
+      end
   end
 
   def repositories_failed_count=(value)
@@ -43,7 +58,7 @@ class GeoNodeStatus
   end
 
   def lfs_objects_count
-    @lfs_objects_count ||= LfsObject.count
+    @lfs_objects_count ||= lfs_objects.count
   end
 
   def lfs_objects_count=(value)
@@ -51,7 +66,15 @@ class GeoNodeStatus
   end
 
   def lfs_objects_synced_count
-    @lfs_objects_synced_count ||= Geo::FileRegistry.where(file_type: :lfs).count
+    @lfs_objects_synced_count ||= begin
+      relation = Geo::FileRegistry.where(file_type: :lfs)
+
+      if restricted_project_ids
+        relation = relation.where(file_id: lfs_objects.pluck(:id))
+      end
+
+      relation.count
+    end
   end
 
   def lfs_objects_synced_count=(value)
@@ -63,7 +86,7 @@ class GeoNodeStatus
   end
 
   def attachments_count
-    @attachments_count ||= Upload.count
+    @attachments_count ||= attachments.count
   end
 
   def attachments_count=(value)
@@ -72,7 +95,7 @@ class GeoNodeStatus
 
   def attachments_synced_count
     @attachments_synced_count ||= begin
-      upload_ids = Upload.pluck(:id)
+      upload_ids = attachments.pluck(:id)
       synced_ids = Geo::FileRegistry.where(file_type: [:attachment, :avatar, :file]).pluck(:file_id)
 
       (synced_ids & upload_ids).length
@@ -94,4 +117,33 @@ class GeoNodeStatus
 
     (synced.to_f / total.to_f) * 100.0
   end
+
+  def attachments
+    @attachments ||=
+      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
+  end
+
+  def lfs_objects
+    @lfs_objects ||=
+      if restricted_project_ids
+        LfsObject.joins(:projects).where(projects: { id: restricted_project_ids })
+      else
+        LfsObject.all
+      end
+  end
+
+  def restricted_project_ids
+    return @restricted_project_ids if defined?(@restricted_project_ids)
+
+    @restricted_project_ids = Gitlab::Geo.current_node.project_ids
+  end
 end

app/workers/geo/file_download_dispatch_worker.rb

@@ -27,7 +27,6 @@ module Geo
           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

spec/factories/geo/file_registry.rb

@@ -8,5 +8,9 @@ FactoryGirl.define do
     trait :avatar do
       file_type :avatar
     end
+
+    trait :lfs do
+      file_type :lfs
+    end
   end
 end

spec/models/geo_node_status_spec.rb

 require 'spec_helper'
 
 describe GeoNodeStatus do
+  let!(:geo_node) { create(:geo_node, :current) }
+  let(:group)      { create(:group) }
+  let!(:project_1) { create(:empty_project, group: group) }
+  let!(:project_2) { create(:empty_project, group: group) }
+  let!(:project_3) { create(:empty_project) }
+  let!(:project_4) { create(:empty_project) }
+
   subject { described_class.new }
 
   describe '#healthy?' do
@@ -40,7 +47,7 @@ describe GeoNodeStatus do
       expect(subject.attachments_synced_count).to eq(0)
 
       upload = Upload.find_by(model: user, uploader: 'AvatarUploader')
-      Geo::FileRegistry.create(file_type: :avatar, file_id: upload.id)
+      create(:geo_file_registry, :avatar, file_id: upload.id)
 
       subject = described_class.new
       expect(subject.attachments_count).to eq(1)
@@ -53,7 +60,7 @@ describe GeoNodeStatus do
       expect(subject.attachments_synced_count).to eq(0)
 
       upload = Upload.find_by(model: user, uploader: 'AvatarUploader')
-      Geo::FileRegistry.create(file_type: :avatar, file_id: upload.id)
+      create(:geo_file_registry, :avatar, file_id: upload.id)
 
       subject = described_class.new
       expect(subject.attachments_count).to eq(1)
@@ -62,51 +69,80 @@ describe GeoNodeStatus do
   end
 
   describe '#attachments_synced_in_percentage' do
-    it 'returns 0 when no objects are available' do
-      subject.attachments_count = 0
-      subject.attachments_synced_count = 0
+    let(:avatar) { fixture_file_upload(Rails.root.join('spec/fixtures/dk.png')) }
+    let(:upload_1) { create(:upload, model: group, path: avatar) }
+    let(:upload_2) { create(:upload, model: project_1, path: avatar) }
+
+    before do
+      create(:upload, model: create(:group), path: avatar)
+      create(:upload, model: project_3, path: avatar)
+    end
 
+    it 'returns 0 when no objects are available' do
       expect(subject.attachments_synced_in_percentage).to eq(0)
     end
 
-    it 'returns the right percentage' do
-      subject.attachments_count = 4
-      subject.attachments_synced_count = 1
+    it 'returns the right percentage with no group restrictions' do
+      create(:geo_file_registry, :avatar, file_id: upload_1.id)
+      create(:geo_file_registry, :avatar, file_id: upload_2.id)
 
-      expect(subject.attachments_synced_in_percentage).to be_within(0.0001).of(25)
+      expect(subject.attachments_synced_in_percentage).to be_within(0.0001).of(50)
+    end
+
+    it 'returns the right percentage with group restrictions' do
+      geo_node.update_attribute(:groups, [group])
+      create(:geo_file_registry, :avatar, file_id: upload_1.id)
+      create(:geo_file_registry, :avatar, file_id: upload_2.id)
+
+      expect(subject.attachments_synced_in_percentage).to be_within(0.0001).of(100)
     end
   end
 
   describe '#lfs_objects_synced_in_percentage' do
-    it 'returns 0 when no objects are available' do
-      subject.lfs_objects_count = 0
-      subject.lfs_objects_synced_count = 0
+    let(:lfs_object_project) { create(:lfs_objects_project, project: project_1) }
 
+    before do
+      allow(ProjectCacheWorker).to receive(:perform_async).and_return(true)
+
+      create(:lfs_objects_project, project: project_1)
+      create_list(:lfs_objects_project, 2, project: project_3)
+    end
+
+    it 'returns 0 when no objects are available' do
       expect(subject.lfs_objects_synced_in_percentage).to eq(0)
     end
 
-    it 'returns the right percentage' do
-      subject.lfs_objects_count = 4
-      subject.lfs_objects_synced_count = 1
+    it 'returns the right percentage with no group restrictions' do
+      create(:geo_file_registry, :lfs, file_id: lfs_object_project.lfs_object_id)
 
       expect(subject.lfs_objects_synced_in_percentage).to be_within(0.0001).of(25)
     end
+
+    it 'returns the right percentage with group restrictions' do
+      geo_node.update_attribute(:groups, [group])
+      create(:geo_file_registry, :lfs, file_id: lfs_object_project.lfs_object_id)
+
+      expect(subject.lfs_objects_synced_in_percentage).to be_within(0.0001).of(50)
+    end
   end
 
   describe '#repositories_synced_in_percentage' do
-    it 'returns 0 when no objects are available' do
-      subject.repositories_count = 0
-      subject.repositories_synced_count = 0
-
+    it 'returns 0 when no projects are available' do
       expect(subject.repositories_synced_in_percentage).to eq(0)
     end
 
-    it 'returns the right percentage' do
-      subject.repositories_count = 4
-      subject.repositories_synced_count = 1
+    it 'returns the right percentage with no group restrictions' do
+      create(:geo_project_registry, :synced, project: project_1)
 
       expect(subject.repositories_synced_in_percentage).to be_within(0.0001).of(25)
     end
+
+    it 'returns the right percentage with group restrictions' do
+      geo_node.update_attribute(:groups, [group])
+      create(:geo_project_registry, :synced, project: project_1)
+
+      expect(subject.repositories_synced_in_percentage).to be_within(0.0001).of(50)
+    end
   end
 
   context 'when no values are available' do