Allow using a custom column value<->class mapping to the single_table…

…_inheritance plugin

This allows you to configure the mapping used.  The default mapping
remains the same, where the column value holds the name of the class
as a string.  However, you can now provide a mapping that allows you
to use integers, non-class name strings, or pretty much anything, to
map the column values for strings.

The most common customization is probably to use integers, which can
be easily done with a hash:

  Employee.plugin :single_table_inheritance, :type,
    :model_map=>{1=>:Staff, 2=>:Manager}

You can also use a pair of custom procs:

  Employee.plugin :single_table_inheritance, :type,
    :model_map=>proc{|v| v.reverse},
    :key_map=>proc{|klass| klass.name.reverse}

Check the new module RDoc for details.

CHANGELOG

@@ -1,5 +1,7 @@
 === HEAD
 
+* Allow using a custom column value<->class mapping to the single_table_inheritance plugin (jeremyevans, tmm1)
+
 * Handle SQL::Identifiers in the schema_dumper extension (jeremyevans) (#304)
 
 * Make sure certain alter table operations clear the schema correctly on MySQL (jeremyevans) (#301)

lib/sequel/plugins/single_table_inheritance.rb

@@ -1,27 +1,80 @@
 module Sequel
   module Plugins
-    # Sequel's built in Single Table Inheritance plugin makes subclasses
-    # of this model only load rows where the given key field matches the
-    # subclass's name.  If the key given has a NULL value or there are
-    # any problems looking up the class, uses the current class.
+    # The single_table_inheritance plugin allows storing all objects
+    # in the same class hierarchy in the same table.  It makes it so
+    # subclasses of this model only load rows related to the subclass,
+    # and when you retrieve rows from the main class, you get instances
+    # of the subclasses (if the rows should use the subclasses's class).
+    #
+    # By default, the plugin assumes that the +sti_key+ column (the first
+    # argument to the plugin) holds the class name as a string.  However,
+    # you can override this by using the <tt>:model_map</tt> option and/or
+    # the <tt>:key_map</tt> option.
     #   
-    # You should only use this in the parent class, not in the subclasses.
+    # You should only load this plugin in the parent class, not in the subclasses.
     #   
     # You shouldn't call set_dataset in the model after applying this
-    # plugin, otherwise subclasses might use the wrong dataset.
+    # plugin, otherwise subclasses might use the wrong dataset. You should
+    # make sure this plugin is loaded before the subclasses. Note that since you
+    # need to load the plugin before the subclasses are created, you can't use
+    # direct class references in the plugin class.  You should specify subclasses
+    # in the plugin call using class name strings or symbols, see usage below.
     #   
     # The filters and row_proc that sti_key sets up in subclasses may not work correctly if
     # those subclasses have further subclasses.  For those middle subclasses,
     # you may need to call set_dataset manually with the correct filter and
     # row_proc.
+    # 
+    # Usage:
+    #
+    #   # Use the default of storing the class name in the sti_key
+    #   # column (:kind in this case)
+    #   Employee.plugin :single_table_inheritance, :kind
+    #
+    #   # Using integers to store the class type, with a :model_map hash
+    #   # and an sti_key of :type
+    #   Employee.plugin :single_table_inheritance, :type,
+    #     :model_map=>{1=>:Staff, 2=>:Manager}
+    #
+    #   # Using non-class name strings
+    #   Employee.plugin :single_table_inheritance, :type,
+    #     :model_map=>{'line staff'=>:Staff, 'supervisor'=>:Manager}
+    #
+    #   # Using custom procs, with :model_map taking column values
+    #   # and yielding either a class, string, symbol, or nil, 
+    #   # and :key_map taking a class object and returning the column
+    #   # value to use
+    #   Employee.plugin :single_table_inheritance, :type,
+    #     :model_map=>proc{|v| v.reverse},
+    #     :key_map=>proc{|klass| klass.name.reverse}
+    #
+    # One minor issue to note is that if you specify the <tt>:key_map</tt>
+    # option as a hash, instead of having it inferred from the <tt>:model_map</tt>,
+    # you should only use class name strings as keys, you should not use symbols
+    # as keys.
     module SingleTableInheritance
-      # Set the sti_key and sti_dataset for the model, and change the
-      # dataset's row_proc so that the dataset yields objects of varying classes,
-      # where the class used has the same name as the key field.
-      def self.configure(model, key)
+      # Setup the necessary STI variables, see the module RDoc for SingleTableInheritance
+      def self.configure(model, key, opts={})
         model.instance_eval do
           @sti_key = key 
           @sti_dataset = dataset
+          @sti_model_map = opts[:model_map] || lambda{|v| v if v && v != ''}
+          @sti_key_map = if km = opts[:key_map]
+            if km.is_a?(Hash)
+              h = Hash.new{|h,k| h[k.to_s] unless k.is_a?(String)}
+              h.merge!(km)
+            else
+              km
+            end
+          elsif sti_model_map.is_a?(Hash)
+            h = Hash.new{|h,k| h[k.to_s] unless k.is_a?(String)}
+            sti_model_map.each do |k,v|
+              h[v.to_s] = k
+            end
+            h
+          else
+            lambda{|klass| klass.name.to_s}
+          end
           dataset.row_proc = lambda{|r| model.sti_load(r)}
         end
       end
@@ -34,38 +87,64 @@ module ClassMethods
         # The column name holding the STI key for this model
         attr_reader :sti_key
 
-        # Copy the sti_key and sti_dataset to the subclasses, and filter the
-        # subclass's dataset so it is restricted to rows where the key column
-        # matches the subclass's name.
+        # A hash/proc with class keys and column value values, mapping
+        # the the class to a particular value given to the sti_key column.
+        # Used to set the column value when creating objects, and for the
+        # filter when retrieving objects in subclasses.
+        attr_reader :sti_key_map
+
+        # A hash/proc with column value keys and class values, mapping
+        # the value of the sti_key column to the appropriate class to use.
+        attr_reader :sti_model_map
+
+        # Copy the necessary attributes to the subclasses, and filter the
+        # subclass's dataset based on the sti_kep_map entry for the class.
         def inherited(subclass)
           super
           sk = sti_key
           sd = sti_dataset
-          subclass.set_dataset(sd.filter(SQL::QualifiedIdentifier.new(table_name, sk)=>subclass.name.to_s), :inherited=>true)
+          skm = sti_key_map
+          smm = sti_model_map
+          subclass.set_dataset(sd.filter(SQL::QualifiedIdentifier.new(table_name, sk)=>skm[subclass]), :inherited=>true)
           subclass.instance_eval do
             @sti_key = sk
             @sti_dataset = sd
+            @sti_key_map = skm
+            @sti_model_map = smm
             @simple_table = nil
           end
         end
 
         # Return an instance of the class specified by sti_key,
         # used by the row_proc.
         def sti_load(r)
-          v = r[sti_key]
-          model = if (v && v != '')
+          sti_class(sti_model_map[r[sti_key]]).load(r)
+        end
+
+        private
+
+        # Return a class object.  If a class is given, return it directly.
+        # Treat strings and symbols as class names.  If nil is given or
+        # an invalid class name string or symbol is used, return self.
+        # Raise an error for other types.
+        def sti_class(v)
+          case v
+          when String, Symbol
             constantize(v) rescue self
-          else
+          when nil
             self
+          when Class
+            v
+          else
+            raise(Error, "Invalid class type used: #{v.inspect}")
           end
-          model.load(r)
         end
       end
 
       module InstanceMethods
-        # Set the sti_key column to the name of the model.
+        # Set the sti_key column based on the sti_key_map.
         def before_create
-          send("#{model.sti_key}=", model.name.to_s) unless send(model.sti_key)
+          send("#{model.sti_key}=", model.sti_key_map[model]) unless send(model.sti_key)
           super
         end
       end

spec/extensions/single_table_inheritance_spec.rb

@@ -93,4 +93,70 @@ def @ds.fetch_rows(sql)
     StiTestSub1.dataset.sql.should == "SELECT * FROM sti_tests WHERE (sti_tests.kind = 'StiTestSub1')"
     StiTestSub2.dataset.sql.should == "SELECT * FROM sti_tests WHERE (sti_tests.kind = 'StiTestSub2')"
   end
+
+  context "with custom options" do
+    before do
+      class ::StiTest2 < Sequel::Model
+        columns :id, :kind
+        def _refresh(x); end
+      end
+    end
+    after do
+      Object.send(:remove_const, :StiTest2)
+      Object.send(:remove_const, :StiTest3)
+      Object.send(:remove_const, :StiTest4)
+    end
+
+    it "should work with custom procs with strings" do
+      StiTest2.plugin :single_table_inheritance, :kind, :model_map=>proc{|v| v == 1 ? 'StiTest3' : 'StiTest4'}, :key_map=>proc{|klass| klass.name == 'StiTest3' ? 1 : 2}
+      class ::StiTest3 < ::StiTest2; end
+      class ::StiTest4 < ::StiTest2; end
+      StiTest2.dataset.row_proc.call(:kind=>0).should be_a_instance_of(StiTest4)
+      StiTest2.dataset.row_proc.call(:kind=>1).should be_a_instance_of(StiTest3)
+      StiTest2.dataset.row_proc.call(:kind=>2).should be_a_instance_of(StiTest4)
+
+      StiTest2.create.kind.should == 2
+      StiTest3.create.kind.should == 1
+      StiTest4.create.kind.should == 2
+    end
+
+    it "should work with custom procs with symbols" do
+      StiTest2.plugin :single_table_inheritance, :kind, :model_map=>proc{|v| v == 1 ? :StiTest3 : :StiTest4}, :key_map=>proc{|klass| klass.name == 'StiTest3' ? 1 : 2}
+      class ::StiTest3 < ::StiTest2; end
+      class ::StiTest4 < ::StiTest2; end
+      StiTest2.dataset.row_proc.call(:kind=>0).should be_a_instance_of(StiTest4)
+      StiTest2.dataset.row_proc.call(:kind=>1).should be_a_instance_of(StiTest3)
+      StiTest2.dataset.row_proc.call(:kind=>2).should be_a_instance_of(StiTest4)
+
+      StiTest2.create.kind.should == 2
+      StiTest3.create.kind.should == 1
+      StiTest4.create.kind.should == 2
+    end
+
+    it "should work with custom hashes" do
+      StiTest2.plugin :single_table_inheritance, :kind, :model_map=>{0=>StiTest2, 1=>:StiTest3, 2=>'StiTest4'}, :key_map=>{StiTest2=>4, 'StiTest3'=>5, 'StiTest4'=>6}
+      class ::StiTest3 < ::StiTest2; end
+      class ::StiTest4 < ::StiTest2; end
+      StiTest2.dataset.row_proc.call(:kind=>0).should be_a_instance_of(StiTest2)
+      StiTest2.dataset.row_proc.call(:kind=>1).should be_a_instance_of(StiTest3)
+      StiTest2.dataset.row_proc.call(:kind=>2).should be_a_instance_of(StiTest4)
+
+      StiTest2.create.kind.should == 4
+      StiTest3.create.kind.should == 5
+      StiTest4.create.kind.should == 6
+    end
+
+    it "should infer key_map from model_map if provided as a hash" do
+      StiTest2.plugin :single_table_inheritance, :kind, :model_map=>{0=>StiTest2, 1=>'StiTest3', 2=>:StiTest4}
+      class ::StiTest3 < ::StiTest2; end
+      class ::StiTest4 < ::StiTest2; end
+      StiTest2.dataset.row_proc.call(:kind=>0).should be_a_instance_of(StiTest2)
+      StiTest2.dataset.row_proc.call(:kind=>1).should be_a_instance_of(StiTest3)
+      StiTest2.dataset.row_proc.call(:kind=>2).should be_a_instance_of(StiTest4)
+
+      StiTest2.create.kind.should == 0
+      StiTest3.create.kind.should == 1
+      StiTest4.create.kind.should == 2
+    end
+  end
 end