Rails 3 + STI: Making Associations Work Properly

I really like the idea of Single Table Inheritance (STI) for all sorts of applications to keep code DRY and make it easier to organize object behavior. The only problem is that Rails 3.0.3 doesn’t fully support STI with association collections.

Let’s say you have a User model that has many badges. The badges will be stored in the badges table but you want to implement each badge in a subclass. All you have to do is make sure there’s a :type field of type string in your badges table and Rails STI support should take care of the rest (well, in theory).

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class User < ActiveRecord::Base
  has_many :badges
end

class Badge < ActiveRecord::Base
  belongs_to :user
  def award
    raise "Must implement in subclass"
  end
end

class Badges::Superhero < Badge
  def award
    user.status = 'superhero'
  end
end

Now you can do cool things like create a new Superhero badge and add it to a user’s badge collection.

1
2
3
user = User.first
badge = Badges::Superhero.new
user.badges << badge

But for some weird reason, you can’t use the best practice of building a badge directly from the user’s badges collection.

1
2
3
user = User.first
badge = user.badges.build(:type => Badges::Superhero)
# badge.class == Badge

This is particularly annoying if you’re trying to create new badges from a form where :type is a drop down menu.

The reason the collection build method doesn’t work as expected is because :type is a protected field and ActiveRecord::AssociationReflection doesn’t fully support STI (at least in Rails 3.0.3).

Not to fret, hacks to the rescue!

You have two options to make STI work as expected.

Option 1: Override the Badge.new method to handle :type

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class Badge < ActiveRecord::Base
  belongs_to :user
  self.abstract_class = true

  class << self
    def new_with_cast(*a, &b)
      if (h = a.first).is_a? Hash and (type = h.symbolize_keys[self.class.inheritance_column.to_sym]) and (klass = type.to_s.constantize) != self
        raise "Must be a subclass of Badge" unless klass < self  # klass should be a descendant of self
        return klass.new_without_cast(*a, &b)
      end
      raise "Badge must be created through a subclass."
      new_without_cast(*a, &b)
    end
    alias_method_chain :new, :cast
  end
end

Option 2: Patch AssociationReflection to behave more intelligently

1
2
3
4
5
6
7
8
9
10
11
12
class ActiveRecord::Reflection::AssociationReflection
  def build_association(*opts)
    col = klass.inheritance_column.to_sym
    if (h = opts.first).is_a? Hash and (type = h.symbolize_keys[col]) and type.class == Class
      opts.first[col].to_s.constantize.new(*opts)
    elsif klass.abstract_class?
      raise "#{klass.to_s} is an abstract class and can not be directly instantiated"
    else
      klass.new(*opts)
    end
  end
end

My preference is Option 2 even though it might break in future releases of Rails. I’d rather have Rails behaving as expected than pepper my models code with repetitive hacks.

The above solutions were inspired from a couple of different posts and sources.

I submitted Option 2 as a patch for Rails.