Merge branch 'alexives/293713/fix_queries_for_snippet_repositories' into 'master'

Add specs for snippet_repository replication helper

See merge request gitlab-org/gitlab!49976

ee/app/models/ee/snippet_repository.rb

@@ -17,33 +17,30 @@ module EE
       def replicables_for_current_secondary(primary_key_in)
         node = ::Gitlab::Geo.current_node
 
-        if !node.selective_sync?
-          all.primary_key_in(primary_key_in)
+        replicables = if !node.selective_sync?
+                        all
                       elsif node.selective_sync_by_namespaces?
-          snippet_repositories_for_selected_namespaces(primary_key_in)
+                        snippet_repositories_for_selected_namespaces
                       elsif node.selective_sync_by_shards?
-          snippet_repositories_for_selected_shards(primary_key_in)
+                        snippet_repositories_for_selected_shards
                       else
                         self.none
                       end
+
+        replicables.primary_key_in(primary_key_in)
       end
 
-      def snippet_repositories_for_selected_namespaces(primary_key_in)
-        personal_snippets = self.where(snippet: ::Snippet.only_personal_snippets.primary_key_in(primary_key_in))
+      def snippet_repositories_for_selected_namespaces
+        personal_snippets = self.joins(:snippet).where(snippet: ::Snippet.only_personal_snippets)
 
-        project_snippets = self.where(snippet: ::Snippet.for_projects(::Gitlab::Geo.current_node.projects.select(:id))
-                                                   .primary_key_in(primary_key_in))
+        project_snippets = self.joins(snippet: :project)
+                               .merge(::Snippet.for_projects(::Gitlab::Geo.current_node.projects.select(:id)))
 
-        # We use `find_by_sql` here in order to perform the union operation and cast the results as
-        # Snippet repositories. If we use a `from_union`, it wraps the query in in a sub-select and
-        # it increases the query time by a surprising amount.
-        self.find_by_sql(::Gitlab::SQL::Union.new([project_snippets, personal_snippets]).to_sql)
+        self.from_union([project_snippets, personal_snippets])
       end
 
-      def snippet_repositories_for_selected_shards(primary_key_in)
-        self
-          .for_repository_storage(::Gitlab::Geo.current_node.selective_sync_shards)
-          .primary_key_in(primary_key_in)
+      def snippet_repositories_for_selected_shards
+        self.for_repository_storage(::Gitlab::Geo.current_node.selective_sync_shards)
       end
     end
 

ee/spec/models/snippet_repository_spec.rb

@@ -11,7 +11,7 @@ RSpec.describe SnippetRepository, :request_store, :geo, type: :model do
     stub_current_geo_node(node)
   end
 
-  describe '#replicables_for_current_secondary' do
+  context 'with 3 groups, 2 projects, and 5 snippets' do
     let(:group_1) { create(:group) }
     let(:group_2) { create(:group) }
     let(:nested_group_1) { create(:group, parent: group_1) }
@@ -30,6 +30,53 @@ RSpec.describe SnippetRepository, :request_store, :geo, type: :model do
     let!(:snippet_repository_4) { create(:snippet_repository) }
     let!(:snippet_repository_5) { create(:snippet_repository, snippet: project_snippet_2) }
 
+    describe '#in_replicables_for_current_secondary?' do
+      it 'all returns true if all are replicated' do
+        [
+          snippet_repository_1,
+          snippet_repository_2,
+          snippet_repository_3,
+          snippet_repository_4,
+          snippet_repository_5
+        ].each do |repository|
+          expect(repository.in_replicables_for_current_secondary?).to be true
+        end
+      end
+
+      context 'with selective sync by namespace' do
+        before do
+          node.update!(selective_sync_type: 'namespaces', namespaces: [group_1])
+        end
+
+        it 'returns true for snippets in the namespace' do
+          expect(snippet_repository_1.in_replicables_for_current_secondary?).to be true
+        end
+
+        it 'returns true for personal snippets' do
+          expect(snippet_repository_2.in_replicables_for_current_secondary?).to be true
+        end
+
+        it 'returns false for project snippets not in an included namespace' do
+          expect(snippet_repository_5.in_replicables_for_current_secondary?).to be false
+        end
+      end
+
+      context 'with selective sync by shard' do
+        before do
+          node.update!(selective_sync_type: 'shards', selective_sync_shards: ['default'])
+        end
+
+        it 'returns true for snippets in the shard' do
+          expect(snippet_repository_1.in_replicables_for_current_secondary?).to be true
+        end
+
+        it 'returns false for project snippets not in an included shard' do
+          expect(snippet_repository_3.in_replicables_for_current_secondary?).to be false
+        end
+      end
+    end
+
+    describe '#replicables_for_current_secondary' do
       it 'returns all snippet_repositories without selective sync' do
         expect(described_class.replicables_for_current_secondary(1..described_class.last.id)).to match_array([
           snippet_repository_1,
@@ -72,4 +119,5 @@ RSpec.describe SnippetRepository, :request_store, :geo, type: :model do
         expect(described_class.replicables_for_current_secondary(1..described_class.last.id)).to be_empty
       end
     end
+  end
 end