Geo: Backfill existing Snippets

ee/app/models/ee/snippet_repository.rb

@@ -6,6 +6,7 @@ module EE
 
     prepended do
       include ::Gitlab::Geo::ReplicableModel
+      include FromUnion
 
       with_replicator Geo::SnippetRepositoryReplicator
     end
@@ -14,8 +15,34 @@ module EE
       # @param primary_key_in [Range, SnippetRepository] arg to pass to primary_key_in scope
       # @return [ActiveRecord::Relation<SnippetRepository>] everything that should be synced to this node, restricted by primary key
       def replicables_for_current_secondary(primary_key_in)
-        # Not implemented yet. Should be responsible for selective sync
-        all
+        node = ::Gitlab::Geo.current_node
+
+        replicables = if !node.selective_sync?
+                        all
+                      elsif node.selective_sync_by_namespaces?
+                        snippet_repositories_for_selected_namespaces
+                      elsif node.selective_sync_by_shards?
+                        snippet_repositories_for_selected_shards
+                      else
+                        self.none
+                      end
+
+        replicables.primary_key_in(primary_key_in)
+      end
+
+      def snippet_repositories_for_selected_namespaces
+        personal_snippets = self.joins(:snippet).where(snippet: ::Snippet.only_personal_snippets)
+
+        project_snippets = self.joins(snippet: :project).where(
+          'snippets.project_id IN(?)',
+          ::Gitlab::Geo.current_node.projects.select(:id)
+        )
+
+        self.from_union([project_snippets, personal_snippets])
+      end
+
+      def snippet_repositories_for_selected_shards
+        self.for_repository_storage(::Gitlab::Geo.current_node.selective_sync_shards)
       end
     end
 

ee/app/models/geo/snippet_repository_registry.rb

@@ -7,4 +7,8 @@ class Geo::SnippetRepositoryRegistry < Geo::BaseRegistry
   MODEL_FOREIGN_KEY = :snippet_repository_id
 
   belongs_to :snippet_repository, class_name: 'SnippetRepository'
+
+  def self.delete_worker_class
+    ::Geo::FrameworkRepositoryDestroyWorker
+  end
 end

ee/app/services/geo/repository_registry_removal_service.rb

@@ -10,7 +10,7 @@ module Geo
 
     # @replicator [Gitlab::Geo::Replicator] Gitlab Geo Replicator
     # @params [Hash] Should include keys: full_path, repository_storage, disk_path
-    def initialize(replicator, params)
+    def initialize(replicator, params = {})
       @replicator = replicator
       @params = params
       @full_path = params[:full_path]
@@ -25,6 +25,10 @@ module Geo
     private
 
     def destroy_repository
+      # We don't have repository location information after main DB record deletion.
+      # The issue https://gitlab.com/gitlab-org/gitlab/-/issues/281430
+      return unless full_path
+
       repository = Repository.new(params[:disk_path], self, shard: params[:repository_storage])
       result = Repositories::DestroyService.new(repository).execute
 

ee/app/workers/all_queues.yml

@@ -387,6 +387,14 @@
   :weight: 1
   :idempotent:
   :tags: []
+- :name: geo:geo_framework_repository_destroy
+  :feature_category: :geo_replication
+  :has_external_dependencies:
+  :urgency: :low
+  :resource_boundary: :unknown
+  :weight: 1
+  :idempotent: true
+  :tags: []
 - :name: geo:geo_hashed_storage_attachments_migration
   :feature_category: :geo_replication
   :has_external_dependencies:

ee/app/workers/geo/framework_repository_destroy_worker.rb

+# frozen_string_literal: true
+
+module Geo
+  class FrameworkRepositoryDestroyWorker
+    include ApplicationWorker
+    include GeoQueue
+    include ::Gitlab::Geo::LogHelpers
+
+    idempotent!
+
+    loggable_arguments 0
+
+    def perform(replicable_name, replicable_id)
+      log_info('Executing Geo::FrameworkRepositoryDestroyWorker', replicable_id: replicable_id, replicable_name: replicable_name)
+
+      replicator = Gitlab::Geo::Replicator.for_replicable_params(replicable_name: replicable_name, replicable_id: replicable_id)
+
+      ::Geo::RepositoryRegistryRemovalService.new(replicator).execute
+    end
+  end
+end

ee/app/workers/geo/secondary/registry_consistency_worker.rb

@@ -24,7 +24,8 @@ module Geo
         Geo::PackageFileRegistry,
         Geo::ProjectRegistry,
         Geo::TerraformStateVersionRegistry,
-        Geo::UploadRegistry
+        Geo::UploadRegistry,
+        Geo::SnippetRepositoryRegistry
       ].freeze
 
       BATCH_SIZE = 10000

ee/spec/models/snippet_repository_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe SnippetRepository, :request_store, :geo, type: :model do
+  include EE::GeoHelpers
+
+  let(:node) { create(:geo_node) }
+
+  before do
+    stub_current_geo_node(node)
+  end
+
+  describe '#replicables_for_current_secondary' do
+    let(:group_1) { create(:group) }
+    let(:group_2) { create(:group) }
+    let(:nested_group_1) { create(:group, parent: group_1) }
+
+    let!(:project_1) { create(:project, group: group_1) }
+    let!(:project_2) { create(:project, group: group_2) }
+
+    let!(:project_snippet_1) { create(:project_snippet, project: project_1) }
+    let!(:project_snippet_2) { create(:project_snippet, project: project_2) }
+    let!(:personal_snippet_1) { create(:personal_snippet) }
+
+    let!(:snippet_repository_1) { create(:snippet_repository, snippet: project_snippet_1) }
+    let!(:snippet_repository_2) { create(:snippet_repository, snippet: personal_snippet_1) }
+
+    let!(:snippet_repository_3) { create(:snippet_repository, shard_name: 'broken') }
+    let!(:snippet_repository_4) { create(:snippet_repository) }
+    let!(:snippet_repository_5) { create(:snippet_repository, snippet: project_snippet_2) }
+
+    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,
+        snippet_repository_2,
+        snippet_repository_3,
+        snippet_repository_4,
+        snippet_repository_5
+      ])
+    end
+
+    context 'with selective sync by namespace' do
+      it 'returns snippet_repositories that belong to the namespaces + personal snippets' do
+        node.update!(selective_sync_type: 'namespaces', namespaces: [group_1])
+
+        expect(described_class.replicables_for_current_secondary(1..described_class.last.id)).to match_array([
+          snippet_repository_1,
+          snippet_repository_2,
+          snippet_repository_3,
+          snippet_repository_4
+        ])
+      end
+    end
+
+    context 'with selective sync by shard' do
+      it 'returns snippet_repositories that belong to the shards' do
+        node.update!(selective_sync_type: 'shards', selective_sync_shards: ['default'])
+
+        expect(described_class.replicables_for_current_secondary(1..described_class.last.id)).to match_array([
+          snippet_repository_1,
+          snippet_repository_2,
+          snippet_repository_4,
+          snippet_repository_5
+        ])
+      end
+    end
+
+    it 'returns nothing if an unrecognised selective sync type is used' do
+      node.update_attribute(:selective_sync_type, 'unknown')
+
+      expect(described_class.replicables_for_current_secondary(1..described_class.last.id)).to be_empty
+    end
+  end
+end

ee/spec/services/geo/registry_consistency_service_spec.rb

@@ -19,6 +19,7 @@ RSpec.describe Geo::RegistryConsistencyService, :geo, :use_clean_rails_memory_st
     { Geo::DesignRegistry => :project_with_design,
       Geo::MergeRequestDiffRegistry => :external_merge_request_diff,
       Geo::PackageFileRegistry => :package_file_with_file,
+      Geo::SnippetRepositoryRegistry => :snippet_repository,
       Geo::TerraformStateVersionRegistry => :terraform_state_version }
       .fetch(registry_class, default_factory_name)
   end
@@ -172,7 +173,7 @@ RSpec.describe Geo::RegistryConsistencyService, :geo, :use_clean_rails_memory_st
           end
 
           before do
-            model_class.where(id: unused_model_ids).delete_all
+            model_class.where(model_class.primary_key => unused_model_ids).delete_all
           end
 
           it 'deletes unused registries', :sidekiq_inline do