1 files changed, 27 insertions, 0 deletions
diff --git a/app/models/concerns/ci/partitionable.rb b/app/models/concerns/ci/partitionable.rb
index 710ee1ba64f..df803180e77 100644
--- a/app/models/concerns/ci/partitionable.rb
+++ b/app/models/concerns/ci/partitionable.rb
@@ -19,7 +19,32 @@ module Ci
     extend ActiveSupport::Concern
     include ::Gitlab::Utils::StrongMemoize
 
+    module Testing
+      InclusionError = Class.new(StandardError)
+
+      PARTITIONABLE_MODELS = %w[
+        CommitStatus
+        Ci::BuildMetadata
+        Ci::Stage
+        Ci::JobArtifact
+        Ci::PipelineVariable
+        Ci::Pipeline
+      ].freeze
+
+      def self.check_inclusion(klass)
+        return if PARTITIONABLE_MODELS.include?(klass.name)
+
+        raise Partitionable::Testing::InclusionError,
+          "#{klass} must be included in PARTITIONABLE_MODELS"
+
+      rescue InclusionError => e
+        Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
+      end
+    end
+
     included do
+      Partitionable::Testing.check_inclusion(self)
+
       before_validation :set_partition_id, on: :create
       validates :partition_id, presence: true
 
@@ -37,6 +62,8 @@ module Ci
       def partitionable(scope:)
         define_method(:partition_scope_value) do
           strong_memoize(:partition_scope_value) do
+            next Ci::Pipeline.current_partition_value if respond_to?(:importing?) && importing?
+
             record = scope.to_proc.call(self)
             record.respond_to?(:partition_id) ? record.partition_id : record
           end