fk: generalize using `AlterTable` and `SchemaCreation`.

6 files changed, 64 insertions, 49 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 47fe501752..ad62eab4d2 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/schema_creation.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/schema_creation.rb
@@ -18,11 +18,25 @@ module ActiveRecord
           add_column_options!(sql, column_options(o))
         end
 
+        def visit_AddForeignKey(o)
+          <<-SQL
+ADD CONSTRAINT #{quote_column_name(o.name)}
+FOREIGN KEY (#{quote_column_name(o.column)})
+  REFERENCES #{quote_table_name(o.to_table)} (#{quote_column_name(o.primary_key)})
+          SQL
+        end
+
+        def visit_DropForeignKey(name)
+          "DROP CONSTRAINT #{name}"
+        end
+
         private
 
           def visit_AlterTable(o)
             sql = "ALTER TABLE #{quote_table_name(o.name)} "
             sql << o.adds.map { |col| visit_AddColumn col }.join(' ')
+            sql << o.foreign_key_adds.map { |fk| visit_AddForeignKey fk }.join(' ')
+            sql << o.foreign_key_drops.map { |fk| visit_DropForeignKey fk }.join(' ')
           end
 
           def visit_ColumnDefinition(o)
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 b132543332..c18ebf1014 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/schema_definitions.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/schema_definitions.rb
@@ -317,14 +317,26 @@ module ActiveRecord
 
     class AlterTable # :nodoc:
       attr_reader :adds
+      attr_reader :foreign_key_adds
+      attr_reader :foreign_key_drops
 
       def initialize(td)
         @td   = td
         @adds = []
+        @foreign_key_adds = []
+        @foreign_key_drops = []
       end
 
       def name; @td.name; end
 
+      def add_foreign_key(to_table, options)
+        @foreign_key_adds << ForeignKeyDefinition.new(name, to_table, options)
+      end
+
+      def drop_foreign_key(name)
+        @foreign_key_drops << name
+      end
+
       def add_column(name, type, options)
         name = name.to_s
         type = type.to_sym
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 55e3fb4477..da6b15dad0 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/schema_statements.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/schema_statements.rb
@@ -625,13 +625,6 @@ module ActiveRecord
       end
       alias :add_belongs_to :add_reference
 
-      def foreign_key_name(table_name, options)
-        options.fetch(:name) do
-          column_name = options.fetch(:column)
-          "#{table_name}_#{column_name}_fk"
-        end
-      end
-
       # Removes the reference(s). Also removes a +type+ column if one exists.
       # <tt>remove_reference</tt>, <tt>remove_references</tt> and <tt>remove_belongs_to</tt> are acceptable.
       #
@@ -649,6 +642,36 @@ module ActiveRecord
       end
       alias :remove_belongs_to :remove_reference
 
+      def foreign_keys(table_name)
+        raise NotImplementedError, "foreign_keys is not implemented"
+      end
+
+      def add_foreign_key(from_table, to_table, options = {})
+        options = {
+          column: options.fetch(:column),
+          primary_key: "id",
+          name: foreign_key_name(from_table, options)
+        }
+        at = create_alter_table from_table
+        at.add_foreign_key to_table, options
+
+        execute schema_creation.accept at
+      end
+
+      def remove_foreign_key(from_table, options = {})
+        at = create_alter_table from_table
+        at.drop_foreign_key foreign_key_name(from_table, options)
+
+        execute schema_creation.accept at
+      end
+
+      def foreign_key_name(table_name, options) # :nodoc:
+        options.fetch(:name) do
+          column_name = options.fetch(:column)
+          "#{table_name}_#{column_name}_fk"
+        end
+      end
+
       def dump_schema_information #:nodoc:
         sm_table = ActiveRecord::Migrator.schema_migrations_table_name
 
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 431db591e6..6ba226765c 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
@@ -10,6 +10,10 @@ module ActiveRecord
           add_column_position!(super, column_options(o))
         end
 
+        def visit_DropForeignKey(name)
+          "DROP FOREIGN KEY #{name}"
+        end
+
         private
 
         def visit_TableDefinition(o)
@@ -525,26 +529,6 @@ module ActiveRecord
         end
       end
 
-      def add_foreign_key(from_table, to_table, options = {})
-        foreign_key_column = options.fetch(:column)
-        referenced_column = "id"
-        foreign_key_name = foreign_key_name(from_table, options)
-        execute <<-SQL
-ALTER TABLE #{quote_table_name(from_table)}
-ADD CONSTRAINT #{foreign_key_name}
-FOREIGN KEY (#{quote_column_name(foreign_key_column)})
-REFERENCES #{quote_table_name(to_table)} (#{quote_column_name(referenced_column)})
-        SQL
-      end
-
-      def remove_foreign_key(from_table, options = {})
-        foreign_key_name = foreign_key_name(from_table, options)
-        execute <<-SQL
-ALTER TABLE #{quote_table_name(from_table)}
-DROP FOREIGN KEY #{foreign_key_name}
-        SQL
-      end
-
       # Maps logical Rails types to MySQL-specific data types.
       def type_to_sql(type, limit = nil, precision = nil, scale = nil)
         case type.to_s
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 7a93d9cde7..c061337e71 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb
@@ -472,26 +472,6 @@ ORDER BY c.conname
           end
         end
 
-        def add_foreign_key(from_table, to_table, options = {})
-          foreign_key_column = options.fetch(:column)
-          referenced_column = "id"
-          foreign_key_name = foreign_key_name(from_table, options)
-          execute <<-SQL
-ALTER TABLE #{quote_table_name(from_table)}
-ADD CONSTRAINT #{foreign_key_name}
-  FOREIGN KEY (#{quote_column_name(foreign_key_column)})
-  REFERENCES #{quote_table_name(to_table)} (#{quote_column_name(referenced_column)})
-          SQL
-        end
-
-        def remove_foreign_key(from_table, options = {})
-          foreign_key_name = foreign_key_name(from_table, options)
-          execute <<-SQL
-ALTER TABLE #{quote_table_name(from_table)}
-DROP CONSTRAINT #{foreign_key_name}
-          SQL
-        end
-
         def index_name_length
           63
         end
diff --git a/activerecord/test/cases/migration/foreign_key_test.rb b/activerecord/test/cases/migration/foreign_key_test.rb
index 655393445d..6ad595668f 100644
--- a/activerecord/test/cases/migration/foreign_key_test.rb
+++ b/activerecord/test/cases/migration/foreign_key_test.rb
@@ -23,8 +23,10 @@ module ActiveRecord
       end
 
       teardown do
-        @connection.execute "DROP TABLE IF EXISTS astronauts"
-        @connection.execute "DROP TABLE IF EXISTS rockets"
+        if defined?(@connection)
+          @connection.execute "DROP TABLE IF EXISTS astronauts"
+          @connection.execute "DROP TABLE IF EXISTS rockets"
+        end
       end
 
       def test_foreign_keys