From 799a2090ce51c0f504594c911cd46c252db1d9a0 Mon Sep 17 00:00:00 2001 From: Jeremy Evans Date: Tue, 25 May 2010 08:47:26 -0700 Subject: [PATCH] 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 | 2 + .../plugins/single_table_inheritance.rb | 119 +++++++++++++++--- .../single_table_inheritance_spec.rb | 66 ++++++++++ 3 files changed, 167 insertions(+), 20 deletions(-) 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