Correctly dump composite primary key

Example:

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

7 files changed, 75 insertions, 36 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/abstract/schema_creation.rb b/activerecord/lib/active_record/connection_adapters/abstract/schema_creation.rb
index 18d943f452..236f577a38 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/schema_creation.rb
+++ b/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)}
diff --git a/activerecord/lib/active_record/connection_adapters/abstract/schema_definitions.rb b/activerecord/lib/active_record/connection_adapters/abstract/schema_definitions.rb
index 2c76dd1518..10329de5f4 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/schema_definitions.rb
+++ b/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
 
diff --git a/activerecord/lib/active_record/connection_adapters/abstract/schema_statements.rb b/activerecord/lib/active_record/connection_adapters/abstract/schema_statements.rb
index 16aefc94ab..33a2afeba2 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/schema_statements.rb
+++ b/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?
diff --git a/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb b/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
index 597f0d0597..bc350a0ceb 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
+++ b/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)
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb b/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb
index 69aa02ccf4..04b1cc479f 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb
+++ b/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.
diff --git a/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb b/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb
index 24fc67938d..452102a357 100644
--- a/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb
+++ b/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:
diff --git a/activerecord/lib/active_record/schema_dumper.rb b/activerecord/lib/active_record/schema_dumper.rb
index fd48ea402a..b1f0be2a25 100644
--- a/activerecord/lib/active_record/schema_dumper.rb
+++ b/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