Add latest changes from gitlab-org/gitlab@master

5 files changed, 84 insertions, 17 deletions

diff --git a/app/models/ci/build.rb b/app/models/ci/build.rb
index 71ca09717d9..34b5b637422 100644
--- a/app/models/ci/build.rb
+++ b/app/models/ci/build.rb
@@ -172,8 +172,6 @@ module Ci
 
     add_authentication_token_field :token, encrypted: :required
 
-    before_save :ensure_token, unless: :assign_token_on_scheduling?
-
     after_save :stick_build_if_status_changed
 
     after_create unless: :importing? do |build|
@@ -247,11 +245,8 @@ module Ci
         !build.waiting_for_deployment_approval? # If false is returned, it stops the transition
       end
 
-      before_transition any => [:pending] do |build, transition|
-        if build.assign_token_on_scheduling?
-          build.ensure_token
-        end
-
+      before_transition any => [:pending] do |build|
+        build.ensure_token
         true
       end
 
@@ -1140,10 +1135,6 @@ module Ci
       end
     end
 
-    def assign_token_on_scheduling?
-      ::Feature.enabled?(:ci_assign_job_token_on_scheduling, project)
-    end
-
     protected
 
     def run_status_commit_hooks!
diff --git a/app/models/ci/sources/pipeline.rb b/app/models/ci/sources/pipeline.rb
index 2df504cd3de..855e68d1db1 100644
--- a/app/models/ci/sources/pipeline.rb
+++ b/app/models/ci/sources/pipeline.rb
@@ -3,6 +3,7 @@
 module Ci
   module Sources
     class Pipeline < Ci::ApplicationRecord
+      include Ci::Partitionable
       include Ci::NamespacedModelName
 
       self.table_name = "ci_sources_pipelines"
@@ -15,6 +16,11 @@ module Ci
       belongs_to :source_bridge, class_name: "Ci::Bridge", foreign_key: :source_job_id
       belongs_to :source_pipeline, class_name: "Ci::Pipeline", foreign_key: :source_pipeline_id
 
+      partitionable scope: :pipeline
+
+      before_validation :set_source_partition_id, on: :create
+      validates :source_partition_id, presence: true
+
       validates :project, presence: true
       validates :pipeline, presence: true
 
@@ -23,6 +29,15 @@ module Ci
       validates :source_pipeline, presence: true
 
       scope :same_project, -> { where(arel_table[:source_project_id].eq(arel_table[:project_id])) }
+
+      private
+
+      def set_source_partition_id
+        return if source_partition_id_changed? && source_partition_id.present?
+        return unless source_job
+
+        self.source_partition_id = source_job.partition_id
+      end
     end
   end
 end
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
diff --git a/app/models/work_items/parent_link.rb b/app/models/work_items/parent_link.rb
index 2c7f4212de3..9df2f026bd0 100644
--- a/app/models/work_items/parent_link.rb
+++ b/app/models/work_items/parent_link.rb
@@ -19,6 +19,8 @@ module WorkItems
     validate :validate_max_children
     validate :validate_confidentiality
 
+    scope :for_parents, ->(parent_ids) { where(work_item_parent_id: parent_ids) }
+
     class << self
       def has_public_children?(parent_id)
         joins(:work_item).where(work_item_parent_id: parent_id, 'issues.confidential': false).exists?