Refactor queries for Geo::ExpireUploadsFinder

These changes follow our guidelines for reusing abstractions

ee/app/finders/geo/expire_uploads_finder.rb

@@ -4,7 +4,7 @@ module Geo
   class ExpireUploadsFinder
     def find_project_uploads(project)
       if Gitlab::Geo::Fdw.enabled?
-        fdw_find_project_uploads(project)
+        Geo::Fdw::Upload.for_model_with_type(project, 'file')
       else
         legacy_find_project_uploads(project)
       end
@@ -12,7 +12,9 @@ module Geo
 
     def find_file_registries_uploads(project)
       if Gitlab::Geo::Fdw.enabled?
-        fdw_find_file_registries_uploads(project)
+        Gitlab::Geo::Fdw::FileRegistryQueryBuilder.new
+          .for_model(project)
+          .with_type('file')
       else
         legacy_find_file_registries_uploads(project)
       end
@@ -20,47 +22,9 @@ module Geo
 
     private
 
-    #
-    # FDW accessors
-    #
-
-    # @return [ActiveRecord::Relation<Geo::Fdw::Upload>]
-    # rubocop: disable CodeReuse/ActiveRecord
-    def fdw_find_project_uploads(project)
-      fdw_table = Geo::Fdw::Upload.table_name
-      upload_type = 'file'
-
-      Geo::Fdw::Upload.joins("JOIN file_registry
-                                ON file_registry.file_id = #{fdw_table}.id
-                               AND #{fdw_table}.model_id='#{project.id}'
-                               AND #{fdw_table}.model_type='#{project.class.name}'
-                               AND file_registry.file_type='#{upload_type}'")
-    end
-    # rubocop: enable CodeReuse/ActiveRecord
-
-    # @return [ActiveRecord::Relation<Geo::FileRegistry>]
-    # rubocop: disable CodeReuse/ActiveRecord
-    def fdw_find_file_registries_uploads(project)
-      fdw_table = Geo::Fdw::Upload.table_name
-      upload_type = 'file'
-
-      Geo::FileRegistry.joins("JOIN #{fdw_table}
-                                 ON file_registry.file_id = #{fdw_table}.id
-                                AND #{fdw_table}.model_id='#{project.id}'
-                                AND #{fdw_table}.model_type='#{project.class.name}'
-                                AND file_registry.file_type='#{upload_type}'")
-    end
-    # rubocop: enable CodeReuse/ActiveRecord
-
-    #
-    # Legacy accessors (non FDW)
-    #
-
-    # @return [ActiveRecord::Relation<Geo::FileRegistry>] list of file registry items
-    # rubocop: disable CodeReuse/ActiveRecord
+    # rubocop:disable CodeReuse/ActiveRecord
     def legacy_find_file_registries_uploads(project)
-      upload_ids = Upload.where(model_type: project.class.name, model_id: project.id).pluck(:id)
-
+      upload_ids = Upload.for_model(project).pluck_primary_key
       return Geo::FileRegistry.none if upload_ids.empty?
 
       values_sql = upload_ids.map { |id| "(#{id})" }.join(',')
@@ -73,13 +37,11 @@ module Geo
          AND file_registry.file_type='#{upload_type}'
       SQL
     end
-    # rubocop: enable CodeReuse/ActiveRecord
+    # rubocop:enable CodeReuse/ActiveRecord
 
-    # @return [ActiveRecord::Relation<Upload>] list of upload files
-    # rubocop: disable CodeReuse/ActiveRecord
+    # rubocop:disable CodeReuse/ActiveRecord
     def legacy_find_project_uploads(project)
       file_registry_ids = legacy_find_file_registries_uploads(project).pluck(:file_id)
-
       return Upload.none if file_registry_ids.empty?
 
       values_sql = file_registry_ids.map { |f_id| "(#{f_id})" }.join(',')
@@ -90,6 +52,6 @@ module Geo
           ON (file_registry.file_id = uploads.id)
       SQL
     end
-    # rubocop: enable CodeReuse/ActiveRecord
+    # rubocop:enable CodeReuse/ActiveRecord
   end
 end

ee/app/models/ee/upload.rb

@@ -11,6 +11,7 @@ module EE
     prepended do
       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 }
     end
 

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

@@ -13,6 +13,12 @@ module Geo
       scope :geo_syncable, -> { with_files_stored_locally }
 
       class << self
+        def for_model_with_type(model, type)
+          inner_join_file_registry
+            .where(model_id: model.id, model_type: model.class.name)
+            .merge(Geo::FileRegistry.with_file_type(type))
+        end
+
         # Searches for a list of uploads based on the query given in `query`.
         #
         # On PostgreSQL this method uses "ILIKE" to perform a case-insensitive
@@ -22,6 +28,21 @@ module Geo
         def search(query)
           fuzzy_search(query, [:path])
         end
+
+        private
+
+        def inner_join_file_registry
+          join_statement =
+            arel_table
+              .join(file_registry_table, Arel::Nodes::InnerJoin)
+              .on(arel_table[:id].eq(file_registry_table[:file_id]))
+
+          joins(join_statement.join_sources)
+        end
+
+        def file_registry_table
+          Geo::FileRegistry.arel_table
+        end
       end
     end
   end

ee/app/models/geo/file_registry.rb

@@ -8,6 +8,7 @@ class Geo::FileRegistry < Geo::BaseRegistry
   scope :failed, -> { where(success: false).where.not(retry_count: nil) }
   scope :never, -> { where(success: false, retry_count: nil) }
   scope :fresh, -> { order(created_at: :desc) }
+  scope :with_file_type, ->(type) { where(file_type: type) }
 
   self.inheritance_column = 'file_type'
 

ee/lib/gitlab/geo/fdw/file_registry_query_builder.rb

+# frozen_string_literal: true
+
+# Builder class to create composable queries using FDW to
+# retrieve file registries.
+#
+# Basic usage:
+#
+#     Gitlab::Geo::Fdw::FileRegistryQueryBuilder
+#       .new
+#       .for_project_with_type(project, 'file')
+#
+module Gitlab
+  module Geo
+    class Fdw
+      class FileRegistryQueryBuilder < SimpleDelegator
+        attr_reader :query
+
+        def initialize(query = nil)
+          @query = query || base_query
+          super(query)
+        end
+
+        # rubocop:disable CodeReuse/ActiveRecord
+        def for_model(model)
+          reflect(
+            query
+              .joins(fdw_inner_join_uploads)
+              .where(
+                fdw_upload_table[:model_id].eq(model.id)
+                  .and(fdw_upload_table[:model_type].eq(model.class.name))
+              )
+          )
+        end
+        # rubocop:enable CodeReuse/ActiveRecord
+
+        def with_type(type)
+          reflect(query.merge(::Geo::FileRegistry.with_file_type(type)))
+        end
+
+        private
+
+        def base_query
+          ::Geo::FileRegistry
+        end
+
+        def reflect(query)
+          self.class.new(query)
+        end
+
+        def file_registry_table
+          ::Geo::FileRegistry.arel_table
+        end
+
+        def fdw_upload_table
+          ::Geo::Fdw::Upload.arel_table
+        end
+
+        def fdw_inner_join_uploads
+          file_registry_table
+            .join(fdw_upload_table, Arel::Nodes::InnerJoin)
+            .on(file_registry_table[:file_id].eq(fdw_upload_table[:id]))
+            .join_sources
+        end
+      end
+    end
+  end
+end

ee/spec/lib/gitlab/geo/fdw/file_registry_query_builder_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+# Disable transactions via :delete method because a foreign table
+# can't see changes inside a transaction of a different connection.
+describe Gitlab::Geo::Fdw::FileRegistryQueryBuilder, :geo, :delete do
+  let(:project) { create(:project) }
+  let(:upload_1) { create(:upload, :issuable_upload, model: project) }
+  let(:upload_2) { create(:upload, :issuable_upload, model: project) }
+  let(:upload_3) { create(:upload, :issuable_upload) }
+  let!(:file_registry_1) { create(:geo_file_registry, file_id: upload_1.id) }
+  let!(:file_registry_2) { create(:geo_file_registry, :attachment, file_id: upload_2.id) }
+  let!(:file_registry_3) { create(:geo_file_registry, file_id: upload_3.id) }
+
+  before do
+    skip('FDW is not configured') if Gitlab::Database.postgresql? && !Gitlab::Geo::Fdw.enabled?
+  end
+
+  describe '#for_model' do
+    it 'returns registries that upload belong to the model' do
+      expect(subject.for_model(project)).to match_ids(file_registry_1, file_registry_2)
+    end
+  end
+
+  describe '#with_type' do
+    it 'returns registries filtered by file_type' do
+      expect(subject.with_type('file')).to match_ids(file_registry_1, file_registry_3)
+    end
+  end
+end