From e419c8cc3fd0c053fdd2847173dbd95fb4524077 Mon Sep 17 00:00:00 2001 From: Jordan Owens Date: Thu, 29 Dec 2016 16:22:42 -0500 Subject: [PATCH] Add support for composite_primary_keys gem --- gemfiles/3.2.gemfile | 5 ++--- gemfiles/4.0.gemfile | 5 ++--- gemfiles/4.1.gemfile | 5 ++--- gemfiles/4.2.gemfile | 9 ++------- gemfiles/5.0.gemfile | 5 ++--- .../adapters/postgresql_adapter.rb | 19 ++++++++++++++++--- lib/activerecord-import/import.rb | 12 +++++++++--- test/import_test.rb | 19 +++++++++++++++++++ test/models/book.rb | 2 ++ test/models/tag.rb | 4 ++++ test/schema/generic_schema.rb | 11 +++++++++++ .../on_duplicate_key_ignore.rb | 12 ++++++++++++ .../on_duplicate_key_update.rb | 12 ++++++++++++ test/test_helper.rb | 1 + 14 files changed, 96 insertions(+), 25 deletions(-) create mode 100644 test/models/tag.rb diff --git a/gemfiles/3.2.gemfile b/gemfiles/3.2.gemfile index 49a57f74..330655f7 100644 --- a/gemfiles/3.2.gemfile +++ b/gemfiles/3.2.gemfile @@ -1,3 +1,2 @@ -platforms :ruby do - gem 'activerecord', '~> 3.2.0' -end +gem 'activerecord', '~> 3.2.0' +gem 'composite_primary_keys', '~> 5.0' diff --git a/gemfiles/4.0.gemfile b/gemfiles/4.0.gemfile index 33d6c241..5200f073 100644 --- a/gemfiles/4.0.gemfile +++ b/gemfiles/4.0.gemfile @@ -1,3 +1,2 @@ -platforms :ruby do - gem 'activerecord', '~> 4.0.0' -end +gem 'activerecord', '~> 4.0.0' +gem 'composite_primary_keys', '~> 6.0' diff --git a/gemfiles/4.1.gemfile b/gemfiles/4.1.gemfile index be5fdf95..79021b8b 100644 --- a/gemfiles/4.1.gemfile +++ b/gemfiles/4.1.gemfile @@ -1,3 +1,2 @@ -platforms :ruby do - gem 'activerecord', '~> 4.1.0' -end +gem 'activerecord', '~> 4.1.0' +gem 'composite_primary_keys', '~> 7.0' diff --git a/gemfiles/4.2.gemfile b/gemfiles/4.2.gemfile index 52563e51..aea3d9ff 100644 --- a/gemfiles/4.2.gemfile +++ b/gemfiles/4.2.gemfile @@ -1,7 +1,2 @@ -platforms :ruby do - gem 'activerecord', '~> 4.2.0' -end - -platforms :jruby do - gem 'activerecord', '~> 4.2.0' -end +gem 'activerecord', '~> 4.2.0' +gem 'composite_primary_keys', '~> 8.0' diff --git a/gemfiles/5.0.gemfile b/gemfiles/5.0.gemfile index 7fdb62ae..ca6a20fb 100644 --- a/gemfiles/5.0.gemfile +++ b/gemfiles/5.0.gemfile @@ -1,3 +1,2 @@ -platforms :ruby do - gem 'activerecord', '~> 5.0.0' -end +gem 'activerecord', '~> 5.0.0' +gem 'composite_primary_keys', '~> 9.0' diff --git a/lib/activerecord-import/adapters/postgresql_adapter.rb b/lib/activerecord-import/adapters/postgresql_adapter.rb index 4f792fa1..4e3d1ed5 100644 --- a/lib/activerecord-import/adapters/postgresql_adapter.rb +++ b/lib/activerecord-import/adapters/postgresql_adapter.rb @@ -44,7 +44,8 @@ def post_sql_statements( table_name, options ) # :nodoc: sql += super(table_name, options) unless options[:no_returning] || options[:primary_key].blank? - sql << "RETURNING #{options[:primary_key]}" + primary_key = Array(options[:primary_key]) + sql << " RETURNING (#{primary_key.join(', ')})" end sql @@ -136,8 +137,20 @@ def sql_for_conflict_target( args = {} ) end def sql_for_default_conflict_target( table_name ) - conflict_target = primary_key( table_name ) - "(#{conflict_target}) " if conflict_target + pks = select_values(<<-SQL.strip_heredoc, "SCHEMA") + WITH pk_constraint AS ( + SELECT conrelid, unnest(conkey) AS connum FROM pg_constraint + WHERE contype = 'p' + AND conrelid = #{quote(quote_table_name(table_name))}::regclass + ), cons AS ( + SELECT conrelid, connum, row_number() OVER() AS rownum FROM pk_constraint + ) + SELECT attr.attname FROM pg_attribute attr + INNER JOIN cons ON attr.attrelid = cons.conrelid AND attr.attnum = cons.connum + ORDER BY cons.rownum + SQL + conflict_target = pks.join(', ') + "(#{conflict_target}) " if conflict_target.present? end # Return true if the statement is a duplicate key record error diff --git a/lib/activerecord-import/import.rb b/lib/activerecord-import/import.rb index 2e1ea742..ce7e6579 100644 --- a/lib/activerecord-import/import.rb +++ b/lib/activerecord-import/import.rb @@ -404,9 +404,15 @@ def import_helper( *args ) # Force the primary key col into the insert if it's not # on the list and we are using a sequence and stuff a nil # value for it into each row so the sequencer will fire later - if !column_names.include?(primary_key) && connection.prefetch_primary_key? && sequence_name - column_names << primary_key - array_of_attributes.each { |a| a << nil } + symbolized_column_names = Array(column_names).map(&:to_sym) + symbolized_primary_key = Array(primary_key).map(&:to_sym) + + if !symbolized_primary_key.to_set.subset?(symbolized_column_names.to_set) && connection.prefetch_primary_key? && sequence_name + column_count = column_names.size + column_names.concat(primary_key).uniq! + columns_added = column_names.size - column_count + new_fields = Array.new(columns_added) + array_of_attributes.each { |a| a.concat(new_fields) } end # record timestamps unless disabled in ActiveRecord::Base diff --git a/test/import_test.rb b/test/import_test.rb index dae9927e..1d4869dd 100644 --- a/test/import_test.rb +++ b/test/import_test.rb @@ -59,6 +59,25 @@ end end + describe "with composite primary keys" do + it "should import models successfully" do + tags = [Tag.new(tag_id: 1, publisher_id: 1, tag: 'Mystery')] + + assert_difference "Tag.count", +1 do + Tag.import tags + end + end + + it "should import array of values successfully" do + columns = [:tag_id, :publisher_id, :tag] + values = [[1, 1, 'Mystery'], [2, 1, 'Science']] + + assert_difference "Tag.count", +2 do + Tag.import columns, values, validate: false + end + end + end + describe "with STI models" do it "should import models successfully" do dictionaries = [Dictionary.new(author_name: "Noah Webster", title: "Webster's Dictionary")] diff --git a/test/models/book.rb b/test/models/book.rb index b51acf6d..4148d41d 100644 --- a/test/models/book.rb +++ b/test/models/book.rb @@ -1,5 +1,7 @@ class Book < ActiveRecord::Base belongs_to :topic, inverse_of: :books + belongs_to :tag, foreign_key: [:tag_id, :parent_id] + has_many :chapters, inverse_of: :book has_many :discounts, as: :discountable has_many :end_notes, inverse_of: :book diff --git a/test/models/tag.rb b/test/models/tag.rb new file mode 100644 index 00000000..3509f0ee --- /dev/null +++ b/test/models/tag.rb @@ -0,0 +1,4 @@ +class Tag < ActiveRecord::Base + self.primary_keys = :tag_id, :publisher_id + has_many :books, inverse_of: :tag +end diff --git a/test/schema/generic_schema.rb b/test/schema/generic_schema.rb index b85df788..a158892d 100644 --- a/test/schema/generic_schema.rb +++ b/test/schema/generic_schema.rb @@ -62,6 +62,8 @@ t.datetime :updated_on t.date :publish_date t.integer :topic_id + t.integer :tag_id + t.integer :publisher_id t.boolean :for_sale, default: true t.integer :status, default: 0 t.string :type @@ -151,4 +153,13 @@ t.text :config t.text :settings end + + execute %( + CREATE TABLE IF NOT EXISTS tags ( + tag_id INT NOT NULL, + publisher_id INT NOT NULL, + tag VARCHAR(50), + PRIMARY KEY (tag_id, publisher_id) + ); + ).split.join(' ').strip end diff --git a/test/support/shared_examples/on_duplicate_key_ignore.rb b/test/support/shared_examples/on_duplicate_key_ignore.rb index f48109c7..f54f1d29 100644 --- a/test/support/shared_examples/on_duplicate_key_ignore.rb +++ b/test/support/shared_examples/on_duplicate_key_ignore.rb @@ -11,6 +11,18 @@ def should_support_on_duplicate_key_ignore Topic.import topics, on_duplicate_key_ignore: true, validate: false end end + + context "with composite primary keys" do + it "should import array of values successfully" do + columns = [:tag_id, :publisher_id, :tag] + values = [[1, 1, 'Mystery'], [1, 1, 'Science']] + + assert_difference "Tag.count", +1 do + Tag.import columns, values, on_duplicate_key_ignore: true, validate: false + end + assert_equal 'Mystery', Tag.first.tag + end + end end context "with :ignore" do diff --git a/test/support/shared_examples/on_duplicate_key_update.rb b/test/support/shared_examples/on_duplicate_key_update.rb index e743383c..ff1f9f42 100644 --- a/test/support/shared_examples/on_duplicate_key_update.rb +++ b/test/support/shared_examples/on_duplicate_key_update.rb @@ -76,6 +76,18 @@ def should_support_basic_on_duplicate_key_update assert_equal 'DISCOUNT2', updated_promotion.code end end + + context "with composite primary keys" do + it "should import array of values successfully" do + columns = [:tag_id, :publisher_id, :tag] + Tag.import columns, [[1, 1, 'Mystery']], validate: false + + assert_difference "Tag.count", +0 do + Tag.import columns, [[1, 1, 'Science']], on_duplicate_key_update: [:tag], validate: false + end + assert_equal 'Science', Tag.first.tag + end + end end context "with :on_duplicate_key_update turned off" do diff --git a/test/test_helper.rb b/test/test_helper.rb index b03559fc..2f91fc70 100644 --- a/test/test_helper.rb +++ b/test/test_helper.rb @@ -26,6 +26,7 @@ require 'timecop' require 'chronic' +require 'composite_primary_keys' require "ruby-debug" if RUBY_VERSION.to_f < 1.9