Add set_constraints helper for PostgreSQL

The docs already talk about how to set up deferrable constraints, but
then rely on the user crafting custom SQL to actually use the feature.
The helper makes it easier to handle juggling multiple specific
constraints and quoting issues.

Author: ccutrer

activerecord/CHANGELOG.md

@@ -1,3 +1,19 @@
+*   Add `set_constraints` helper for PostgreSQL
+
+    ```ruby
+    Post.create!(user_id: -1) # => ActiveRecord::InvalidForeignKey
+
+    Post.transaction do
+      Post.connection.set_constraints(:deferred)
+      p = Post.create!(user_id: -1)
+      u = User.create!
+      p.user = u
+      p.save!
+    end
+    ```
+
+    *Cody Cutrer*
+
 *   Include `ActiveModel::API` in `ActiveRecord::Base`
 
     *Sean Doyle*

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

@@ -132,6 +132,27 @@ def build_explain_clause(options = [])
           "EXPLAIN (#{options.join(", ").upcase})"
         end
 
+        # Set when constraints will be checked for the current transaction.
+        #
+        # Not passing any specific constraint names will set the value for all deferrable constraints.
+        #
+        # [<tt>deferred</tt>]
+        #   Valid values are +:deferred+ or +:immediate+.
+        #
+        # See https://www.postgresql.org/docs/current/sql-set-constraints.html
+        def set_constraints(deferred, *constraints)
+          unless %i[deferred immediate].include?(deferred)
+            raise ArgumentError, "deferred must be :deferred or :immediate"
+          end
+
+          constraints = if constraints.empty?
+            "ALL"
+          else
+            constraints.map { |c| quote_table_name(c) }.join(", ")
+          end
+          execute("SET CONSTRAINTS #{constraints} #{deferred.to_s.upcase}")
+        end
+
         private
           IDLE_TRANSACTION_STATUSES = [PG::PQTRANS_IDLE, PG::PQTRANS_INTRANS, PG::PQTRANS_INERROR]
           private_constant :IDLE_TRANSACTION_STATUSES

activerecord/test/cases/adapters/postgresql/deferred_constraints_test.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+require "cases/helper"
+require "models/author"
+
+class PostgresqlDeferredConstraintsTest < ActiveRecord::PostgreSQLTestCase
+  def setup
+    @connection = ActiveRecord::Base.connection
+    @fk = @connection.foreign_keys("authors").first.name
+    @other_fk = @connection.foreign_keys("lessons_students").first.name
+  end
+
+  def test_defer_constraints
+    assert_raises ActiveRecord::InvalidForeignKey do
+      @connection.set_constraints(:deferred)
+      assert_nothing_raised do
+        Author.create!(author_address_id: -1, name: "John Doe")
+      end
+      @connection.set_constraints(:immediate)
+    end
+  end
+
+  def test_defer_constraints_with_specific_fk
+    assert_raises ActiveRecord::InvalidForeignKey do
+      @connection.set_constraints(:deferred, @fk)
+      assert_nothing_raised do
+        Author.create!(author_address_id: -1, name: "John Doe")
+      end
+      @connection.set_constraints(:immediate, @fk)
+    end
+  end
+
+  def test_defer_constraints_with_multiple_fks
+    assert_raises ActiveRecord::InvalidForeignKey do
+      @connection.set_constraints(:deferred, @other_fk, @fk)
+      assert_nothing_raised do
+        Author.create!(author_address_id: -1, name: "John Doe")
+      end
+      @connection.set_constraints(:immediate, @other_fk, @fk)
+    end
+  end
+
+  def test_defer_constraints_only_defers_single_fk
+    @connection.set_constraints(:deferred, @other_fk)
+    assert_raises ActiveRecord::InvalidForeignKey do
+      Author.create!(author_address_id: -1, name: "John Doe")
+    end
+  end
+
+  def test_set_constraints_requires_valid_value
+    assert_raises ArgumentError do
+      @connection.set_constraints(:invalid)
+    end
+  end
+end

activerecord/test/cases/migration/exclusion_constraint_test.rb

@@ -162,7 +162,7 @@ def test_added_deferrable_initially_immediate_exclusion_constraint
 
           assert_nothing_raised do
             Invoice.transaction(requires_new: true) do
-              Invoice.connection.exec_query("SET CONSTRAINTS invoices_date_overlap DEFERRED")
+              Invoice.connection.set_constraints(:deferred, "invoices_date_overlap")
               Invoice.create!(start_date: "2020-12-31", end_date: "2021-01-01")
               invoice.update!(end_date: "2020-12-31")
 

activerecord/test/schema/schema.rb

@@ -97,7 +97,7 @@
   create_table :author_addresses, force: true do |t|
   end
 
-  add_foreign_key :authors, :author_addresses
+  add_foreign_key :authors, :author_addresses, deferrable: :immediate
 
   create_table :author_favorites, force: true do |t|
     t.column :author_id, :integer
@@ -706,7 +706,7 @@
     t.integer :college_id
   end
 
-  add_foreign_key :lessons_students, :students, on_delete: :cascade
+  add_foreign_key :lessons_students, :students, on_delete: :cascade, deferrable: :immediate
 
   create_table :lint_models, force: true
 

guides/source/active_record_postgresql.md

@@ -633,11 +633,11 @@ ActiveRecord::Base.connection.transaction do
 end
 ```
 
-When the `:deferrable` option is set to `:immediate`, let the foreign keys keep the default behavior of checking the constraint immediately, but allow manually deferring the checks using `SET CONSTRAINTS ALL DEFERRED` within a transaction. This will cause the foreign keys to be checked when the transaction is committed:
+When the `:deferrable` option is set to `:immediate`, let the foreign keys keep the default behavior of checking the constraint immediately, but allow manually deferring the checks using `set_constraints` within a transaction. This will cause the foreign keys to be checked when the transaction is committed:
 
 ```ruby
-ActiveRecord::Base.transaction do
-  ActiveRecord::Base.connection.execute("SET CONSTRAINTS ALL DEFERRED")
+ActiveRecord::Base.connection.transaction do
+  ActiveRecord::Base.connection.set_constraints(:deferred)
   person = Person.create(alias_id: SecureRandom.uuid, name: "John Doe")
   Alias.create(id: person.alias_id, person_id: person.id, name: "jaydee")
 end