lib/railroady/models_diagram.rb in railroady-1.0.2 vs lib/railroady/models_diagram.rb in railroady-1.0.3

@@ -52,34 +52,36 @@
       if @options.brief || current_class.abstract_class?
         node_type = 'model-brief'
       else 
         node_type = 'model'
 
-        # Collect model's content columns
-
-        content_columns = current_class.content_columns
+        # Collect model's content columns or all columns if all_columns flag is passed
+        if @options.all_columns
+          columns = current_class.columns
+        else
+          columns = current_class.content_columns
+        end
 	
         if @options.hide_magic 
           # From patch #13351
           # http://wiki.rubyonrails.org/rails/pages/MagicFieldNames
           magic_fields = [
           "created_at", "created_on", "updated_at", "updated_on",
           "lock_version", "type", "id", "position", "parent_id", "lft", 
           "rgt", "quote", "template"
           ]
           magic_fields << current_class.table_name + "_count" if current_class.respond_to? 'table_name' 
-          content_columns = current_class.content_columns.select {|c| ! magic_fields.include? c.name}
-        else
-          content_columns = current_class.content_columns
+          columns = current_class.content_columns.select {|c| ! magic_fields.include? c.name}
         end
-        
-        content_columns.each do |a|
-          content_column = a.name
-          content_column += ' :' + a.type.to_s unless @options.hide_types
-          node_attribs << content_column
+
+        columns.each do |a|
+          column = a.name
+          column += ' :' + a.type.to_s unless @options.hide_types
+          node_attribs << column
         end
       end
+
       @graph.add_node [node_type, current_class.name, node_attribs]
       generated = true
       # Process class associations
       associations = current_class.reflect_on_all_associations
       if @options.inheritance && ! @options.transitive
@@ -120,10 +122,10 @@
 
     # Only non standard association names needs a label
     
     # from patch #12384
     # if assoc.class_name == assoc.name.to_s.singularize.camelize
-    assoc_class_name = (assoc.class_name.respond_to? 'underscore') ? assoc.class_name.underscore.singularize.camelize : assoc.class_name 
+    assoc_class_name = (assoc.class_name.respond_to? 'underscore') ? assoc.class_name.underscore.camelize : assoc.class_name
     if assoc_class_name == assoc.name.to_s.singularize.camelize
       assoc_name = ''
     else
       assoc_name = assoc.name.to_s
     end