1 files changed, 37 insertions, 7 deletions
diff --git a/app/models/concerns/ci/partitionable/switch.rb b/app/models/concerns/ci/partitionable/switch.rb
index c1bbd107e9f..6195f92114f 100644
--- a/app/models/concerns/ci/partitionable/switch.rb
+++ b/app/models/concerns/ci/partitionable/switch.rb
@@ -2,6 +2,8 @@
 
 module Ci
   module Partitionable
+    MUTEX = Mutex.new
+
     module Switch
       extend ActiveSupport::Concern
 
@@ -14,18 +16,39 @@ module Ci
                             predicate_builder cached_find_by_statement].freeze
 
       included do |base|
-        partitioned = Class.new(base) do
-          self.table_name = base.routing_table_name
+        install_partitioned_class(base)
+      end
+
+      class_methods do
+        # `Class.new(partitionable_model)` triggers `partitionable_model.inherited`
+        # and we need the mutex to break the recursion without adding extra accessors
+        # on the model. This will be used during code loading, not runtime.
+        #
+        def install_partitioned_class(partitionable_model)
+          Partitionable::MUTEX.synchronize do
+            partitioned = Class.new(partitionable_model) do
+              self.table_name = partitionable_model.routing_table_name
+
+              def self.routing_class?
+                true
+              end
+
+              def self.sti_name
+                superclass.sti_name
+              end
+            end
 
-          def self.routing_class?
-            true
+            partitionable_model.const_set(:Partitioned, partitioned)
           end
         end
 
-        base.const_set(:Partitioned, partitioned)
-      end
+        def inherited(child_class)
+          super
+          return if Partitionable::MUTEX.owned?
+
+          install_partitioned_class(child_class)
+        end
 
-      class_methods do
         def routing_class?
           false
         end
@@ -51,6 +74,13 @@ module Ci
             end
           end
         end
+
+        def type_condition(table = arel_table)
+          sti_column = table[inheritance_column]
+          sti_names  = ([self] + descendants).map(&:sti_name).uniq
+
+          predicate_builder.build(sti_column, sti_names)
+        end
       end
     end
   end