diff --git a/README.md b/README.md
index bac488d5..32ea9578 100644
--- a/README.md
+++ b/README.md
@@ -251,6 +251,7 @@ you can do so with a simple environment variable, instead of editing the
--ignore-unknown-models don't display warnings for bad model files
--with-comment include database comments in model annotations
--with-comment-column include database comments in model annotations, as its own column, after all others
+ --group-sti-columns group annotation columns by STI class ownership
### Option: `additional_file_patterns`
diff --git a/lib/annotate/annotate_models.rb b/lib/annotate/annotate_models.rb
index dc2901a3..2e5f85ec 100644
--- a/lib/annotate/annotate_models.rb
+++ b/lib/annotate/annotate_models.rb
@@ -3,6 +3,7 @@
require 'bigdecimal'
require 'annotate/constants'
+require 'annotate/sti_columns'
require_relative 'annotate_models/file_patterns'
module AnnotateModels
@@ -171,27 +172,43 @@ def get_schema_info(klass, header, options = {}) # rubocop:disable Metrics/Metho
# Output annotation
bare_max_attrs_length = cols_meta.map { |_, m| m[:simple_formatted_attrs].length }.max
- cols.each do |col|
- col_type = cols_meta[col.name][:col_type]
- attrs = cols_meta[col.name][:attrs]
- col_name = cols_meta[col.name][:col_name]
- simple_formatted_attrs = cols_meta[col.name][:simple_formatted_attrs]
- col_comment = cols_meta[col.name][:col_comment]
-
- if options[:format_rdoc]
- info << sprintf("# %-#{max_size}.#{max_size}s%s", "*#{col_name}*::", attrs.unshift(col_type).join(", ")).rstrip + "\n"
- elsif options[:format_yard]
- info << sprintf("# @!attribute #{col_name}") + "\n"
- ruby_class = col.respond_to?(:array) && col.array ? "Array<#{map_col_type_to_ruby_classes(col_type)}>": map_col_type_to_ruby_classes(col_type)
- info << sprintf("# @return [#{ruby_class}]") + "\n"
- elsif options[:format_markdown]
- name_remainder = max_size - col_name.length - non_ascii_length(col_name)
- type_remainder = (md_type_allowance - 2) - col_type.length
- info << (sprintf("# **`%s`**%#{name_remainder}s | `%s`%#{type_remainder}s | `%s`", col_name, " ", col_type, " ", attrs.join(", ").rstrip)).gsub('``', ' ').rstrip + "\n"
- elsif with_comments_column
- info << format_default(col_name, max_size, col_type, bare_type_allowance, simple_formatted_attrs, bare_max_attrs_length, col_comment)
- else
- info << format_default(col_name, max_size, col_type, bare_type_allowance, simple_formatted_attrs)
+ if options[:group_sti_columns]
+ grouped_cols = Annotate::StiColumns.partition(klass, cols)
+ else
+ grouped_cols = [[nil, cols]]
+ end
+
+ grouped_cols.each do |group_label, group_columns|
+ if group_label
+ if options[:format_markdown]
+ info << "#\n# ### #{group_label}\n#\n"
+ else
+ info << "#\n# -- #{group_label} --\n#\n"
+ end
+ end
+
+ group_columns.each do |col|
+ col_type = cols_meta[col.name][:col_type]
+ attrs = cols_meta[col.name][:attrs]
+ col_name = cols_meta[col.name][:col_name]
+ simple_formatted_attrs = cols_meta[col.name][:simple_formatted_attrs]
+ col_comment = cols_meta[col.name][:col_comment]
+
+ if options[:format_rdoc]
+ info << sprintf("# %-#{max_size}.#{max_size}s%s", "*#{col_name}*::", attrs.unshift(col_type).join(", ")).rstrip + "\n"
+ elsif options[:format_yard]
+ info << sprintf("# @!attribute #{col_name}") + "\n"
+ ruby_class = col.respond_to?(:array) && col.array ? "Array<#{map_col_type_to_ruby_classes(col_type)}>": map_col_type_to_ruby_classes(col_type)
+ info << sprintf("# @return [#{ruby_class}]") + "\n"
+ elsif options[:format_markdown]
+ name_remainder = max_size - col_name.length - non_ascii_length(col_name)
+ type_remainder = (md_type_allowance - 2) - col_type.length
+ info << (sprintf("# **`%s`**%#{name_remainder}s | `%s`%#{type_remainder}s | `%s`", col_name, " ", col_type, " ", attrs.join(", ").rstrip)).gsub('``', ' ').rstrip + "\n"
+ elsif with_comments_column
+ info << format_default(col_name, max_size, col_type, bare_type_allowance, simple_formatted_attrs, bare_max_attrs_length, col_comment)
+ else
+ info << format_default(col_name, max_size, col_type, bare_type_allowance, simple_formatted_attrs)
+ end
end
end
@@ -733,7 +750,7 @@ def annotate_model_file(annotated, file, header, options)
klass = get_model_class(file)
do_annotate = klass.is_a?(Class) &&
klass < ActiveRecord::Base &&
- (!options[:exclude_sti_subclasses] || !(klass.superclass < ActiveRecord::Base && klass.table_name == klass.superclass.table_name)) &&
+ (!options[:exclude_sti_subclasses] || !Annotate::StiColumns.subclass?(klass)) &&
!klass.abstract_class? &&
klass.table_exists?
diff --git a/lib/annotate/constants.rb b/lib/annotate/constants.rb
index 0d322565..cfe64ffc 100644
--- a/lib/annotate/constants.rb
+++ b/lib/annotate/constants.rb
@@ -18,7 +18,7 @@ module Constants
:trace, :timestamp, :exclude_serializers, :classified_sort,
:show_foreign_keys, :show_complete_foreign_keys,
:exclude_scaffolds, :exclude_controllers, :exclude_helpers,
- :exclude_sti_subclasses, :ignore_unknown_models, :with_comment, :with_comment_column,
+ :exclude_sti_subclasses, :group_sti_columns, :ignore_unknown_models, :with_comment, :with_comment_column,
:show_check_constraints
].freeze
diff --git a/lib/annotate/parser.rb b/lib/annotate/parser.rb
index ad85caf5..8722c807 100644
--- a/lib/annotate/parser.rb
+++ b/lib/annotate/parser.rb
@@ -309,6 +309,11 @@ def add_options_to_parser(option_parser) # rubocop:disable Metrics/MethodLength,
"include database comments in model annotations, as its own column, after all others") do
env['with_comment_column'] = 'true'
end
+
+ option_parser.on('--group-sti-columns',
+ "group annotation columns by STI class ownership") do
+ env['group_sti_columns'] = 'true'
+ end
end
end
end
diff --git a/lib/annotate/sti_columns.rb b/lib/annotate/sti_columns.rb
new file mode 100644
index 00000000..d600087f
--- /dev/null
+++ b/lib/annotate/sti_columns.rb
@@ -0,0 +1,121 @@
+require 'set'
+
+module Annotate
+ module StiColumns
+ class << self
+ def subclass?(klass)
+ klass.superclass < ActiveRecord::Base &&
+ klass.table_name == klass.superclass.table_name
+ end
+
+ def base_class?(klass)
+ klass.column_names.include?(klass.inheritance_column) &&
+ sti_descendants(klass).any?
+ end
+
+ def columns_referenced_in(klass)
+ cols = Set.new
+ cols.merge(klass.validators.flat_map { |v| v.attributes.map(&:to_s) })
+ if klass.respond_to?(:reflect_on_all_associations)
+ cols.merge(
+ klass.reflect_on_all_associations(:belongs_to).flat_map { |a|
+ fk = [a.foreign_key.to_s]
+ fk << "#{a.name}_type" if a.respond_to?(:polymorphic?) && a.polymorphic?
+ fk
+ }
+ )
+ end
+ cols.merge(klass.defined_enums.keys) if klass.respond_to?(:defined_enums)
+ if klass.respond_to?(:stored_attributes) && klass.stored_attributes.any?
+ cols.merge(klass.stored_attributes.keys.map(&:to_s))
+ end
+ cols
+ end
+
+ def columns_owned_by(klass)
+ return Set.new unless subclass?(klass)
+
+ columns_referenced_in(klass) - columns_referenced_in(klass.superclass)
+ end
+
+ def partition(klass, cols)
+ if subclass?(klass)
+ partition_for_subclass(klass, cols)
+ elsif base_class?(klass)
+ partition_for_base_class(klass, cols)
+ else
+ [[nil, cols]]
+ end
+ end
+
+ private
+
+ # Returns all descendants of klass that share the same table (STI).
+ # Uses descendants (all levels) rather than subclasses (direct only)
+ # to support multi-level STI hierarchies.
+ #
+ # Attempts eager loading first to ensure all subclasses are visible.
+ def sti_descendants(klass)
+ ensure_models_loaded
+ if klass.respond_to?(:descendants)
+ klass.descendants.select { |d| d.table_name == klass.table_name }
+ else
+ klass.subclasses.select { |d| d.table_name == klass.table_name }
+ end
+ end
+
+ def ensure_models_loaded
+ return if @models_loaded
+
+ if defined?(Rails) && Rails.respond_to?(:application) && Rails.application
+ Rails.application.eager_load! unless Rails.application.config.eager_load
+ end
+ @models_loaded = true
+ end
+
+ def partition_for_subclass(klass, cols)
+ owned = columns_owned_by(klass)
+ base_name = klass.superclass.name.demodulize
+ shared = cols.reject { |c| owned.include?(c.name.to_s) }
+ specific = cols.select { |c| owned.include?(c.name.to_s) }
+
+ if specific.empty?
+ $stderr.puts "Warning: --group-sti-columns could not determine which columns belong to #{klass.name}."
+ $stderr.puts " Add validations, associations, or enums to #{klass.name} to improve grouping."
+ end
+
+ groups = []
+ groups << ["#{base_name} columns", shared] if shared.any?
+ groups << ["#{klass.name.demodulize} columns", specific] if specific.any?
+ groups
+ end
+
+ def partition_for_base_class(klass, cols)
+ all_descendants = sti_descendants(klass).sort_by(&:name)
+ ownership = {}
+
+ all_descendants.each do |sub|
+ columns_owned_by(sub).each do |col_name|
+ ownership[col_name] ||= sub.name.demodulize
+ end
+ end
+
+ base_cols = cols.reject { |c| ownership.key?(c.name.to_s) }
+ groups = [["#{klass.name.demodulize} columns", base_cols]]
+
+ all_descendants.each do |sub|
+ sub_name = sub.name.demodulize
+ sub_cols = cols.select { |c| ownership[c.name.to_s] == sub_name }
+ groups << ["#{sub_name} columns", sub_cols] if sub_cols.any?
+ end
+
+ if ownership.empty?
+ $stderr.puts "Warning: --group-sti-columns found STI subclasses for #{klass.name} but no columns could be assigned."
+ $stderr.puts " Add validations, associations, or enums to subclasses to improve grouping."
+ end
+
+ groups
+ end
+ end
+ end
+end
diff --git a/lib/generators/annotate/templates/auto_annotate_models.rake b/lib/generators/annotate/templates/auto_annotate_models.rake
index 61cdcd7a..1cc5d6d5 100644
--- a/lib/generators/annotate/templates/auto_annotate_models.rake
+++ b/lib/generators/annotate/templates/auto_annotate_models.rake
@@ -34,6 +34,7 @@ if Rails.env.development?
'exclude_controllers' => 'true',
'exclude_helpers' => 'true',
'exclude_sti_subclasses' => 'false',
+ 'group_sti_columns' => 'false',
'ignore_model_sub_dir' => 'false',
'ignore_columns' => nil,
'ignore_routes' => nil,
diff --git a/lib/tasks/annotate_models.rake b/lib/tasks/annotate_models.rake
index 776f97ba..99d6dc23 100644
--- a/lib/tasks/annotate_models.rake
+++ b/lib/tasks/annotate_models.rake
@@ -35,6 +35,7 @@ task annotate_models: :environment do
options[:exclude_controllers] = Annotate::Helpers.true?(ENV.fetch('exclude_controllers', 'true'))
options[:exclude_helpers] = Annotate::Helpers.true?(ENV.fetch('exclude_helpers', 'true'))
options[:exclude_sti_subclasses] = Annotate::Helpers.true?(ENV['exclude_sti_subclasses'])
+ options[:group_sti_columns] = Annotate::Helpers.true?(ENV['group_sti_columns'])
options[:ignore_model_sub_dir] = Annotate::Helpers.true?(ENV['ignore_model_sub_dir'])
options[:format_bare] = Annotate::Helpers.true?(ENV['format_bare'])
options[:format_rdoc] = Annotate::Helpers.true?(ENV['format_rdoc'])
diff --git a/spec/lib/annotate/annotate_models_spec.rb b/spec/lib/annotate/annotate_models_spec.rb
index 09647461..faa8cfba 100644
--- a/spec/lib/annotate/annotate_models_spec.rb
+++ b/spec/lib/annotate/annotate_models_spec.rb
@@ -1942,6 +1942,143 @@ def mock_column(name, type, options = {})
end
end
+ describe '.get_schema_info with group_sti_columns' do
+ let(:all_columns) do
+ [
+ mock_column(:id, :integer),
+ mock_column(:type, :string, limit: 255),
+ mock_column(:name, :string, limit: 100),
+ mock_column(:num_doors, :integer),
+ mock_column(:payload_capacity, :integer)
+ ]
+ end
+
+ def build_sti_base(columns, name: 'Vehicle')
+ klass = mock_class(:vehicles, :id, columns)
+ allow(klass).to receive(:superclass).and_return(ActiveRecord::Base)
+ allow(klass).to receive(:<).with(ActiveRecord::Base).and_return(true)
+ allow(klass).to receive(:name).and_return(name)
+ allow(klass).to receive(:inheritance_column).and_return('type')
+ allow(klass).to receive(:subclasses).and_return([])
+
+ # Introspection
+ allow(klass).to receive(:validators).and_return([
+ double('Validator', attributes: [:name])
+ ])
+ allow(klass).to receive(:reflect_on_all_associations).with(:belongs_to).and_return([])
+ allow(klass).to receive(:defined_enums).and_return({})
+ allow(klass).to receive(:stored_attributes).and_return({})
+
+ klass
+ end
+
+ def build_sti_subclass(parent, columns, name:, owned_validator_attrs: [])
+ klass = mock_class(:vehicles, :id, columns)
+ allow(klass).to receive(:superclass).and_return(parent)
+ allow(klass).to receive(:<).with(ActiveRecord::Base).and_return(true)
+ allow(klass).to receive(:name).and_return(name)
+ allow(klass).to receive(:inheritance_column).and_return('type')
+ allow(klass).to receive(:subclasses).and_return([])
+
+ # Make parent look like an AR subclass for sti_subclass? check
+ allow(parent).to receive(:<).with(ActiveRecord::Base).and_return(true)
+
+ # Introspection: parent validators + own
+ own_validators = owned_validator_attrs.map { |attr| double('Validator', attributes: [attr]) }
+ allow(klass).to receive(:validators).and_return(parent.validators + own_validators)
+ allow(klass).to receive(:reflect_on_all_associations).with(:belongs_to).and_return([])
+ allow(klass).to receive(:defined_enums).and_return({})
+ allow(klass).to receive(:stored_attributes).and_return({})
+
+ klass
+ end
+
+ context 'when annotating an STI subclass' do
+ it 'groups columns with section headers' do
+ vehicle = build_sti_base(all_columns)
+ car = build_sti_subclass(vehicle, all_columns, name: 'Car',
+ owned_validator_attrs: [:num_doors])
+ allow(vehicle).to receive(:subclasses).and_return([car])
+
+ result = AnnotateModels.get_schema_info(car, 'Schema Info', group_sti_columns: true)
+
+ expect(result).to include('# -- Vehicle columns --')
+ expect(result).to include('# -- Car columns --')
+ expect(result).to include('# num_doors')
+
+ # Verify num_doors appears after Car header
+ car_section = result.index('# -- Car columns --')
+ num_doors_pos = result.index('# num_doors')
+ expect(num_doors_pos).to be > car_section
+ end
+ end
+
+ context 'when annotating an STI base class' do
+ it 'groups columns by owning subclass' do
+ vehicle = build_sti_base(all_columns)
+ car = build_sti_subclass(vehicle, all_columns, name: 'Car',
+ owned_validator_attrs: [:num_doors])
+ truck = build_sti_subclass(vehicle, all_columns, name: 'Truck',
+ owned_validator_attrs: [:payload_capacity])
+ allow(vehicle).to receive(:subclasses).and_return([car, truck])
+
+ result = AnnotateModels.get_schema_info(vehicle, 'Schema Info', group_sti_columns: true)
+
+ expect(result).to include('# -- Vehicle columns --')
+ expect(result).to include('# -- Car columns --')
+ expect(result).to include('# -- Truck columns --')
+
+ # Base columns should be in Vehicle section
+ vehicle_section = result.index('# -- Vehicle columns --')
+ car_section = result.index('# -- Car columns --')
+ id_pos = result.index('# id')
+ expect(id_pos).to be > vehicle_section
+ expect(id_pos).to be < car_section
+ end
+ end
+
+ context 'when using markdown format' do
+ it 'uses markdown headers for section labels' do
+ vehicle = build_sti_base(all_columns)
+ car = build_sti_subclass(vehicle, all_columns, name: 'Car',
+ owned_validator_attrs: [:num_doors])
+ allow(vehicle).to receive(:subclasses).and_return([car])
+
+ result = AnnotateModels.get_schema_info(car, 'Schema Info',
+ group_sti_columns: true, format_markdown: true)
+
+ expect(result).to include('# ### Vehicle columns')
+ expect(result).to include('# ### Car columns')
+ end
+ end
+
+ context 'when group_sti_columns is false' do
+ it 'does not add section headers' do
+ result = AnnotateModels.get_schema_info(
+ mock_class(:users, :id, [mock_column(:id, :integer), mock_column(:name, :string, limit: 50)]),
+ 'Schema Info',
+ group_sti_columns: false
+ )
+
+ expect(result).not_to include('-- ')
+ end
+ end
+
+ context 'when model is not STI' do
+ it 'does not add section headers' do
+ klass = mock_class(:users, :id, [mock_column(:id, :integer), mock_column(:name, :string, limit: 50)])
+ allow(klass).to receive(:superclass).and_return(ActiveRecord::Base)
+ allow(klass).to receive(:column_names).and_return(%w[id name])
+ allow(klass).to receive(:inheritance_column).and_return('type')
+ allow(klass).to receive(:subclasses).and_return([])
+
+ result = AnnotateModels.get_schema_info(klass, 'Schema Info', group_sti_columns: true)
+
+ expect(result).not_to include('-- ')
+ end
+ end
+ end
+
describe '.set_defaults' do
subject do
Annotate::Helpers.true?(ENV['show_complete_foreign_keys'])
@@ -2584,6 +2721,46 @@ class Foo < ActiveRecord::Base
end
end
+ context 'when annotation contains STI group headers' do
+ let :filename do
+ 'sti_grouped.rb'
+ end
+
+ let :file_content do
+ <<~EOS
+ # == Schema Information
+ #
+ # Table name: vehicles
+ #
+ #
+ # -- Vehicle columns --
+ #
+ # id :bigint not null, primary key
+ # type :string
+ # name :string not null
+ #
+ # -- Car columns --
+ #
+ # num_doors :integer
+ #
+
+ class Car < Vehicle
+ end
+ EOS
+ end
+
+ let :expected_result do
+ <<~EOS
+ class Car < Vehicle
+ end
+ EOS
+ end
+
+ it 'removes annotation including group headers' do
+ expect(file_content_after_removal).to eq expected_result
+ end
+ end
+
context 'when annotation is before main content and CRLF is used for line breaks' do
let :filename do
'before.rb'
diff --git a/spec/lib/annotate/sti_columns_spec.rb b/spec/lib/annotate/sti_columns_spec.rb
new file mode 100644
index 00000000..0b0a145d
--- /dev/null
+++ b/spec/lib/annotate/sti_columns_spec.rb
@@ -0,0 +1,261 @@
+require_relative '../../spec_helper'
+require 'active_record'
+require 'annotate/sti_columns'
+
+describe Annotate::StiColumns do
+ before(:each) do
+ # Reset the models_loaded flag between tests
+ described_class.instance_variable_set(:@models_loaded, true)
+ end
+
+ def mock_ar_class(name:, table_name:, superclass: ActiveRecord::Base, validators: [], belongs_to: [], enums: {}, stored_attrs: {}, column_names: [], inheritance_column: 'type', descendants: [], subclasses: nil)
+ klass = double(name)
+ allow(klass).to receive(:name).and_return(name)
+ allow(klass).to receive(:table_name).and_return(table_name)
+ allow(klass).to receive(:superclass).and_return(superclass)
+ allow(klass).to receive(:<).with(ActiveRecord::Base).and_return(true)
+ allow(klass).to receive(:column_names).and_return(column_names)
+ allow(klass).to receive(:inheritance_column).and_return(inheritance_column)
+ allow(klass).to receive(:descendants).and_return(descendants)
+ allow(klass).to receive(:subclasses).and_return(subclasses || descendants)
+
+ allow(klass).to receive(:validators).and_return(
+ validators.map { |attr| double('Validator', attributes: [attr]) }
+ )
+ allow(klass).to receive(:reflect_on_all_associations).with(:belongs_to).and_return(
+ belongs_to.map { |fk| double('Association', foreign_key: fk, name: fk.to_s.sub(/_id$/, ''), polymorphic?: false) }
+ )
+ allow(klass).to receive(:defined_enums).and_return(enums)
+ allow(klass).to receive(:stored_attributes).and_return(stored_attrs)
+
+ klass
+ end
+
+ def mock_column(name)
+ double('Column', name: name.to_s)
+ end
+
+ describe '.columns_referenced_in' do
+ it 'collects columns from validators' do
+ klass = mock_ar_class(name: 'Car', table_name: 'vehicles', validators: [:num_doors, :color])
+ expect(described_class.columns_referenced_in(klass)).to eq Set.new(%w[num_doors color])
+ end
+
+ it 'collects foreign keys from belongs_to associations' do
+ klass = mock_ar_class(name: 'Car', table_name: 'vehicles', belongs_to: [:manufacturer_id])
+ expect(described_class.columns_referenced_in(klass)).to eq Set.new(%w[manufacturer_id])
+ end
+
+ it 'collects enum columns' do
+ klass = mock_ar_class(name: 'Car', table_name: 'vehicles', enums: { 'fuel_type' => { gas: 0, diesel: 1 } })
+ expect(described_class.columns_referenced_in(klass)).to eq Set.new(%w[fuel_type])
+ end
+
+ it 'collects stored attribute columns' do
+ klass = mock_ar_class(name: 'Car', table_name: 'vehicles', stored_attrs: { settings: [:color, :theme] })
+ expect(described_class.columns_referenced_in(klass)).to eq Set.new(%w[settings])
+ end
+
+ it 'combines all sources' do
+ klass = mock_ar_class(name: 'Car', table_name: 'vehicles',
+ validators: [:name], belongs_to: [:owner_id],
+ enums: { 'status' => {} }, stored_attrs: { prefs: [:lang] })
+ expect(described_class.columns_referenced_in(klass)).to eq Set.new(%w[name owner_id status prefs])
+ end
+ end
+
+ describe '.columns_owned_by' do
+ it 'returns columns referenced by subclass but not by superclass' do
+ vehicle = mock_ar_class(name: 'Vehicle', table_name: 'vehicles', validators: [:name])
+ car = mock_ar_class(name: 'Car', table_name: 'vehicles', superclass: vehicle,
+ validators: [:name, :num_doors])
+
+ expect(described_class.columns_owned_by(car)).to eq Set.new(%w[num_doors])
+ end
+
+ it 'returns empty set for non-STI classes' do
+ klass = mock_ar_class(name: 'User', table_name: 'users', superclass: ActiveRecord::Base)
+ expect(described_class.columns_owned_by(klass)).to eq Set.new
+ end
+
+ it 'returns empty set when subclass references no new columns' do
+ vehicle = mock_ar_class(name: 'Vehicle', table_name: 'vehicles', validators: [:name])
+ car = mock_ar_class(name: 'Car', table_name: 'vehicles', superclass: vehicle,
+ validators: [:name])
+
+ expect(described_class.columns_owned_by(car)).to eq Set.new
+ end
+ end
+
+ describe '.partition' do
+ let(:id_col) { mock_column(:id) }
+ let(:type_col) { mock_column(:type) }
+ let(:name_col) { mock_column(:name) }
+ let(:num_doors_col) { mock_column(:num_doors) }
+ let(:payload_col) { mock_column(:payload_capacity) }
+ let(:color_col) { mock_column(:color) }
+ let(:battery_col) { mock_column(:battery_kwh) }
+ let(:settings_col) { mock_column(:settings) }
+
+ let(:all_columns) { [id_col, type_col, name_col, num_doors_col, payload_col, color_col] }
+
+ context 'for a non-STI class' do
+ it 'returns a single group with no label' do
+ klass = mock_ar_class(name: 'User', table_name: 'users',
+ column_names: %w[id name])
+ cols = [id_col, name_col]
+
+ result = described_class.partition(klass, cols)
+
+ expect(result).to eq [[nil, cols]]
+ end
+ end
+
+ context 'for an STI subclass' do
+ it 'separates owned columns from shared columns' do
+ vehicle = mock_ar_class(name: 'Vehicle', table_name: 'vehicles', validators: [:name],
+ column_names: %w[id type name num_doors payload_capacity color])
+ car = mock_ar_class(name: 'Car', table_name: 'vehicles', superclass: vehicle,
+ validators: [:name, :num_doors])
+
+ result = described_class.partition(car, all_columns)
+
+ labels = result.map(&:first)
+ expect(labels).to eq ['Vehicle columns', 'Car columns']
+
+ vehicle_col_names = result[0][1].map { |c| c.name }
+ car_col_names = result[1][1].map { |c| c.name }
+ expect(vehicle_col_names).to include('id', 'type', 'name', 'payload_capacity', 'color')
+ expect(car_col_names).to eq ['num_doors']
+ end
+
+ it 'returns only shared group when subclass owns no columns' do
+ vehicle = mock_ar_class(name: 'Vehicle', table_name: 'vehicles', validators: [:name])
+ car = mock_ar_class(name: 'Car', table_name: 'vehicles', superclass: vehicle,
+ validators: [:name])
+ cols = [id_col, type_col, name_col]
+
+ result = described_class.partition(car, cols)
+
+ expect(result.length).to eq 1
+ expect(result[0][0]).to eq 'Vehicle columns'
+ expect(result[0][1]).to eq cols
+ end
+
+ it 'warns when subclass owns no columns' do
+ vehicle = mock_ar_class(name: 'Vehicle', table_name: 'vehicles', validators: [:name])
+ car = mock_ar_class(name: 'Car', table_name: 'vehicles', superclass: vehicle,
+ validators: [:name])
+ cols = [id_col, type_col, name_col]
+
+ expect($stderr).to receive(:puts).with(/could not determine which columns belong to Car/)
+ expect($stderr).to receive(:puts).with(/Add validations/)
+
+ described_class.partition(car, cols)
+ end
+ end
+
+ context 'for an STI base class' do
+ it 'groups columns by owning subclass' do
+ vehicle = mock_ar_class(name: 'Vehicle', table_name: 'vehicles', validators: [:name],
+ column_names: %w[id type name num_doors payload_capacity color])
+ car = mock_ar_class(name: 'Car', table_name: 'vehicles', superclass: vehicle,
+ validators: [:name, :num_doors])
+ truck = mock_ar_class(name: 'Truck', table_name: 'vehicles', superclass: vehicle,
+ validators: [:name, :payload_capacity])
+ allow(vehicle).to receive(:descendants).and_return([car, truck])
+
+ result = described_class.partition(vehicle, all_columns)
+
+ labels = result.map(&:first)
+ expect(labels).to eq ['Vehicle columns', 'Car columns', 'Truck columns']
+
+ base_col_names = result[0][1].map { |c| c.name }
+ expect(base_col_names).to include('id', 'type', 'name', 'color')
+ expect(base_col_names).not_to include('num_doors', 'payload_capacity')
+ end
+
+ it 'puts unclaimed columns in the base group' do
+ vehicle = mock_ar_class(name: 'Vehicle', table_name: 'vehicles', validators: [],
+ column_names: %w[id type name color])
+ car = mock_ar_class(name: 'Car', table_name: 'vehicles', superclass: vehicle,
+ validators: [:num_doors])
+ allow(vehicle).to receive(:descendants).and_return([car])
+
+ cols = [id_col, type_col, name_col, color_col, num_doors_col]
+ result = described_class.partition(vehicle, cols)
+
+ base_col_names = result[0][1].map { |c| c.name }
+ expect(base_col_names).to include('id', 'type', 'name', 'color')
+ end
+
+ it 'assigns a column to the first subclass that claims it' do
+ vehicle = mock_ar_class(name: 'Vehicle', table_name: 'vehicles', validators: [],
+ column_names: %w[id type color])
+ car = mock_ar_class(name: 'Car', table_name: 'vehicles', superclass: vehicle,
+ validators: [:color])
+ truck = mock_ar_class(name: 'Truck', table_name: 'vehicles', superclass: vehicle,
+ validators: [:color])
+ allow(vehicle).to receive(:descendants).and_return([car, truck])
+
+ cols = [id_col, type_col, color_col]
+ result = described_class.partition(vehicle, cols)
+
+ # Car comes first alphabetically, so it claims 'color'
+ car_group = result.find { |label, _| label == 'Car columns' }
+ truck_group = result.find { |label, _| label == 'Truck columns' }
+ expect(car_group[1].map { |c| c.name }).to include('color')
+ expect(truck_group).to be_nil
+ end
+
+ it 'warns when no columns can be assigned to subclasses' do
+ vehicle = mock_ar_class(name: 'Vehicle', table_name: 'vehicles', validators: [:name],
+ column_names: %w[id type name])
+ car = mock_ar_class(name: 'Car', table_name: 'vehicles', superclass: vehicle,
+ validators: [:name])
+ allow(vehicle).to receive(:descendants).and_return([car])
+
+ expect($stderr).to receive(:puts).with(/found STI subclasses for Vehicle but no columns could be assigned/)
+ expect($stderr).to receive(:puts).with(/Add validations/)
+
+ described_class.partition(vehicle, [id_col, type_col, name_col])
+ end
+
+ it 'assigns stored attributes backing columns to the subclass correctly' do
+ vehicle = mock_ar_class(name: 'Vehicle', table_name: 'vehicles', validators: [:name],
+ column_names: %w[id type name settings])
+ car = mock_ar_class(name: 'Car', table_name: 'vehicles', superclass: vehicle,
+ validators: [:name], stored_attrs: { settings: [:color, :theme] })
+ allow(vehicle).to receive(:descendants).and_return([car])
+
+ cols = [id_col, type_col, name_col, settings_col]
+ result = described_class.partition(vehicle, cols)
+
+ car_group = result.find { |label, _| label == 'Car columns' }
+ expect(car_group).not_to be_nil
+ expect(car_group[1].map { |c| c.name }).to eq ['settings']
+ end
+
+ it 'includes multi-level descendants via descendants' do
+ vehicle = mock_ar_class(name: 'Vehicle', table_name: 'vehicles', validators: [:name],
+ column_names: %w[id type name num_doors battery_kwh])
+ car = mock_ar_class(name: 'Car', table_name: 'vehicles', superclass: vehicle,
+ validators: [:name, :num_doors])
+ electric_car = mock_ar_class(name: 'ElectricCar', table_name: 'vehicles', superclass: car,
+ validators: [:name, :num_doors, :battery_kwh])
+ # descendants returns all levels, subclasses returns only direct
+ allow(vehicle).to receive(:descendants).and_return([car, electric_car])
+ allow(car).to receive(:descendants).and_return([electric_car])
+
+ cols = [id_col, type_col, name_col, num_doors_col, battery_col]
+ result = described_class.partition(vehicle, cols)
+
+ labels = result.map(&:first)
+ expect(labels).to include('Vehicle columns', 'Car columns', 'ElectricCar columns')
+
+ electric_group = result.find { |label, _| label == 'ElectricCar columns' }
+ expect(electric_group[1].map { |c| c.name }).to eq ['battery_kwh']
+ end
+ end
+ end
+end