Keep locking_version when updating partitioned records

Optimistic locking is overriding the query_constraints method
and our patch was overriding it instead of overriding the
original method and the locking_version attribute was
lost during the update

Changelog: fixed

gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning.rb

@@ -34,9 +34,12 @@ module Partitioning
         ::ActiveRecord::Reflection::MacroReflection.prepend(
           ActiveRecord::GitlabPatches::Partitioning::Reflection::MacroReflection
         )
-        ::ActiveRecord::Base.prepend(
+        ::ActiveRecord::Persistence.prepend(
           ActiveRecord::GitlabPatches::Partitioning::Base
         )
+        ::ActiveRecord::Persistence::ClassMethods.prepend(
+          ActiveRecord::GitlabPatches::Partitioning::Base::ClassMethods
+        )
       end
     end
   end

gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/base.rb

@@ -4,19 +4,18 @@
   raise 'New version of active-record detected, please remove or update this patch'
 end
 
+# rubocop:disable Gitlab/ModuleWithInstanceVariables
 module ActiveRecord
   module GitlabPatches
     module Partitioning
       module Base
-        extend ActiveSupport::Concern
-
         def _query_constraints_hash
           constraints_hash = super
 
           return constraints_hash unless self.class.use_partition_id_filter?
 
           if self.class.query_constraints_list.nil?
-            { @primary_key => id_in_database } # rubocop:disable Gitlab/ModuleWithInstanceVariables
+            { @primary_key => id_in_database }
           else
             self.class.query_constraints_list.index_with do |column_name|
               attribute_in_database(column_name)
@@ -24,7 +23,7 @@ def _query_constraints_hash
           end
         end
 
-        class_methods do
+        module ClassMethods
           def use_partition_id_filter?
             false
           end
@@ -47,3 +46,4 @@ def query_constraints_list # :nodoc:
     end
   end
 end
+# rubocop:enable Gitlab/ModuleWithInstanceVariables

gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/locking_spec.rb

+# frozen_string_literal: true
+
+RSpec.describe 'ActiveRecord::GitlabPatches::Partitioning::Associations::Locking', :partitioning do
+  let!(:job) { LockingJob.create!(partition_id: 100) }
+
+  describe 'optimistic locking' do
+    it 'does not use lock version on unrelated updates' do
+      update_statement = <<~SQL.squish
+        UPDATE "locking_jobs" SET "name" = 'test'
+        WHERE "locking_jobs"."id" = #{job.id} AND "locking_jobs"."partition_id" = #{job.partition_id}
+      SQL
+
+      result = QueryRecorder.log do
+        job.update!(name: 'test')
+      end
+
+      expect(result).to include(update_statement)
+    end
+
+    it 'uses lock version when status changes' do
+      update_statement = <<~SQL.squish
+        UPDATE "locking_jobs"
+        SET "status" = 1, "name" = 'test', "lock_version" = 1
+        WHERE "locking_jobs"."id" = 1 AND "locking_jobs"."partition_id" = 100 AND "locking_jobs"."lock_version" = 0
+      SQL
+
+      result = QueryRecorder.log do
+        job.update!(name: 'test', status: :completed)
+      end
+
+      expect(result).to include(update_statement)
+    end
+  end
+end

gems/activerecord-gitlab/spec/support/database.rb

@@ -24,6 +24,14 @@
         t.integer :partition_id
         t.boolean :test_flag, default: false
       end
+
+      create_table :locking_jobs, force: true do |t|
+        t.integer :pipeline_id
+        t.integer :partition_id
+        t.integer :lock_version, default: 0, null: false
+        t.integer :status, default: 0, null: false
+        t.string :name
+      end
     end
   end
 end

gems/activerecord-gitlab/spec/support/models.rb

@@ -48,3 +48,14 @@ class Metadata < PartitionedRecord
   belongs_to :job,
     ->(metadata) { where(partition_id: metadata.partition_id) }
 end
+
+class LockingJob < PartitionedRecord
+  self.table_name = :locking_jobs
+  query_constraints :id, :partition_id
+
+  enum status: { created: 0, completed: 1 }
+
+  def locking_enabled?
+    will_save_change_to_status?
+  end
+end