Add support for composite_primary_keys gem

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'

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'

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'

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'

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'

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

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

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")]

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

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

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

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

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

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