Extract `primary_key` method into `ColumnMethods`

4 files changed, 68 insertions, 42 deletions

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 a2777fcd0a..a1c7d69b67 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/schema_definitions.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/schema_definitions.rb
@@ -142,6 +142,14 @@ module ActiveRecord
       end
     end
 
+    module ColumnMethods
+      # Appends a primary key definition to the table definition.
+      # Can be called multiple times, but this is probably not a good idea.
+      def primary_key(name, type = :primary_key, **options)
+        column(name, type, options.merge(primary_key: true))
+      end
+    end
+
     # Represents the schema of an SQL table in an abstract way. This class
     # provides methods for manipulating the schema representation.
     #
@@ -163,6 +171,8 @@ module ActiveRecord
     # The table definitions
     # The Columns are stored as a ColumnDefinition in the +columns+ attribute.
     class TableDefinition
+      include ColumnMethods
+
       # An array of ColumnDefinition objects, representing the column changes
       # that have been defined.
       attr_accessor :indexes
@@ -181,12 +191,6 @@ module ActiveRecord
 
       def columns; @columns_hash.values; end
 
-      # Appends a primary key definition to the table definition.
-      # Can be called multiple times, but this is probably not a good idea.
-      def primary_key(name, type = :primary_key, options = {})
-        column(name, type, options.merge(:primary_key => true))
-      end
-
       # Returns a ColumnDefinition for the column with name +name+.
       def [](name)
         @columns_hash[name.to_s]
@@ -462,6 +466,7 @@ module ActiveRecord
     # Available transformations are:
     #
     #   change_table :table do |t|
+    #     t.primary_key
     #     t.column
     #     t.index
     #     t.rename_index
@@ -490,6 +495,8 @@ module ActiveRecord
     #   end
     #
     class Table
+      include ColumnMethods
+
       attr_reader :name
 
       def initialize(table_name, base)
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 8db4bcd7e3..611bdb1758 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
@@ -6,13 +6,21 @@ module ActiveRecord
     class AbstractMysqlAdapter < AbstractAdapter
       include Savepoints
 
-      class TableDefinition < ActiveRecord::ConnectionAdapters::TableDefinition
-        def primary_key(name, type = :primary_key, options = {})
-          options[:auto_increment] ||= type == :bigint
+      module ColumnMethods
+        def primary_key(name, type = :primary_key, **options)
+          options[:auto_increment] = true if type == :bigint
           super
         end
       end
 
+      class TableDefinition < ActiveRecord::ConnectionAdapters::TableDefinition
+        include ColumnMethods
+      end
+
+      class Table < ActiveRecord::ConnectionAdapters::Table
+        include ColumnMethods
+      end
+
       class SchemaCreation < AbstractAdapter::SchemaCreation
         def visit_AddColumn(o)
           add_column_position!(super, column_options(o))
@@ -57,6 +65,10 @@ module ActiveRecord
         end
       end
 
+      def update_table_definition(table_name, base) # :nodoc:
+        Table.new(table_name, base)
+      end
+
       def schema_creation
         SchemaCreation.new self
       end
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql/schema_definitions.rb b/activerecord/lib/active_record/connection_adapters/postgresql/schema_definitions.rb
index b9078d4c86..5362b1cffe 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql/schema_definitions.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql/schema_definitions.rb
@@ -2,6 +2,39 @@ module ActiveRecord
   module ConnectionAdapters
     module PostgreSQL
       module ColumnMethods
+        # Defines the primary key field.
+        # Use of the native PostgreSQL UUID type is supported, and can be used
+        # by defining your tables as such:
+        #
+        #   create_table :stuffs, id: :uuid do |t|
+        #     t.string :content
+        #     t.timestamps
+        #   end
+        #
+        # By default, this will use the +uuid_generate_v4()+ function from the
+        # +uuid-ossp+ extension, which MUST be enabled on your database. To enable
+        # the +uuid-ossp+ extension, you can use the +enable_extension+ method in your
+        # migrations. To use a UUID primary key without +uuid-ossp+ enabled, you can
+        # set the +:default+ option to +nil+:
+        #
+        #   create_table :stuffs, id: false do |t|
+        #     t.primary_key :id, :uuid, default: nil
+        #     t.uuid :foo_id
+        #     t.timestamps
+        #   end
+        #
+        # You may also pass a different UUID generation function from +uuid-ossp+
+        # or another library.
+        #
+        # Note that setting the UUID primary key default value to +nil+ will
+        # require you to assure that you always provide a UUID value before saving
+        # a record (as primary keys cannot be +nil+). This might be done via the
+        # +SecureRandom.uuid+ method and a +before_save+ callback, for instance.
+        def primary_key(name, type = :primary_key, **options)
+          options[:default] = options.fetch(:default, 'uuid_generate_v4()') if type == :uuid
+          super
+        end
+
         def xml(*args)
           options = args.extract_options!
           column(args[0], :xml, options)
@@ -96,39 +129,6 @@ module ActiveRecord
       class TableDefinition < ActiveRecord::ConnectionAdapters::TableDefinition
         include ColumnMethods
 
-        # Defines the primary key field.
-        # Use of the native PostgreSQL UUID type is supported, and can be used
-        # by defining your tables as such:
-        #
-        #   create_table :stuffs, id: :uuid do |t|
-        #     t.string :content
-        #     t.timestamps
-        #   end
-        #
-        # By default, this will use the +uuid_generate_v4()+ function from the
-        # +uuid-ossp+ extension, which MUST be enabled on your database. To enable
-        # the +uuid-ossp+ extension, you can use the +enable_extension+ method in your
-        # migrations. To use a UUID primary key without +uuid-ossp+ enabled, you can
-        # set the +:default+ option to +nil+:
-        #
-        #   create_table :stuffs, id: false do |t|
-        #     t.primary_key :id, :uuid, default: nil
-        #     t.uuid :foo_id
-        #     t.timestamps
-        #   end
-        #
-        # You may also pass a different UUID generation function from +uuid-ossp+
-        # or another library.
-        #
-        # Note that setting the UUID primary key default value to +nil+ will
-        # require you to assure that you always provide a UUID value before saving
-        # a record (as primary keys cannot be +nil+). This might be done via the
-        # +SecureRandom.uuid+ method and a +before_save+ callback, for instance.
-        def primary_key(name, type = :primary_key, options = {})
-          options[:default] = options.fetch(:default, 'uuid_generate_v4()') if type == :uuid
-          super
-        end
-
         def new_column_definition(name, type, options) # :nodoc:
           column = super
           column.array = options[:array]
diff --git a/activerecord/test/cases/migration/change_table_test.rb b/activerecord/test/cases/migration/change_table_test.rb
index 7010af5434..08bb207dbb 100644
--- a/activerecord/test/cases/migration/change_table_test.rb
+++ b/activerecord/test/cases/migration/change_table_test.rb
@@ -100,6 +100,13 @@ module ActiveRecord
         end
       end
 
+      def test_primary_key_creates_primary_key_column
+        with_change_table do |t|
+          @connection.expect :add_column, nil, [:delete_me, :id, :primary_key, primary_key: true, first: true]
+          t.primary_key :id, first: true
+        end
+      end
+
       def test_integer_creates_integer_column
         with_change_table do |t|
           @connection.expect :add_column, nil, [:delete_me, :foo, :integer, {}]