Rails 7.1 composite keys in partitioned tables & returning :id

[pkopac]

SomeModel represents partitioned table (by partition_key), where primary key is %w[id partition_key] so when I do this:

        SomeModel.import!(
          array_of_attributes,
          on_duplicate_key_update: { conflict_target: %i[id partition_key], columns: %i[...] },
          returning: %i[id some other columns]
        ).results

it worked fine in Rails 7.0, but in Rails 7.1 when calling id= it expects eg. a Hash, not just the value of id. So, for compatibility I added this:

module CompositeKeyCompatibility
  extend ActiveSupport::Concern

  # Compatibility override for Rails 7.1 composite key and activerecord-import (1.7.0) gem
  # When using returning it tries to set id=123, but this method requires a Hash.
  def id=(value)
    value = self.class.primary_key.map { |col| col == 'id' ? value : public_send(col) } unless value.is_a? Enumerable
    super
  end
end

But the library should probably do something similar so that it works out of the box.