Add latest changes from gitlab-org/gitlab@master

2 files changed, 65 insertions, 6 deletions

diff --git a/app/models/concerns/ci/partitionable.rb b/app/models/concerns/ci/partitionable.rb
index 88692a735ba..d6ba0f4488f 100644
--- a/app/models/concerns/ci/partitionable.rb
+++ b/app/models/concerns/ci/partitionable.rb
@@ -37,6 +37,7 @@ module Ci
         Ci::PendingBuild
         Ci::RunningBuild
         Ci::PipelineVariable
+        Ci::Sources::Pipeline
         Ci::Stage
         Ci::UnitTestFailure
       ].freeze
@@ -67,14 +68,31 @@ module Ci
     end
 
     class_methods do
-      def partitionable(scope:, through: nil)
-        if through
-          define_singleton_method(:routing_table_name) { through[:table] }
-          define_singleton_method(:routing_table_name_flag) { through[:flag] }
+      def partitionable(scope:, through: nil, partitioned: false)
+        handle_partitionable_through(through)
+        handle_partitionable_dml(partitioned)
+        handle_partitionable_scope(scope)
+      end
 
-          include Partitionable::Switch
-        end
+      private
+
+      def handle_partitionable_through(options)
+        return unless options
+
+        define_singleton_method(:routing_table_name) { options[:table] }
+        define_singleton_method(:routing_table_name_flag) { options[:flag] }
+
+        include Partitionable::Switch
+      end
+
+      def handle_partitionable_dml(partitioned)
+        define_singleton_method(:partitioned?) { partitioned }
+        return unless partitioned
+
+        include Partitionable::PartitionedFilter
+      end
 
+      def handle_partitionable_scope(scope)
         define_method(:partition_scope_value) do
           strong_memoize(:partition_scope_value) do
             next Ci::Pipeline.current_partition_value if respond_to?(:importing?) && importing?
diff --git a/app/models/concerns/ci/partitionable/partitioned_filter.rb b/app/models/concerns/ci/partitionable/partitioned_filter.rb
new file mode 100644
index 00000000000..4adae3be26a
--- /dev/null
+++ b/app/models/concerns/ci/partitionable/partitioned_filter.rb
@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module Ci
+  module Partitionable
+    # Used to patch the save, update, delete, destroy methods to use the
+    # partition_id attributes for their SQL queries.
+    module PartitionedFilter
+      extend ActiveSupport::Concern
+
+      if Rails::VERSION::MAJOR >= 7
+        # These methods are updated in Rails 7 to use `_primary_key_constraints_hash`
+        # by default, so this patch will no longer be required.
+        #
+        # rubocop:disable Gitlab/NoCodeCoverageComment
+        # :nocov:
+        raise "`#{__FILE__}` should be double checked" if Rails.env.test?
+
+        warn "Update `#{__FILE__}`. Patches Rails internals for partitioning"
+        # :nocov:
+        # rubocop:enable Gitlab/NoCodeCoverageComment
+      else
+        def _update_row(attribute_names, attempted_action = "update")
+          self.class._update_record(
+            attributes_with_values(attribute_names),
+            _primary_key_constraints_hash
+          )
+        end
+
+        def _delete_row
+          self.class._delete_record(_primary_key_constraints_hash)
+        end
+      end
+
+      # Introduced in Rails 7, but updated to include `partition_id` filter.
+      # https://github.com/rails/rails/blob/a4dbb153fd390ac31bb9808809e7ac4d3a2c5116/activerecord/lib/active_record/persistence.rb#L1031-L1033
+      def _primary_key_constraints_hash
+        { @primary_key => id_in_database, partition_id: partition_id }  # rubocop:disable Gitlab/ModuleWithInstanceVariables
+      end
+    end
+  end
+end