Add selective sync support for FDW queries to count synced attachments

This changes fixes a bunch of CodeReuse/ActiveRecord violations and make
the queries much easier to understand.

ee/app/finders/geo/attachment_registry_finder.rb

@@ -6,42 +6,30 @@ module Geo
       syncable.count
     end
 
-    def syncable
-      if use_legacy_queries_for_selective_sync?
-        legacy_finder.syncable
-      elsif selective_sync?
-        fdw_all.geo_syncable
-      else
-        Upload.geo_syncable
-      end
-    end
-
     def count_synced
-      if aggregate_pushdown_supported?
-        find_synced.count
-      else
-        legacy_find_synced.count
-      end
+      attachments_synced.count
     end
 
     def count_failed
-      if aggregate_pushdown_supported?
-        find_failed.count
-      else
-        legacy_find_failed.count
-      end
+      attachments_failed.count
     end
 
     def count_synced_missing_on_primary
-      if aggregate_pushdown_supported? && !use_legacy_queries?
-        fdw_find_synced_missing_on_primary.count
-      else
-        legacy_find_synced_missing_on_primary.count
-      end
+      attachments_synced_missing_on_primary.count
     end
 
     def count_registry
-      Geo::FileRegistry.attachments.count
+      registries_for_attachments.count
+    end
+
+    def syncable
+      if use_legacy_queries_for_selective_sync?
+        legacy_finder.syncable
+      elsif selective_sync?
+        attachments.syncable
+      else
+        Upload.syncable
+      end
     end
 
     # Find limited amount of non replicated attachments.
@@ -82,7 +70,9 @@ module Geo
 
     # rubocop: disable CodeReuse/ActiveRecord
     def find_retryable_failed_registries(batch_size:, except_file_ids: [])
-      find_failed_registries
+      Geo::FileRegistry
+        .attachments
+        .failed
         .retry_due
         .file_id_not_in(except_file_ids)
         .limit(batch_size)
@@ -91,7 +81,10 @@ module Geo
 
     # rubocop: disable CodeReuse/ActiveRecord
     def find_retryable_synced_missing_on_primary_registries(batch_size:, except_file_ids: [])
-      find_synced_missing_on_primary_registries
+      Geo::FileRegistry
+        .attachments
+        .synced
+        .missing_on_primary
         .retry_due
         .file_id_not_in(except_file_ids)
         .limit(batch_size)
@@ -110,10 +103,51 @@ module Geo
       @fdw_geo_node ||= Geo::Fdw::GeoNode.find(current_node.id)
     end
 
+    def registries_for_attachments
+      if use_legacy_queries_for_selective_sync?
+        legacy_finder.registries_for_attachments
+      else
+        attachments
+          .inner_join_file_registry
+          .merge(Geo::FileRegistry.attachments)
+      end
+    end
+
+    def attachments_synced
+      if use_legacy_queries_for_selective_sync?
+        legacy_finder.attachments_synced
+      else
+        registries_for_attachments
+          .syncable
+          .merge(Geo::FileRegistry.synced)
+      end
+    end
+
+    def attachments_failed
+      if use_legacy_queries_for_selective_sync?
+        legacy_finder.attachments_failed
+      else
+        registries_for_attachments
+          .syncable
+          .merge(Geo::FileRegistry.failed)
+      end
+    end
+
+    def attachments_synced_missing_on_primary
+      if use_legacy_queries_for_selective_sync?
+        legacy_finder.attachments_synced_missing_on_primary
+      else
+        registries_for_attachments
+          .syncable
+          .merge(Geo::FileRegistry.synced)
+          .merge(Geo::FileRegistry.missing_on_primary)
+      end
+    end
+
     # rubocop: disable CodeReuse/ActiveRecord
-    def fdw_all
+    def attachments
       if selective_sync?
-        Geo::Fdw::Upload.where(fdw_group_uploads.or(fdw_project_uploads).or(fdw_other_uploads))
+        Geo::Fdw::Upload.where(group_uploads.or(project_uploads).or(other_uploads))
       else
         Geo::Fdw::Upload.all
       end
@@ -121,7 +155,7 @@ module Geo
     # rubocop: enable CodeReuse/ActiveRecord
 
     # rubocop: disable CodeReuse/ActiveRecord
-    def fdw_group_uploads
+    def group_uploads
       namespace_ids =
         if current_node.selective_sync_by_namespaces?
           Gitlab::ObjectHierarchy.new(fdw_geo_node.namespaces).base_and_descendants.select(:id)
@@ -141,7 +175,7 @@ module Geo
     end
     # rubocop: enable CodeReuse/ActiveRecord
 
-    def fdw_project_uploads
+    def project_uploads
       project_ids = fdw_geo_node.projects.select(:id)
 
       # This query was intentionally converted to a raw one to get it work in Rails 5.0.
@@ -152,7 +186,7 @@ module Geo
       fdw_upload_table[:model_type].eq('Project').and(project_ids_in_sql)
     end
 
-    def fdw_other_uploads
+    def other_uploads
       fdw_upload_table[:model_type].not_in(%w[Namespace Project])
     end
 
@@ -160,84 +194,21 @@ module Geo
       Geo::Fdw::Upload.arel_table
     end
 
-    def find_synced
-      if use_legacy_queries?
-        legacy_find_synced
-      else
-        fdw_find_synced
-      end
-    end
-
-    def find_failed
-      if use_legacy_queries?
-        legacy_find_failed
-      else
-        fdw_find_failed
-      end
-    end
-
-    def find_synced_registries
-      Geo::FileRegistry.attachments.synced
-    end
-
-    def find_failed_registries
-      Geo::FileRegistry.attachments.failed
-    end
-
-    def find_synced_missing_on_primary_registries
-      find_synced_registries.missing_on_primary
-    end
-
-    def fdw_find_synced
-      fdw_find_syncable.merge(Geo::FileRegistry.synced)
-    end
-
-    def fdw_find_failed
-      fdw_find_syncable.merge(Geo::FileRegistry.failed)
-    end
-
-    def fdw_find_syncable
-      fdw_all
-        .inner_join_file_registry
-        .geo_syncable
-        .merge(Geo::FileRegistry.attachments)
-    end
-
     def fdw_find_unsynced(except_file_ids:)
-      fdw_all
+      attachments
         .missing_file_registry
-        .geo_syncable
+        .syncable
         .id_not_in(except_file_ids)
     end
 
-    def fdw_find_synced_missing_on_primary
-      fdw_find_synced.merge(Geo::FileRegistry.missing_on_primary)
-    end
-
     def fdw_find_migrated_local(except_file_ids:)
-      fdw_all
+      attachments
         .inner_join_file_registry
         .with_files_stored_remotely
         .merge(Geo::FileRegistry.attachments)
         .id_not_in(except_file_ids)
     end
 
-    def legacy_find_synced
-      legacy_inner_join_registry_ids(
-        syncable,
-        find_synced_registries.pluck_file_key,
-        Upload
-      )
-    end
-
-    def legacy_find_failed
-      legacy_inner_join_registry_ids(
-        syncable,
-        find_failed_registries.pluck_file_key,
-        Upload
-      )
-    end
-
     def legacy_find_unsynced(except_file_ids:)
       registry_file_ids = Geo::FileRegistry.attachments.pluck_file_key | except_file_ids
 
@@ -257,13 +228,5 @@ module Geo
         Upload
       )
     end
-
-    def legacy_find_synced_missing_on_primary
-      legacy_inner_join_registry_ids(
-        syncable,
-        find_synced_missing_on_primary_registries.pluck_file_key,
-        Upload
-      )
-    end
   end
 end

ee/app/finders/geo/legacy_attachment_registry_finder.rb

@@ -3,7 +3,7 @@
 module Geo
   class LegacyAttachmentRegistryFinder < RegistryFinder
     def syncable
-      attachments.geo_syncable
+      attachments.syncable
     end
 
     # rubocop:disable CodeReuse/ActiveRecord
@@ -16,6 +16,41 @@ module Geo
     end
     # rubocop:enable CodeReuse/ActiveRecord
 
+    def attachments_synced
+      legacy_inner_join_registry_ids(
+        syncable,
+        Geo::FileRegistry.attachments.synced.pluck_file_key,
+        Upload
+      )
+    end
+
+    def attachments_failed
+      legacy_inner_join_registry_ids(
+        syncable,
+        Geo::FileRegistry.attachments.failed.pluck_file_key,
+        Upload
+      )
+    end
+
+    def attachments_synced_missing_on_primary
+      legacy_inner_join_registry_ids(
+        syncable,
+        Geo::FileRegistry.attachments.synced.missing_on_primary.pluck_file_key,
+        Upload
+      )
+    end
+
+    def registries_for_attachments
+      return Geo::FileRegistry.attachments unless selective_sync?
+
+      legacy_inner_join_registry_ids(
+        Geo::FileRegistry.attachments,
+        attachments.pluck_primary_key,
+        Geo::FileRegistry,
+        foreign_key: :file_id
+      )
+    end
+
     private
 
     # rubocop:disable CodeReuse/ActiveRecord

ee/app/models/ee/upload.rb

@@ -12,7 +12,7 @@ module EE
       after_destroy :log_geo_deleted_event
 
       scope :for_model, ->(model) { where(model_id: model.id, model_type: model.class.name) }
-      scope :geo_syncable, -> { with_files_stored_locally }
+      scope :syncable, -> { with_files_stored_locally }
     end
 
     def log_geo_deleted_event

ee/app/models/geo/fdw/upload.rb

@@ -10,7 +10,7 @@ module Geo
 
       self.table_name = Gitlab::Geo::Fdw.foreign_table_name('uploads')
 
-      scope :geo_syncable, -> { with_files_stored_locally }
+      scope :syncable, -> { with_files_stored_locally }
 
       class << self
         def for_model_with_type(model, type)

ee/changelogs/unreleased/10999-geo-implement-selective-sync-support-for-the-fdw-queries-to-count-attachments.yml

+---
+title: Geo - Implement selective sync support for the FDW queries to count attachments
+merge_request: 11518
+author:
+type: changed