1 files changed, 37 insertions, 10 deletions

diff --git a/app/models/ci/stage.rb b/app/models/ci/stage.rb
index 3d2df9a45ef..becb8f204bf 100644
--- a/app/models/ci/stage.rb
+++ b/app/models/ci/stage.rb
@@ -7,11 +7,8 @@ module Ci
     include Ci::HasStatus
     include Gitlab::OptimisticLocking
     include Presentable
-    include SafelyChangeColumnDefault
     include IgnorableColumns
 
-    columns_changing_default :partition_id
-
     ignore_column :pipeline_id_convert_to_bigint, remove_with: '16.6', remove_after: '2023-10-22'
 
     partitionable scope: :pipeline
@@ -21,19 +18,45 @@ module Ci
     belongs_to :project
     belongs_to :pipeline
 
-    has_many :statuses, class_name: 'CommitStatus', foreign_key: :stage_id, inverse_of: :ci_stage
-    has_many :latest_statuses, -> { ordered.latest },
+    has_many :statuses,
+      ->(stage) { in_partition(stage) },
+      class_name: 'CommitStatus',
+      foreign_key: :stage_id,
+      partition_foreign_key: :partition_id,
+      inverse_of: :ci_stage
+    has_many :latest_statuses,
+      ->(stage) { in_partition(stage).ordered.latest },
       class_name: 'CommitStatus',
       foreign_key: :stage_id,
+      partition_foreign_key: :partition_id,
       inverse_of: :ci_stage
-    has_many :retried_statuses, -> { ordered.retried },
+    has_many :retried_statuses,
+      ->(stage) { in_partition(stage).ordered.retried },
       class_name: 'CommitStatus',
       foreign_key: :stage_id,
+      partition_foreign_key: :partition_id,
+      inverse_of: :ci_stage
+    has_many :processables,
+      ->(stage) { in_partition(stage) },
+      class_name: 'Ci::Processable',
+      foreign_key: :stage_id,
+      partition_foreign_key: :partition_id,
+      inverse_of: :ci_stage
+    has_many :builds,
+      ->(stage) { in_partition(stage) },
+      foreign_key: :stage_id,
+      partition_foreign_key: :partition_id,
+      inverse_of: :ci_stage
+    has_many :bridges,
+      ->(stage) { in_partition(stage) },
+      foreign_key: :stage_id,
+      partition_foreign_key: :partition_id,
+      inverse_of: :ci_stage
+    has_many :generic_commit_statuses,
+      ->(stage) { in_partition(stage) },
+      foreign_key: :stage_id,
+      partition_foreign_key: :partition_id,
       inverse_of: :ci_stage
-    has_many :processables, class_name: 'Ci::Processable', foreign_key: :stage_id, inverse_of: :ci_stage
-    has_many :builds, foreign_key: :stage_id, inverse_of: :ci_stage
-    has_many :bridges, foreign_key: :stage_id, inverse_of: :ci_stage
-    has_many :generic_commit_statuses, foreign_key: :stage_id, inverse_of: :ci_stage
 
     scope :ordered, -> { order(position: :asc) }
     scope :in_pipelines, ->(pipelines) { where(pipeline: pipelines) }
@@ -107,6 +130,10 @@ module Ci
       end
     end
 
+    def self.use_partition_id_filter?
+      Ci::Pipeline.use_partition_id_filter?
+    end
+
     def set_status(new_status)
       retry_optimistic_lock(self, name: 'ci_stage_set_status') do
         case new_status