From 115d54e9a009e46c30e32dc930fa8ebae8674e23 Mon Sep 17 00:00:00 2001 From: Carlos Musetti Date: Tue, 21 Apr 2026 11:18:47 -0300 Subject: [PATCH 1/2] Fix handling of enum attributes --- lib/staging_table/session.rb | 28 +++-- rbs_collection.yaml | 2 - spec/staging_table/enum_attributes_spec.rb | 127 +++++++++++++++++++++ 3 files changed, 148 insertions(+), 9 deletions(-) create mode 100644 spec/staging_table/enum_attributes_spec.rb diff --git a/lib/staging_table/session.rb b/lib/staging_table/session.rb index c620d38..46809bb 100644 --- a/lib/staging_table/session.rb +++ b/lib/staging_table/session.rb @@ -171,16 +171,30 @@ def ensure_table_created! raise TableError, "Staging table has not been created. You must call #create_table or use StagingTable.stage with a block before inserting or transferring data." unless @table_created end + def normalizable?(records) + records.is_a?(ActiveRecord::Relation) || records.respond_to?(:to_a) + end + def normalize_records(records) - if records.is_a?(ActiveRecord::Relation) - records.map(&:attributes) - elsif records.respond_to?(:to_a) - records.to_a.map do |record| - record.is_a?(ActiveRecord::Base) ? record.attributes : record + if normalizable?(records) + return records.map { |record| normalize_record(record) } + end + + records + end + + def normalize_record(record) + return record unless record.is_a?(ActiveRecord::Base) + + if record.respond_to?(:attributes_for_database) + return record.attributes_for_database + elsif record.respond_to?(:read_attribute_for_database) + return record.attributes.keys.each_with_object({}) do |attribute, normalized_attributes| + normalized[attribute] = record.read_attribute_for_database(attribute) end - else - records end + + record.attributes_before_type_cast end end end diff --git a/rbs_collection.yaml b/rbs_collection.yaml index 31c6319..d329757 100644 --- a/rbs_collection.yaml +++ b/rbs_collection.yaml @@ -13,6 +13,4 @@ path: .gem_rbs_collection gems: - name: activerecord - name: activesupport - # Ignore gems without RBS definitions in the collection - name: prism - ignore: true diff --git a/spec/staging_table/enum_attributes_spec.rb b/spec/staging_table/enum_attributes_spec.rb new file mode 100644 index 0000000..be8570d --- /dev/null +++ b/spec/staging_table/enum_attributes_spec.rb @@ -0,0 +1,127 @@ +# frozen_string_literal: true + +require "spec_helper" + +class EnumTestRecord < ActiveRecord::Base + self.table_name = "enum_test_records" + + if ActiveRecord.gem_version >= Gem::Version.new("7.0") + enum :foo, {bar: 2, baz: 3} + else + enum foo: {bar: 2, baz: 3} + end +end + +# Regression coverage for the bug where `Session#normalize_records` calls +# `attributes_before_type_cast` on AR objects, returning the *enum label* +# (e.g. "bar") instead of the underlying database integer (e.g. 2). +RSpec.describe "StagingTable with enum attributes" do + shared_examples "preserves enum integer values" do + let(:enum_values) do + EnumTestRecord.defined_enums.fetch("foo").transform_values(&:to_i) + end + + let(:session) do + StagingTable::Session.new(EnumTestRecord) + end + + before do + ActiveRecord::Base.connection.create_table(EnumTestRecord.table_name, force: true) do |t| + t.string :name + t.string :email + t.integer :foo + t.timestamps null: true + end + EnumTestRecord.reset_column_information + end + + after do + session.drop_table if session.instance_variable_get(:@table_created) + ActiveRecord::Base.connection.drop_table(EnumTestRecord.table_name, if_exists: true) + end + + it "stages the enum's database integer, not the label, when given AR objects" do + record = EnumTestRecord.new( + id: 1, + name: "Alice", + email: "alice@example.com", + foo: :bar + ) + + session.create_table + + expect { session.insert([record]) }.not_to raise_error + + raw_value = session.staging_model.connection.select_value( + "SELECT foo FROM #{session.staging_model.table_name}" + ) + + expect(Integer(raw_value)).to eq(enum_values.fetch("bar")) + end + + it "round-trips an enum value from AR object through staging into the source table" do + original = EnumTestRecord.new( + id: 2, + name: "Bob", + email: "bob@example.com", + foo: :bar + ) + + session.create_table + session.insert([original]) + result = session.transfer + + expect(result.inserted).to eq(1) + expect(EnumTestRecord.count).to eq(1) + expect(EnumTestRecord.first.foo).to eq("bar") + end + + it "stages the enum integer when given an ActiveRecord::Relation" do + EnumTestRecord.create!(name: "Carol", email: "carol@example.com", foo: :bar) + EnumTestRecord.create!(name: "Dave", email: "dave@example.com", foo: :baz) + + session.create_table + + expect { session.insert(EnumTestRecord.all) }.not_to raise_error + + staged_values = session.staging_model.connection.select_values( + "SELECT foo FROM #{session.staging_model.table_name} ORDER BY foo" + ).map { |value| Integer(value) } + + expect(staged_values).to eq([ + enum_values.fetch("bar"), + enum_values.fetch("baz") + ]) + end + + it "still accepts plain hashes with the integer value" do + session.create_table + + expect { + session.insert([{ + name: "Eve", + email: "eve@example.com", + foo: enum_values.fetch("bar") + }]) + }.not_to raise_error + + raw_value = session.staging_model.connection.select_value( + "SELECT foo FROM #{session.staging_model.table_name}" + ) + + expect(Integer(raw_value)).to eq(enum_values.fetch("bar")) + end + end + + context "with PostgreSQL", :postgresql do + include_examples "preserves enum integer values" + end + + context "with MySQL", :mysql do + include_examples "preserves enum integer values" + end + + context "with SQLite", :sqlite do + include_examples "preserves enum integer values" + end +end From d9ca85b7ccd7c2cd41557fe823b052d89d86d121 Mon Sep 17 00:00:00 2001 From: Santiago Calvo Date: Tue, 21 Apr 2026 15:52:32 +0100 Subject: [PATCH 2/2] fix: adjust accumulator attribute to normalized_attributes --- lib/staging_table/session.rb | 20 +++----------------- spec/staging_table/enum_attributes_spec.rb | 2 +- 2 files changed, 4 insertions(+), 18 deletions(-) diff --git a/lib/staging_table/session.rb b/lib/staging_table/session.rb index 46809bb..cd2ad71 100644 --- a/lib/staging_table/session.rb +++ b/lib/staging_table/session.rb @@ -171,30 +171,16 @@ def ensure_table_created! raise TableError, "Staging table has not been created. You must call #create_table or use StagingTable.stage with a block before inserting or transferring data." unless @table_created end - def normalizable?(records) - records.is_a?(ActiveRecord::Relation) || records.respond_to?(:to_a) - end - def normalize_records(records) - if normalizable?(records) - return records.map { |record| normalize_record(record) } - end + return records unless records.is_a?(ActiveRecord::Relation) || records.is_a?(Array) - records + records.map { |record| normalize_record(record) } end def normalize_record(record) return record unless record.is_a?(ActiveRecord::Base) - if record.respond_to?(:attributes_for_database) - return record.attributes_for_database - elsif record.respond_to?(:read_attribute_for_database) - return record.attributes.keys.each_with_object({}) do |attribute, normalized_attributes| - normalized[attribute] = record.read_attribute_for_database(attribute) - end - end - - record.attributes_before_type_cast + record.attributes_for_database end end end diff --git a/spec/staging_table/enum_attributes_spec.rb b/spec/staging_table/enum_attributes_spec.rb index be8570d..78b32d5 100644 --- a/spec/staging_table/enum_attributes_spec.rb +++ b/spec/staging_table/enum_attributes_spec.rb @@ -36,7 +36,7 @@ class EnumTestRecord < ActiveRecord::Base end after do - session.drop_table if session.instance_variable_get(:@table_created) + session.drop_table ActiveRecord::Base.connection.drop_table(EnumTestRecord.table_name, if_exists: true) end