Allow enabling for only certain STI models

README.md

@@ -4,12 +4,12 @@
 Given the following class definitions:
 
 ```ruby
-class Address
+class Address < ActiveRecord::Base
   belongs_to :addressable, :polymorphic => true
 end
 
-class Person
-  has_many :addresses, :as => addressable
+class Person < ActiveRecord::Base
+  has_many :addresses, :as => :addressable
 end
 
 class Vendor < Person
@@ -33,14 +33,14 @@ will output:
 #<Address id: 1, addressable_id: 1, addressable_type: 'Person' ...>
 ```
 
-Notice that addressable_type column is Person even though the actual class is Vendor.
+Notice that `addressable_type` column is `Person` even though the actual class is `Vendor`.
 
 Normally, this isn't a problem, however, it can have negative performance characteristics in certain circumstances. The most obvious one is that
-a join with persons or an extra query is required to find out the actual type of addressable.
+a join with persons or an extra query is required to find out the actual type of `addressable`.
 
-This gem adds the ActiveRecord::Base.store_base_sti_class configuration option. It defaults to true for backwards compatibility. Setting it to false will alter ActiveRecord's behavior to store the actual class in polymorphic _type columns when STI is used.
+This gem adds the ActiveRecord::Base.store_base_sti_class configuration option. It defaults to true for backwards compatibility. Setting it to false will alter ActiveRecord's behavior to store the actual class in polymorphic `_type` columns when STI is used.
 
-In the example above, if the ActiveRecord::Base.store_base_sti_class is false, the output will be,
+In the example above, if the `ActiveRecord::Base.store_base_sti_class is false`, the output will be,
 ```
   #<Vendor id: 1, type: "Vendor" ...>
   #<Address id: 1, addressable_id: 1, addressable_type: 'Vendor' ...>
@@ -54,18 +54,31 @@ Add the following line to your Gemfile,
 gem 'store_base_sti_class'
 ```
 
-then bundle install. Once you have the gem installed, add the following to one of the initializers (or make a new one) in config/initializers,
+then bundle install. Once you have the gem installed, add the following to one of the initializers (or make a new one) in `config/initializers`,
 
-  ActiveRecord::Base.store_base_sti_class = false
+```ruby
+ActiveRecord::Base.store_base_sti_class = false
+```
+
+When changing this behavior, you will have write a migration to update all of your existing `_type` columns accordingly. You may also need to change your application if it explicitly relies on the `_type` columns.
+
+If you only want to store the actual STI subclass type for certain classes and let all others use the
+default behavior of storing the STI base class, just set `store_sti_classes_for` to which classes should
+have the actual STI class stored.
 
-When changing this behavior, you will have write a migration to update all of your existing _type columns accordingly. You may also need to change your application if it explicitly relies on the _type columns.
+So in the example above, if you *only* wanted it to store the STI type for Person and its subclasses
+(Vendor), you would set:
+
+```ruby
+ActiveRecord::Base.store_sti_classes_for = ['Person']
+```
 
 ## Notes
 
 This gem incorporates work from:
 - https://github.com/codepodu/store_base_sti_class_for_4_0
 
-It currently works with ActiveRecord 4.0.x through 5.0.x. If you need support for ActiveRecord 3.x, use a pre-1.0 version of the gem.
+It currently works with ActiveRecord 4.0.x through 5.1.x. If you need support for ActiveRecord 3.x, use a pre-1.0 version of the gem.
 
 ## Copyright
 

lib/store_base_sti_class_for_5_1.rb

@@ -4,8 +4,38 @@
   module ActiveRecord
 
     class Base
-      class_attribute :store_base_sti_class
-      self.store_base_sti_class = true
+      class_attribute :_store_sti_classes_for, instance_accessor: false
+      def self.store_sti_classes_for
+        _store_sti_classes_for
+      end
+      # Override the setter so we can validate the input
+      def self.store_sti_classes_for=(new)
+        raise ArgumentError, "store_sti_classes_for must be set to an array or :all but was #{new.inspect}" unless new == :all or new.is_a? Array
+        new.map(&:constantize).each { |klass|
+          if klass != klass.base_class
+            raise ArgumentError, " You tried to set store_sti_classes_for to #{klass}, but store_sti_classes_for should only be set to an array of *base* classes for which you want to store the STI class (itself or any of its STI subclasses) in any _type columns. Did you mean '#{klass.base_class}'?"
+          end
+        } if new.is_a? Array
+        self._store_sti_classes_for = new
+      end
+      self.store_sti_classes_for = []
+
+      def self.store_sti_class?(klass)
+        return true if store_sti_classes_for == :all
+        klass = klass.is_a?(Class) ? klass : klass.constantize
+        store_sti_classes_for.include? klass.base_class.name
+      end
+
+      # For backwards compatibility
+      def self.store_base_sti_class=(new)
+        if new == true
+          self.store_sti_classes_for = []
+        elsif new == false
+          self.store_sti_classes_for = :all
+        else
+          raise ArgumentError, "store_base_sti_class can only be set to true or false but tried setting to #{new}"
+        end
+      end
     end
 
     module Associations
@@ -23,7 +53,7 @@ def creation_attributes
               # original:
               # attributes[reflection.type] = owner.class.base_class.name
 
-              attributes[reflection.type] = ActiveRecord::Base.store_base_sti_class ? owner.class.base_class.name : owner.class.name
+              attributes[reflection.type] = ActiveRecord::Base.store_sti_class?(owner.class) ? owner.class.name : owner.class.base_class.name
               # END PATCH
             end
           end
@@ -70,7 +100,7 @@ def join_constraints(foreign_table, foreign_klass, join_type, tables, chain)
                 # START PATCH
                 # original:
                 # value = foreign_klass.base_class.name
-                value = ActiveRecord::Base.store_base_sti_class ? foreign_klass.base_class.name : foreign_klass.name
+                value = ActiveRecord::Base.store_sti_class?(foreign_klass) ? foreign_klass.name : foreign_klass.base_class.name
                 # END PATCH
                 column = klass.columns_hash[reflection.type.to_s]
 
@@ -99,7 +129,7 @@ def replace_keys(record)
           # original:
           # owner[reflection.foreign_type] = record.class.base_class.name
 
-          owner[reflection.foreign_type] = ActiveRecord::Base.store_base_sti_class ? record.class.base_class.name : record.class.name
+          owner[reflection.foreign_type] = ActiveRecord::Base.store_sti_class?(record.class) ? record.class.name : record.class.base_class.name
 
           # END PATCH
         end
@@ -145,7 +175,7 @@ def build_scope
               # original:
               # scope.where!(klass.table_name => { reflection.type => model.base_class.sti_name })
 
-              scope.where!(klass.table_name => { reflection.type => ActiveRecord::Base.store_base_sti_class ? model.base_class.sti_name : model.sti_name })
+              scope.where!(klass.table_name => { reflection.type => ActiveRecord::Base.store_sti_class?(model) ? model.sti_name : model.base_class.sti_name })
 
               # END PATCH
             end
@@ -166,10 +196,10 @@ def through_scope
               # original: scope.where! reflection.foreign_type => options[:source_type]
 
               adjusted_foreign_type =
-                if ActiveRecord::Base.store_base_sti_class
-                  options[:source_type]
-                else
+                if ActiveRecord::Base.store_sti_class?(options[:source_type])
                   ([options[:source_type].constantize] + options[:source_type].constantize.descendants).map(&:to_s)
+                else
+                  options[:source_type]
                 end
 
               scope.where! reflection.foreign_type => adjusted_foreign_type
@@ -202,15 +232,15 @@ def self.get_bind_values(owner, chain)
           if last_reflection.type
             # START PATCH
             # original: binds << owner.class.base_class.name
-            binds << (ActiveRecord::Base.store_base_sti_class ? owner.class.base_class.name : owner.class.name)
+            binds << (ActiveRecord::Base.store_sti_class?(owner.class) ? owner.class.name : owner.class.base_class.name)
             # END PATCH
           end
 
           chain.each_cons(2).each do |reflection, next_reflection|
             if reflection.type
               # START PATCH
               # original: binds << next_reflection.klass.base_class.name
-              binds << (ActiveRecord::Base.store_base_sti_class ? next_reflection.klass.base_class.name : next_reflection.klass.name)
+              binds << (ActiveRecord::Base.store_sti_class?(next_reflection.klass) ? next_reflection.klass.name : next_reflection.klass.base_class.name)
               # END PATCH
             end
           end
@@ -231,11 +261,11 @@ def next_chain_scope(scope, table, reflection, foreign_table, next_reflection)
             # original:
             # value = transform_value(next_reflection.klass.base_class.name)
             # scope = scope.where(table.name => { reflection.type => value })
-            if ActiveRecord::Base.store_base_sti_class
-              value = transform_value(next_reflection.klass.base_class.name)
-            else
-              klass = next_reflection.klass
+            klass = next_reflection.klass
+            if ActiveRecord::Base.store_sti_class?(klass)
               value = ([klass] + klass.descendants).map(&:name)
+            else
+              value = transform_value(klass.base_class.name)
             end
             scope = scope.where(table.name => { reflection.type => value })
             # END PATCH
@@ -255,7 +285,7 @@ def last_chain_scope(scope, table, reflection, owner)
           if reflection.type
             # BEGIN PATCH
             # polymorphic_type = transform_value(owner.class.base_class.name)
-            polymorphic_type = transform_value(ActiveRecord::Base.store_base_sti_class ? owner.class.base_class.name : owner.class.name)
+            polymorphic_type = transform_value(ActiveRecord::Base.store_sti_class?(owner.class) ? owner.class.name : owner.class.base_class.name)
             # END PATCH
             scope = scope.where(table.name => { reflection.type => polymorphic_type })
           end
@@ -290,7 +320,7 @@ def construct_join_attributes(*records)
             #  records.map { |record| record.class.base_class.name }
 
             join_attributes[source_reflection.foreign_type] =
-              records.map { |record| ActiveRecord::Base.store_base_sti_class ? record.class.base_class.name : record.class.name }
+              records.map { |record| ActiveRecord::Base.store_sti_class?(record.class) ? record.class.name : record.class.base_class.name }
 
             # END PATCH
           end
@@ -314,10 +344,12 @@ def build_through_record(record)
             through_record.send("#{source_reflection.name}=", record)
 
             # START PATCH
-            if ActiveRecord::Base.store_base_sti_class
-              if options[:source_type]
-                through_record.send("#{source_reflection.foreign_type}=", options[:source_type])
-              end
+            # original:
+            # if options[:source_type]
+            #   through_record.send("#{source_reflection.foreign_type}=", options[:source_type])
+            # end
+            if options[:source_type] && !ActiveRecord::Base.store_sti_class?(options[:source_type])
+              through_record.send("#{source_reflection.foreign_type}=", options[:source_type])
             end
             # END PATCH
 
@@ -335,10 +367,10 @@ def source_type_info
 
           # START PATCH
           adjusted_source_type =
-            if ActiveRecord::Base.store_base_sti_class
-              source_type
-            else
+            if ActiveRecord::Base.store_sti_class?(source_type)
               ([source_type.constantize] + source_type.constantize.descendants).map(&:to_s)
+            else
+              source_type
             end
           # END PATCH
 

store_base_sti_class.gemspec

@@ -20,7 +20,7 @@ Gem::Specification.new do |s|
   s.homepage = 'http://github.com/appfolio/store_base_sti_class'
   s.licenses = ['MIT']
   s.rubygems_version = '2.2.2'
-  s.summary = 'Modifies ActiveRecord 4.0.x - 5.0.x with the ability to store the actual class (instead of the base class) in polymorhic _type columns when using STI'
+  s.summary = 'Modifies ActiveRecord 4.0.x - 5.1.x with the ability to store the actual class (instead of the base class) in polymorhic _type columns when using STI'
 
   s.add_runtime_dependency(%q<activerecord>, ['>= 4.0'])
   s.add_development_dependency(%q<minitest>, ['>= 4.0'])

test/models.rb

@@ -48,3 +48,14 @@ class Tag < ActiveRecord::Base
 
 class SpecialTag < Tag
 end
+
+class Address < ActiveRecord::Base
+  belongs_to :addressable, :polymorphic => true
+end
+
+class Person < ActiveRecord::Base
+  has_many :addresses, :as => :addressable
+end
+
+class Vendor < Person
+end

test/schema.rb

@@ -21,7 +21,7 @@
 
     t.integer :polytag_id
     t.string  :polytag_type
-    
+
     t.string  :taggable_type
     t.integer :taggable_id
   end
@@ -32,4 +32,15 @@
     t.integer :taggings_count, :default => 0
   end
 
+  create_table :addresses, :force => true do |t|
+    t.string  :city
+
+    t.integer :addressable_id
+    t.string  :addressable_type
+  end
+
+  create_table :people, :force => true do |t|
+    t.string  :type
+  end
+
 end

test/test_configuration_options.rb

@@ -0,0 +1,41 @@
+require 'helper'
+
+class TestClassVariables < StoreBaseSTIClass::TestCase
+
+  def setup
+    @old_store_sti_classes_for               = ActiveRecord::Base.store_sti_classes_for
+  end
+
+  def teardown
+    ActiveRecord::Base.store_sti_classes_for = @old_store_sti_classes_for
+  end
+
+  def test_setting_store_base_sti_class
+    ActiveRecord::Base.store_base_sti_class = false
+    assert_equal :all, ActiveRecord::Base.store_sti_classes_for
+
+    ActiveRecord::Base.store_base_sti_class = true
+    assert_equal [], ActiveRecord::Base.store_sti_classes_for
+  end
+
+  def test_setting_store_sti_classes_for
+    assert_nothing_raised do
+      ActiveRecord::Base.store_sti_classes_for = ['Post']
+    end
+    assert_equal ['Post'], ActiveRecord::Base.store_sti_classes_for
+
+    assert_raises(ArgumentError) do
+      ActiveRecord::Base.store_sti_classes_for = ['SpecialPost']
+    end
+
+    assert_nothing_raised do
+      ActiveRecord::Base.store_sti_classes_for = []
+    end
+    assert_equal [], ActiveRecord::Base.store_sti_classes_for
+
+    assert_raises(ArgumentError) do
+      ActiveRecord::Base.store_sti_classes_for = :none
+    end
+  end
+
+end