Merge pull request #27636 from mtsmfm/disable-referential-integrity-without-superuser-privilege-take-2

Use `SET CONSTRAINTS` for `disable_referential_integrity` without superuser privileges (take 2)

2 files changed, 49 insertions, 4 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/postgresql/referential_integrity.rb b/activerecord/lib/active_record/connection_adapters/postgresql/referential_integrity.rb
index 44a7338bf5..730e7c7137 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql/referential_integrity.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql/referential_integrity.rb
@@ -6,8 +6,50 @@ module ActiveRecord
           true
         end
 
-        def disable_referential_integrity # :nodoc:
+        def disable_referential_integrity(&block) # :nodoc:
           if supports_disable_referential_integrity?
+            if supports_alter_constraint?
+              disable_referential_integrity_with_alter_constraint(&block)
+            else
+              disable_referential_integrity_with_disable_trigger(&block)
+            end
+          else
+            yield
+          end
+        end
+
+        private
+
+          def disable_referential_integrity_with_alter_constraint
+            tables_constraints = execute(<<-SQL).values
+              SELECT table_name, constraint_name
+              FROM information_schema.table_constraints
+              WHERE constraint_type = 'FOREIGN KEY'
+              AND is_deferrable = 'NO'
+            SQL
+
+            execute(
+              tables_constraints.collect { |table, constraint|
+                "ALTER TABLE #{quote_table_name(table)} ALTER CONSTRAINT #{constraint} DEFERRABLE"
+              }.join(";")
+            )
+
+            begin
+              transaction do
+                execute("SET CONSTRAINTS ALL DEFERRED")
+
+                yield
+              end
+            ensure
+              execute(
+                tables_constraints.collect { |table, constraint|
+                  "ALTER TABLE #{quote_table_name(table)} ALTER CONSTRAINT #{constraint} NOT DEFERRABLE"
+                }.join(";")
+              )
+            end
+          end
+
+          def disable_referential_integrity_with_disable_trigger
             original_exception = nil
 
             begin
@@ -39,10 +81,7 @@ Rails needs superuser privileges to disable referential integrity.
               end
             rescue ActiveRecord::ActiveRecordError
             end
-          else
-            yield
           end
-        end
       end
     end
   end
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
index 73cf3ac862..6273614f0c 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
@@ -314,6 +314,12 @@ module ActiveRecord
         postgresql_version >= 90400
       end
 
+      def supports_alter_constraint?
+        # PostgreSQL 9.4 introduces ALTER TABLE ... ALTER CONSTRAINT but it has a bug and fixed in 9.4.2
+        # https://www.postgresql.org/docs/9.4/static/release-9-4-2.html
+        postgresql_version >= 90402
+      end
+
       def get_advisory_lock(lock_id) # :nodoc:
         unless lock_id.is_a?(Integer) && lock_id.bit_length <= 63
           raise(ArgumentError, "Postgres requires advisory lock ids to be a signed 64 bit integer")