lib/table_helper/header.rb in table_helper-0.2.0 vs lib/table_helper/header.rb in table_helper-0.2.1

@@ -31,10 +31,20 @@
       # If we know what class the objects in the collection are and we can
       # figure out what columns are defined in that class, then we can
       # pre-fill the header with those columns so that the user doesn't
       # have to
       klass = table.klass
-      column(*klass.column_names.map(&:to_sym)) if klass && klass.respond_to?(:column_names)
+      if klass && klass.respond_to?(:column_names)
+        if !table.empty? && klass < ActiveRecord::Base
+          # Make sure only the attributes that have been loaded are used
+          column_names = table.collection.first.attributes.keys
+        else
+          # Assume all attributes are loaded
+          column_names = klass.column_names
+        end
+        
+        column(*column_names.map(&:to_sym))
+      end
       
       @customized = false
     end
     
     # The current columns in this header, in the order in which they will be