Correctly dump composite primary key

Example:

    create_table :barcodes, primary_key: ["region", "code"] do |t|
      t.string :region
      t.integer :code
    end

activerecord/CHANGELOG.md

@@ -1,3 +1,14 @@
+*   Correctly dump composite primary key.
+
+    Example:
+
+        create_table :barcodes, primary_key: ["region", "code"] do |t|
+          t.string :region
+          t.integer :code
+        end
+
+    *Ryuta Kamizono*
+
 *   Lookup the attribute name for `restrict_with_error` messages on the
     model class that defines the association.
 

activerecord/lib/active_record/connection_adapters/abstract/schema_creation.rb

@@ -38,14 +38,21 @@ module ActiveRecord
           end
 
           def visit_TableDefinition(o)
-            create_sql = "CREATE#{' TEMPORARY' if o.temporary} TABLE "
-            create_sql << "#{quote_table_name(o.name)} "
-            create_sql << "(#{o.columns.map { |c| accept c }.join(', ')}) " unless o.as
+            create_sql = "CREATE#{' TEMPORARY' if o.temporary} TABLE #{quote_table_name(o.name)} "
+
+            statements = o.columns.map { |c| accept c }
+            statements << accept(o.primary_keys) if o.primary_keys
+
+            create_sql << "(#{statements.join(', ')}) " if statements.present?
             create_sql << "#{o.options}"
             create_sql << " AS #{@conn.to_sql(o.as)}" if o.as
             create_sql
           end
 
+          def visit_PrimaryKeyDefinition(o)
+            "PRIMARY KEY (#{o.name.join(', ')})"
+          end
+
           def visit_AddForeignKey(o)
             sql = <<-SQL.strip_heredoc
               ADD CONSTRAINT #{quote_column_name(o.name)}

activerecord/lib/active_record/connection_adapters/abstract/schema_definitions.rb

@@ -23,6 +23,9 @@ module ActiveRecord
     class ChangeColumnDefinition < Struct.new(:column, :name) #:nodoc:
     end
 
+    class PrimaryKeyDefinition < Struct.new(:name) # :nodoc:
+    end
+
     class ForeignKeyDefinition < Struct.new(:from_table, :to_table, :options) #:nodoc:
       def name
         options[:name]
@@ -207,6 +210,7 @@ module ActiveRecord
         @columns_hash = {}
         @indexes = {}
         @foreign_keys = {}
+        @primary_keys = nil
         @native = types
         @temporary = temporary
         @options = options
@@ -214,6 +218,11 @@ module ActiveRecord
         @name = name
       end
 
+      def primary_keys(name = nil) # :nodoc:
+        @primary_keys = PrimaryKeyDefinition.new(name) if name
+        @primary_keys
+      end
+
       # Returns an array of ColumnDefinition objects for the columns of the table.
       def columns; @columns_hash.values; end
 

activerecord/lib/active_record/connection_adapters/abstract/schema_statements.rb

@@ -93,6 +93,12 @@ module ActiveRecord
                                       (!options.key?(:null)      || c.null == options[:null]) }
       end
 
+      # Returns just a table's primary key
+      def primary_key(table_name)
+        pks = primary_keys(table_name)
+        pks.first if pks.one?
+      end
+
       # Creates a new table with the name +table_name+. +table_name+ may either
       # be a String or a Symbol.
       #
@@ -220,7 +226,11 @@ module ActiveRecord
             Base.get_primary_key table_name.to_s.singularize
           end
 
-          td.primary_key pk, options.fetch(:id, :primary_key), options
+          if pk.is_a?(Array)
+            td.primary_keys pk
+          else
+            td.primary_key pk, options.fetch(:id, :primary_key), options
+          end
         end
 
         yield td if block_given?

activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb

@@ -57,6 +57,7 @@ module ActiveRecord
           create_sql = "CREATE#{' TEMPORARY' if o.temporary} TABLE #{quote_table_name(name)} "
 
           statements = o.columns.map { |c| accept c }
+          statements << accept(o.primary_keys) if o.primary_keys
           statements.concat(o.indexes.map { |column_name, options| index_in_create(name, column_name, options) })
 
           create_sql << "(#{statements.join(', ')}) " if statements.present?
@@ -735,21 +736,25 @@ module ActiveRecord
 
       # Returns a table's primary key and belonging sequence.
       def pk_and_sequence_for(table)
-        execute_and_free("SHOW CREATE TABLE #{quote_table_name(table)}", 'SCHEMA') do |result|
-          create_table = each_hash(result).first[:"Create Table"]
-          if create_table.to_s =~ /PRIMARY KEY\s+(?:USING\s+\w+\s+)?\((.+)\)/
-            keys = $1.split(",").map { |key| key.delete('`"') }
-            keys.length == 1 ? [keys.first, nil] : nil
-          else
-            nil
-          end
+        if pk = primary_key(table)
+          [ pk, nil ]
         end
       end
 
-      # Returns just a table's primary key
-      def primary_key(table)
-        pk_and_sequence = pk_and_sequence_for(table)
-        pk_and_sequence && pk_and_sequence.first
+      def primary_keys(table_name) # :nodoc:
+        raise ArgumentError unless table_name.present?
+
+        schema, name = table_name.to_s.split('.', 2)
+        schema, name = @config[:database], schema unless name # A table was provided without a schema
+
+        select_values(<<-SQL.strip_heredoc, 'SCHEMA')
+          SELECT column_name
+          FROM information_schema.key_column_usage
+          WHERE constraint_name = 'PRIMARY'
+            AND table_schema = #{quote(schema)}
+            AND table_name = #{quote(name)}
+          ORDER BY ordinal_position
+        SQL
       end
 
       def case_sensitive_modifier(node, table_attribute)

activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb

@@ -353,17 +353,19 @@ module ActiveRecord
           nil
         end
 
-        # Returns just a table's primary key
-        def primary_key(table)
-          pks = query(<<-end_sql, 'SCHEMA')
-            SELECT attr.attname
-            FROM pg_attribute attr
-            INNER JOIN pg_constraint cons ON attr.attrelid = cons.conrelid AND attr.attnum = any(cons.conkey)
-            WHERE cons.contype = 'p'
-              AND cons.conrelid = '#{quote_table_name(table)}'::regclass
-          end_sql
-          return nil unless pks.count == 1
-          pks[0][0]
+        def primary_keys(table_name) # :nodoc:
+          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_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
         end
 
         # Renames a table.

activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb

@@ -369,10 +369,9 @@ module ActiveRecord
         end
       end
 
-      def primary_key(table_name) #:nodoc:
+      def primary_keys(table_name) # :nodoc:
         pks = table_structure(table_name).select { |f| f['pk'] > 0 }
-        return nil unless pks.count == 1
-        pks[0]['name']
+        pks.sort_by { |f| f['pk'] }.map { |f| f['name'] }
       end
 
       def remove_index(table_name, options = {}) #:nodoc:

activerecord/lib/active_record/schema_dumper.rb

@@ -112,20 +112,27 @@ HEADER
           tbl = StringIO.new
 
           # first dump primary key column
-          pk = @connection.primary_key(table)
+          if @connection.respond_to?(:primary_keys)
+            pk = @connection.primary_keys(table)
+            pk = pk.first unless pk.size > 1
+          else
+            pk = @connection.primary_key(table)
+          end
 
           tbl.print "  create_table #{remove_prefix_and_suffix(table).inspect}"
-          pkcol = columns.detect { |c| c.name == pk }
-          if pkcol
-            if pk != 'id'
-              tbl.print %Q(, primary_key: "#{pk}")
-            end
+
+          case pk
+          when String
+            tbl.print ", primary_key: #{pk.inspect}" unless pk == 'id'
+            pkcol = columns.detect { |c| c.name == pk }
             pkcolspec = @connection.column_spec_for_primary_key(pkcol)
             if pkcolspec
               pkcolspec.each do |key, value|
                 tbl.print ", #{key}: #{value}"
               end
             end
+          when Array
+            tbl.print ", primary_key: #{pk.inspect}"
           else
             tbl.print ", id: false"
           end

activerecord/test/cases/primary_keys_test.rb

@@ -241,6 +241,33 @@ class PrimaryKeyAnyTypeTest < ActiveRecord::TestCase
   end
 end
 
+class CompositePrimaryKeyTest < ActiveRecord::TestCase
+  include SchemaDumpingHelper
+
+  self.use_transactional_tests = false
+
+  def setup
+    @connection = ActiveRecord::Base.connection
+    @connection.create_table(:barcodes, primary_key: ["region", "code"], force: true) do |t|
+      t.string :region
+      t.integer :code
+    end
+  end
+
+  def teardown
+    @connection.drop_table(:barcodes, if_exists: true)
+  end
+
+  def test_composite_primary_key
+    assert_equal ["region", "code"], @connection.primary_keys("barcodes")
+  end
+
+  def test_collectly_dump_composite_primary_key
+    schema = dump_table_schema "barcodes"
+    assert_match %r{create_table "barcodes", primary_key: \["region", "code"\]}, schema
+  end
+end
+
 if current_adapter?(:MysqlAdapter, :Mysql2Adapter)
   class PrimaryKeyWithAnsiQuotesTest < ActiveRecord::TestCase
     self.use_transactional_tests = false

activerecord/test/cases/test_case.rb

@@ -103,7 +103,7 @@ module ActiveRecord
     # ignored SQL, or better yet, use a different notification for the queries
     # instead examining the SQL content.
     oracle_ignored     = [/^select .*nextval/i, /^SAVEPOINT/, /^ROLLBACK TO/, /^\s*select .* from all_triggers/im, /^\s*select .* from all_constraints/im, /^\s*select .* from all_tab_cols/im]
-    mysql_ignored      = [/^SHOW TABLES/i, /^SHOW FULL FIELDS/, /^SHOW CREATE TABLE /i, /^SHOW VARIABLES /]
+    mysql_ignored      = [/^SHOW TABLES/i, /^SHOW FULL FIELDS/, /^SHOW CREATE TABLE /i, /^SHOW VARIABLES /, /^\s*SELECT column_name\b.*\bFROM information_schema\.key_column_usage\b/im]
     postgresql_ignored = [/^\s*select\b.*\bfrom\b.*pg_namespace\b/im, /^\s*select tablename\b.*from pg_tables\b/im, /^\s*select\b.*\battname\b.*\bfrom\b.*\bpg_attribute\b/im, /^SHOW search_path/i]
     sqlite3_ignored =    [/^\s*SELECT name\b.*\bFROM sqlite_master/im, /^\s*SELECT sql\b.*\bFROM sqlite_master/im]