Move overridden DB Load Balancing methods (project-related)

Move Project#mark_primary_write_location to Core
Move non-stickiness in ProjectFeatureUsage to Core

app/models/project.rb

@@ -2360,7 +2360,7 @@ class Project < ApplicationRecord
   end
 
   def mark_primary_write_location
-    # Overriden in EE
+    ::Gitlab::Database::LoadBalancing::Sticking.mark_primary_write_location(:project, self.id)
   end
 
   def toggle_ci_cd_settings!(settings_attribute)

app/models/project_feature_usage.rb

@@ -20,14 +20,16 @@ class ProjectFeatureUsage < ApplicationRecord
   end
 
   def log_jira_dvcs_integration_usage(cloud: true)
-    integration_field = self.class.jira_dvcs_integration_field(cloud: cloud)
+    ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes do
+      integration_field = self.class.jira_dvcs_integration_field(cloud: cloud)
 
-    # The feature usage is used only once later to query the feature usage in a
-    # long date range. Therefore, we just need to update the timestamp once per
-    # day
-    return if persisted? && updated_today?(integration_field)
+      # The feature usage is used only once later to query the feature usage in a
+      # long date range. Therefore, we just need to update the timestamp once per
+      # day
+      break if persisted? && updated_today?(integration_field)
 
-    persist_jira_dvcs_usage(integration_field)
+      persist_jira_dvcs_usage(integration_field)
+    end
   end
 
   private

ee/app/models/ee/project.rb

@@ -410,11 +410,6 @@ module EE
         feature_available?(:github_project_service_integration)
     end
 
-    override :mark_primary_write_location
-    def mark_primary_write_location
-      ::Gitlab::Database::LoadBalancing::Sticking.mark_primary_write_location(:project, self.id)
-    end
-
     override :add_import_job
     def add_import_job
       return if gitlab_custom_project_template_import?

ee/app/models/ee/project_feature_usage.rb

-# frozen_string_literal: true
-
-module EE
-  module ProjectFeatureUsage
-    extend ActiveSupport::Concern
-    extend ::Gitlab::Utils::Override
-
-    override :log_jira_dvcs_integration_usage
-    def log_jira_dvcs_integration_usage(**options)
-      ::Gitlab::Database::LoadBalancing::Session.without_sticky_writes do
-        super(**options)
-      end
-    end
-  end
-end

ee/spec/models/project_feature_usage_spec.rb

-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe ProjectFeatureUsage, :request_store do
-  include_context 'clear DB Load Balancing configuration'
-
-  describe '#log_jira_dvcs_integration_usage' do
-    let!(:project) { create(:project) }
-
-    subject { project.feature_usage }
-
-    context 'database load balancing is configured' do
-      before do
-        # Do not pollute AR for other tests, but rather simulate effect of configure_proxy.
-        allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
-        ::Gitlab::Database::LoadBalancing.configure_proxy
-        allow(ActiveRecord::Base).to receive(:connection).and_return(::Gitlab::Database::LoadBalancing.proxy)
-        ::Gitlab::Database::LoadBalancing::Session.clear_session
-      end
-
-      it 'logs Jira DVCS Cloud last sync' do
-        freeze_time do
-          subject.log_jira_dvcs_integration_usage
-
-          expect(subject.jira_dvcs_server_last_sync_at).to be_nil
-          expect(subject.jira_dvcs_cloud_last_sync_at).to be_like_time(Time.current)
-        end
-      end
-
-      it 'does not stick to primary' do
-        expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_performed_write
-        expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_using_primary
-
-        subject.log_jira_dvcs_integration_usage
-
-        expect(::Gitlab::Database::LoadBalancing::Session.current).to be_performed_write
-        expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_using_primary
-      end
-    end
-
-    context 'database load balancing is not cofigured' do
-      it 'logs Jira DVCS Cloud last sync' do
-        freeze_time do
-          subject.log_jira_dvcs_integration_usage
-
-          expect(subject.jira_dvcs_server_last_sync_at).to be_nil
-          expect(subject.jira_dvcs_cloud_last_sync_at).to be_like_time(Time.current)
-        end
-      end
-    end
-  end
-end

ee/spec/models/project_spec.rb

@@ -2717,14 +2717,6 @@ RSpec.describe Project do
     end
   end
 
-  describe '#mark_primary_write_location' do
-    it 'marks the location with project ID' do
-      expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:mark_primary_write_location).with(:project, project.id)
-
-      project.mark_primary_write_location
-    end
-  end
-
   describe '#add_template_export_job' do
     it 'starts project template export job' do
       user = create(:user)

spec/models/project_feature_usage_spec.rb

@@ -126,4 +126,54 @@ RSpec.describe ProjectFeatureUsage, type: :model do
       end
     end
   end
+
+  context 'ProjectFeatureUsage with DB Load Balancing', :request_store do
+    include_context 'clear DB Load Balancing configuration'
+
+    describe '#log_jira_dvcs_integration_usage' do
+      let!(:project) { create(:project) }
+
+      subject { project.feature_usage }
+
+      context 'database load balancing is configured' do
+        before do
+          # Do not pollute AR for other tests, but rather simulate effect of configure_proxy.
+          allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
+          ::Gitlab::Database::LoadBalancing.configure_proxy
+          allow(ActiveRecord::Base).to receive(:connection).and_return(::Gitlab::Database::LoadBalancing.proxy)
+          ::Gitlab::Database::LoadBalancing::Session.clear_session
+        end
+
+        it 'logs Jira DVCS Cloud last sync' do
+          freeze_time do
+            subject.log_jira_dvcs_integration_usage
+
+            expect(subject.jira_dvcs_server_last_sync_at).to be_nil
+            expect(subject.jira_dvcs_cloud_last_sync_at).to be_like_time(Time.current)
+          end
+        end
+
+        it 'does not stick to primary' do
+          expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_performed_write
+          expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_using_primary
+
+          subject.log_jira_dvcs_integration_usage
+
+          expect(::Gitlab::Database::LoadBalancing::Session.current).to be_performed_write
+          expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_using_primary
+        end
+      end
+
+      context 'database load balancing is not cofigured' do
+        it 'logs Jira DVCS Cloud last sync' do
+          freeze_time do
+            subject.log_jira_dvcs_integration_usage
+
+            expect(subject.jira_dvcs_server_last_sync_at).to be_nil
+            expect(subject.jira_dvcs_cloud_last_sync_at).to be_like_time(Time.current)
+          end
+        end
+      end
+    end
+  end
 end

spec/models/project_spec.rb

@@ -2854,6 +2854,16 @@ RSpec.describe Project, factory_default: :keep do
     end
   end
 
+  describe '#mark_primary_write_location' do
+    let(:project) { create(:project) }
+
+    it 'marks the location with project ID' do
+      expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:mark_primary_write_location).with(:project, project.id)
+
+      project.mark_primary_write_location
+    end
+  end
+
   describe '#mark_stuck_remote_mirrors_as_failed!' do
     it 'fails stuck remote mirrors' do
       project = create(:project, :repository, :remote_mirror)