Add bulk versions of all constraint management for Postgres

Author: ccutrer

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

@@ -311,8 +311,8 @@ def exclusion_constraint(*args)
         #  t.remove_exclusion_constraint(name: "price_check")
         #
         # See {connection.remove_exclusion_constraint}[rdoc-ref:SchemaStatements#remove_exclusion_constraint]
-        def remove_exclusion_constraint(*args)
-          @base.remove_exclusion_constraint(name, *args)
+        def remove_exclusion_constraint(*args, **options)
+          @base.remove_exclusion_constraint(name, *args, **options)
         end
 
         # Adds a unique constraint.
@@ -329,8 +329,8 @@ def unique_constraint(*args)
         #  t.remove_unique_constraint(name: "unique_position")
         #
         # See {connection.remove_unique_constraint}[rdoc-ref:SchemaStatements#remove_unique_constraint]
-        def remove_unique_constraint(*args)
-          @base.remove_unique_constraint(name, *args)
+        def remove_unique_constraint(*args, **options)
+          @base.remove_unique_constraint(name, *args, **options)
         end
       end
 

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

@@ -1010,6 +1010,76 @@ def change_column_null_for_alter(table_name, column_name, null, default = nil)
             end
           end
 
+          def add_foreign_key_for_alter(from_table, to_table, **options)
+            assert_valid_deferrable(options[:deferrable])
+            return unless use_foreign_keys?
+
+            return if options[:if_not_exists] == true && foreign_key_exists?(from_table, to_table, **options.slice(:column))
+
+            options = foreign_key_options(from_table, to_table, options)
+            td = create_table_definition from_table
+            fk = td.new_foreign_key_definition to_table, options
+
+            "ADD #{schema_creation.accept(fk)}"
+          end
+
+          def remove_foreign_key_for_alter(from_table, to_table = nil, **options)
+            return unless use_foreign_keys?
+            return if options.delete(:if_exists) == true && !foreign_key_exists?(from_table, to_table)
+    
+            fk_name_to_delete = foreign_key_for!(from_table, to_table: to_table, **options).name
+            remove_constraint_for_alter(from_table, fk_name_to_delete, **options)
+          end
+
+          def add_check_constraint_for_alter(table_name, expression, if_not_exists: false, **options)
+            options = check_constraint_options(table_name, expression, options)
+            return if if_not_exists && check_constraint_exists?(table_name, **options)
+
+            td = create_table_definition(table_name)
+            constraint = td.new_check_constraint_definition(expression, options)
+
+            "ADD #{schema_creation.accept(constraint)}"
+          end
+
+          def remove_constraint_for_alter(table_name, constraint, if_exists: false, **options)
+            "DROP CONSTRAINT#{' IF EXISTS' if if_exists} #{quote_column_name(constraint)}"
+          end
+
+          def remove_check_constraint_for_alter(table_name, expression = nil, **options)
+            chk_name_to_delete = check_constraint_name(table_name, expression: expression, **options)
+
+            remove_constraint_for_alter(table_name, chk_name_to_delete, **options)
+          end
+
+          def add_exclusion_constraint_for_alter(table_name, expression, options)
+            options = exclusion_constraint_options(table_name, expression, options)
+            td = create_table_definition(table_name)
+            constraint = td.new_exclusion_constraint_definition(expression, **options)
+
+            "ADD #{schema_creation.accept(constraint)}"
+          end
+
+          def remove_exclusion_constraint_for_alter(table_name, expression = nil, **options)
+            excl_name_to_delete = exclusion_constraint_name(table_name, expression: expression, **options)
+
+            remove_constraint_for_alter(table_name, excl_name_to_delete, **options)
+          end
+
+          def add_unique_constraint_for_alter(table_name, column_name, options = {})
+            options = unique_constraint_options(table_name, column_name, options)
+
+            td = create_table_definition(table_name)
+            constraint = td.new_unique_constraint_definition(column_name, options)
+
+            "ADD #{schema_creation.accept(constraint)}"
+          end
+
+          def remove_unique_constraint_for_alter(table_name, column_name = nil, **options)
+            unique_name_to_delete = unique_constraint_name(table_name, column: column_name, **options)
+
+            remove_constraint_for_alter(table_name, unique_name_to_delete, **options)
+          end
+
           def add_index_opclass(quoted_columns, **options)
             opclasses = options_for_index_columns(options[:opclass])
             quoted_columns.each do |name, column|

activerecord/test/cases/migration_test.rb

@@ -1214,13 +1214,15 @@ def test_add_and_remove_index
   class BulkAlterTableMigrationsTest < ActiveRecord::TestCase
     def setup
       @connection = Person.lease_connection
+      @connection.create_table(:delete_me2, force: true) { |t| }
       @connection.create_table(:delete_me, force: true) { |t| }
       Person.reset_column_information
       Person.reset_sequence_name
     end
 
     teardown do
       Person.lease_connection.drop_table(:delete_me) rescue nil
+      Person.lease_connection.drop_table(:delete_me2) rescue nil
     end
 
     def test_adding_multiple_columns
@@ -1485,6 +1487,43 @@ def test_updating_auto_increment
       end
     end
 
+    if current_adapter?(:PostgreSQLAdapter)
+      def test_constraints
+        conn = Person.lease_connection
+
+        conn.add_reference :delete_me, :delete_me2
+
+        assert_queries_count(1, include_schema: true) do
+          with_bulk_change_table do |t|
+            t.foreign_key :delete_me2
+            t.check_constraint "id IS NOT NULL", name: "id_chk"
+            t.exclusion_constraint "id WITH =", name: "id_exclusion"
+            t.unique_constraint "id", name: "id_uniq"
+          end
+        end
+
+        assert conn.foreign_key_exists?(:delete_me, :delete_me2)
+        assert conn.check_constraint_exists?(:delete_me, name: "id_chk")
+        assert conn.exclusion_constraints(:delete_me).any? { |c| c.name == "id_exclusion" }
+        assert conn.unique_constraints(:delete_me).any? { |c| c.name == "id_uniq" }
+
+        assert_queries_count(2, include_schema: true) do # one extra query to find the FK name
+          with_bulk_change_table do |t|
+            t.remove_foreign_key :delete_me2
+            t.remove_check_constraint name: "id_chk"
+            t.remove_exclusion_constraint name: "id_exclusion"
+            t.remove_unique_constraint name: "id_uniq", if_exists: true
+            t.remove_unique_constraint name: "non_existinent", if_exists: true
+          end
+        end
+
+        assert_not conn.foreign_key_exists?(:delete_me, :delete_me2)
+        assert_not conn.check_constraint_exists?(:delete_me, name: "id_chk")
+        assert_not conn.exclusion_constraints(:delete_me).any? { |c| c.name == "id_exclusion" }
+        assert_not conn.unique_constraints(:delete_me).any? { |c| c.name == "id_uniq" }
+      end
+    end
+
     def test_changing_index
       with_bulk_change_table do |t|
         t.string :username