1 files changed, 10 insertions, 5 deletions
diff --git a/app/models/ci/build_need.rb b/app/models/ci/build_need.rb
index 1831b7868f9..54a54c42fd1 100644
--- a/app/models/ci/build_need.rb
+++ b/app/models/ci/build_need.rb
@@ -4,14 +4,16 @@ module Ci
   class BuildNeed < Ci::ApplicationRecord
     include Ci::Partitionable
     include IgnorableColumns
-    include SafelyChangeColumnDefault
     include BulkInsertSafe
 
     MAX_JOB_NAME_LENGTH = 255
 
-    columns_changing_default :partition_id
-
-    belongs_to :build, class_name: "Ci::Processable", foreign_key: :build_id, inverse_of: :needs
+    belongs_to :build,
+      ->(need) { in_partition(need) },
+      class_name: 'Ci::Processable',
+      foreign_key: :build_id,
+      partition_foreign_key: :partition_id,
+      inverse_of: :needs
 
     partitionable scope: :build
 
@@ -19,7 +21,10 @@ module Ci
     validates :name, presence: true, length: { maximum: MAX_JOB_NAME_LENGTH }
     validates :optional, inclusion: { in: [true, false] }
 
-    scope :scoped_build, -> { where("#{Ci::Build.quoted_table_name}.id = #{quoted_table_name}.build_id") }
+    scope :scoped_build, -> {
+      where(arel_table[:build_id].eq(Ci::Build.arel_table[:id]))
+      .where(arel_table[:partition_id].eq(Ci::Build.arel_table[:partition_id]))
+    }
     scope :artifacts, -> { where(artifacts: true) }
   end
 end