diff --git a/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning.rb b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning.rb
index cf0bd4849e2d71f612383595fc204fbab27cc0f7..2f3dde2a871372a8e2cb64738cad024032065bfd 100644
--- a/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning.rb
+++ b/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
diff --git a/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/base.rb b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/base.rb
index 0c8a248b984fe9a9ca643426e810b245fc4c7d97..c4b3528a526cc1e8ce267f6c1c54c58da25e311b 100644
--- a/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/base.rb
+++ b/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
diff --git a/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/locking_spec.rb b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/locking_spec.rb
new file mode 100644
index 0000000000000000000000000000000000000000..df9241e0983e7b08ddd4314caf073e050c827a3a
--- /dev/null
+++ b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/locking_spec.rb
@@ -0,0 +1,34 @@
+# 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
diff --git a/gems/activerecord-gitlab/spec/support/database.rb b/gems/activerecord-gitlab/spec/support/database.rb
index 998d945c311505c72107c8de71a8dd1c8fc1e8d6..7b4de82312d8bae2ff1a7d20d1022bf424349bed 100644
--- a/gems/activerecord-gitlab/spec/support/database.rb
+++ b/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
diff --git a/gems/activerecord-gitlab/spec/support/models.rb b/gems/activerecord-gitlab/spec/support/models.rb
index c0017656ea814d87fef9e81d73e06abdcc5a88e3..445dc8b2ac24b6b1dc89bc653f9d3290e451ee1b 100644
--- a/gems/activerecord-gitlab/spec/support/models.rb
+++ b/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