Merge branch '11000-geo-implement-selective-sync-support-for-the-uploads-fdw-queries' into 'master'

Geo - Refactor FDW queries in the Geo::ExpireUploadsFinder

Closes #11000

See merge request gitlab-org/gitlab-ee!10654

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,53 +12,19 @@ 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
     end
 
-    #
-    # FDW accessors
-    #
+    private
 
-    # @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(',')
@@ -71,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(',')
@@ -88,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/base_query_builder.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Geo
+    class Fdw
+      class BaseQueryBuilder < SimpleDelegator
+        def initialize(query = nil)
+          @query = query || base
+          super(query)
+        end
+
+        private
+
+        attr_reader :query
+
+        def base
+          raise NotImplementedError
+        end
+
+        def reflect(query)
+          self.class.new(query)
+        end
+      end
+    end
+  end
+end

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 < BaseQueryBuilder
+        # 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
+          ::Geo::FileRegistry.select(file_registry_table[Arel.star])
+        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/lib/gitlab/geo/fdw/project_registry_query_builder.rb

@@ -13,14 +13,7 @@
 module Gitlab
   module Geo
     class Fdw
-      class ProjectRegistryQueryBuilder < SimpleDelegator
-        attr_reader :query
-
-        def initialize(query = nil)
-          @query = query || base_query
-          super(query)
-        end
-
+      class ProjectRegistryQueryBuilder < BaseQueryBuilder
         # rubocop:disable CodeReuse/ActiveRecord
         def registries_pending_verification
           reflect(
@@ -54,14 +47,10 @@ module Gitlab
 
         private
 
-        def base_query
+        def base
           ::Geo::ProjectRegistry.select(project_registries_table[Arel.star])
         end
 
-        def reflect(query)
-          self.class.new(query)
-        end
-
         def project_registries_table
           ::Geo::ProjectRegistry.arel_table
         end

ee/spec/finders/geo/expire_uploads_finder_spec.rb

 require 'spec_helper'
 
 describe Geo::ExpireUploadsFinder, :geo do
+  include EE::GeoHelpers
+
   let(:project) { create(:project) }
 
   # Disable transactions via :delete method because a foreign table
@@ -11,23 +13,14 @@ describe Geo::ExpireUploadsFinder, :geo do
     end
 
     describe '#find_project_uploads' do
-      let(:project) { build_stubbed(:project) }
-
-      it 'delegates to #fdw_find_project_uploads' do
-        expect(subject).to receive(:fdw_find_project_uploads).with(project)
-
-        subject.find_project_uploads(project)
-      end
-    end
-
-    describe '#fdw_find_project_uploads' do
       context 'filtering per project uploads' do
         it 'returns only objects associated with the project' do
           other_upload = create(:upload, :issuable_upload)
           upload = create(:upload, :issuable_upload, model: project)
           create(:geo_file_registry, file_id: upload.id)
           create(:geo_file_registry, file_id: other_upload.id)
-          uploads = subject.fdw_find_project_uploads(project)
+
+          uploads = subject.find_project_uploads(project)
 
           expect(uploads.count).to eq(1)
           expect(uploads.first.id).to eq(upload.id)
@@ -39,7 +32,8 @@ describe Geo::ExpireUploadsFinder, :geo do
           create(:upload, :issuable_upload, model: project)
           upload = create(:upload, :issuable_upload, model: project)
           create(:geo_file_registry, file_id: upload.id, success: false)
-          uploads = subject.fdw_find_project_uploads(project)
+
+          uploads = subject.find_project_uploads(project)
 
           expect(uploads.count).to eq(1)
           expect(uploads.first.id).to eq(upload.id)
@@ -48,23 +42,14 @@ describe Geo::ExpireUploadsFinder, :geo do
     end
 
     describe '#find_file_registries_uploads' do
-      let(:project) { build_stubbed(:project) }
-
-      it 'delegates to #fdw_find_file_registries_uploads' do
-        expect(subject).to receive(:fdw_find_file_registries_uploads).with(project)
-
-        subject.find_file_registries_uploads(project)
-      end
-    end
-
-    describe '#fdw_find_file_registries_uploads' do
       context 'filtering per project uploads' do
         it 'returns only objects associated with the project' do
           other_upload = create(:upload, :issuable_upload)
           upload = create(:upload, :issuable_upload, model: project)
           create(:geo_file_registry, file_id: other_upload.id)
           file_registry = create(:geo_file_registry, file_id: upload.id)
-          files = subject.fdw_find_file_registries_uploads(project)
+
+          files = subject.find_file_registries_uploads(project)
 
           expect(files.count).to eq(1)
           expect(files.first.id).to eq(file_registry.id)
@@ -75,27 +60,18 @@ describe Geo::ExpireUploadsFinder, :geo do
 
   context 'Legacy' do
     before do
-      allow(Gitlab::Geo::Fdw).to receive(:enabled?).and_return(false)
+      stub_fdw_disabled
     end
 
     describe '#find_project_uploads' do
-      let(:project) { build_stubbed(:project) }
-
-      it 'delegates to #legacy_find_project_uploads' do
-        expect(subject).to receive(:legacy_find_project_uploads).with(project)
-
-        subject.find_project_uploads(project)
-      end
-    end
-
-    describe '#legacy_find_project_uploads' do
       context 'filtering per project uploads' do
         it 'returns only objects associated with the project' do
           other_upload = create(:upload, :issuable_upload)
           upload = create(:upload, :issuable_upload, model: project)
           create(:geo_file_registry, file_id: upload.id)
           create(:geo_file_registry, file_id: other_upload.id)
-          uploads = subject.legacy_find_project_uploads(project)
+
+          uploads = subject.find_project_uploads(project)
 
           expect(uploads.count).to eq(1)
           expect(uploads.first.id).to eq(upload.id)
@@ -107,7 +83,8 @@ describe Geo::ExpireUploadsFinder, :geo do
           create(:upload, :issuable_upload, model: project)
           upload = create(:upload, :issuable_upload, model: project)
           create(:geo_file_registry, file_id: upload.id, success: false)
-          uploads = subject.legacy_find_project_uploads(project)
+
+          uploads = subject.find_project_uploads(project)
 
           expect(uploads.count).to eq(1)
           expect(uploads.first.id).to eq(upload.id)
@@ -116,23 +93,14 @@ describe Geo::ExpireUploadsFinder, :geo do
     end
 
     describe '#find_file_registries_uploads' do
-      let(:project) { build_stubbed(:project) }
-
-      it 'delegates to #legacy_find_file_registries_uploads' do
-        expect(subject).to receive(:legacy_find_file_registries_uploads).with(project)
-
-        subject.find_file_registries_uploads(project)
-      end
-    end
-
-    describe '#legacy_find_file_registries_uploads' do
       context 'filtering per project uploads' do
         it 'returns only objects associated with the project' do
           other_upload = create(:upload, :issuable_upload)
           upload = create(:upload, :issuable_upload, model: project)
           create(:geo_file_registry, file_id: other_upload.id)
           file_registry = create(:geo_file_registry, file_id: upload.id)
-          files = subject.legacy_find_file_registries_uploads(project)
+
+          files = subject.find_file_registries_uploads(project)
 
           expect(files.count).to eq(1)
           expect(files.first.id).to eq(file_registry.id)

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