1 files changed, 24 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?