diff --git a/CHANGELOG b/CHANGELOG
index b08bbddbbf..d6188fcb0c 100644
--- a/CHANGELOG
+++ b/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)
diff --git a/lib/sequel/plugins/single_table_inheritance.rb b/lib/sequel/plugins/single_table_inheritance.rb
index 09c310876e..6c180ba996 100644
--- a/lib/sequel/plugins/single_table_inheritance.rb
+++ b/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 :model_map option and/or
+ # the :key_map 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 :key_map
+ # option as a hash, instead of having it inferred from the :model_map,
+ # 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,17 +87,30 @@ 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
@@ -52,20 +118,33 @@ def inherited(subclass)
# 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
diff --git a/spec/extensions/single_table_inheritance_spec.rb b/spec/extensions/single_table_inheritance_spec.rb
index e55247fd69..05c5ddb2ef 100644
--- a/spec/extensions/single_table_inheritance_spec.rb
+++ b/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