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)

1 files changed, 43 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