Add Database#{defer,immediate}_constraints on PostgreSQL for changing handling of deferrable constraints in a transaction

This allows you to easily check deferrable constraints at a
specific point in a transaction, instead of waiting until the
end of the transaction.  For deferrable constraints that are
initially immediate, this allows you to deter the checking of
the constraints.

I thought about using a single set_constraints method for this,
but I think two separate methods results in more readable code.

Author: jeremyevans

CHANGELOG

@@ -1,3 +1,7 @@
+=== master
+
+* Add Database#{defer,immediate}_constraints on PostgreSQL for changing handling of deferrable constraints in a transaction (jeremyevans)
+
 === 5.74.0 (2023-11-01)
 
 * Make generated columns show up in Database#schema when using SQLite 3.37+ (jeremyevans) (#2087)

lib/sequel/adapters/shared/postgres.rb

@@ -498,6 +498,25 @@ def database_type
         :postgres
       end
 
+      # For constraints that are deferrable, defer constraints until 
+      # transaction commit. Options:
+      #
+      # :constraints :: An identifier of the constraint, or an array of
+      #                 identifiers for constraints, to apply this
+      #                 change to specific constraints.
+      # :server :: The server/shard on which to run the query.
+      #
+      # Examples:
+      #
+      #   DB.defer_constraints
+      #   # SET CONSTRAINTS ALL DEFERRED
+      #
+      #   DB.defer_constraints(constraints: [:c1, Sequel[:sc][:c2]])
+      #   # SET CONSTRAINTS "c1", "sc"."s2" DEFERRED
+      def defer_constraints(opts=OPTS)
+        _set_constraints(' DEFERRED', opts)
+      end
+
       # Use PostgreSQL's DO syntax to execute an anonymous code block.  The code should
       # be the literal code string to use in the underlying procedural language.  Options:
       #
@@ -611,6 +630,24 @@ def freeze
         super
       end
 
+      # Immediately apply deferrable constraints.
+      #
+      # :constraints :: An identifier of the constraint, or an array of
+      #                 identifiers for constraints, to apply this
+      #                 change to specific constraints.
+      # :server :: The server/shard on which to run the query.
+      #
+      # Examples:
+      #
+      #   DB.immediate_constraints
+      #   # SET CONSTRAINTS ALL IMMEDIATE
+      #
+      #   DB.immediate_constraints(constraints: [:c1, Sequel[:sc][:c2]])
+      #   # SET CONSTRAINTS "c1", "sc"."s2" IMMEDIATE
+      def immediate_constraints(opts=OPTS)
+        _set_constraints(' IMMEDIATE', opts)
+      end
+
       # Use the pg_* system tables to determine indexes on a table
       def indexes(table, opts=OPTS)
         m = output_identifier_meth
@@ -1038,6 +1075,23 @@ def _schema_ds
         end
       end
 
+      # Internals of defer_constraints/immediate_constraints
+      def _set_constraints(type, opts)
+        execute_ddl(_set_constraints_sql(type, opts), opts)
+      end
+
+      # SQL to use for SET CONSTRAINTS
+      def _set_constraints_sql(type, opts)
+        sql = String.new
+        sql << "SET CONSTRAINTS "
+        if constraints = opts[:constraints]
+          dataset.send(:source_list_append, sql, Array(constraints))
+        else
+          sql << "ALL"
+        end
+        sql << type
+      end
+
       def alter_table_add_column_sql(table, op)
         "ADD COLUMN#{' IF NOT EXISTS' if op[:if_not_exists]} #{column_definition_sql(op)}"
       end

spec/adapters/postgres_spec.rb

@@ -479,6 +479,47 @@ def c.exec_prepared(*); super; nil end
     end
   end
 
+  it "should have #immediate_constraints and #defer_constraints for deferring/checking deferrable constraints" do
+    @db.create_table(:tmp_dolls) do
+      primary_key :id
+      foreign_key(:x, :tmp_dolls, :foreign_key_constraint_name=>:x_fk, :deferrable=>true)
+      foreign_key(:y, :tmp_dolls, :foreign_key_constraint_name=>:y_fk, :deferrable=>true)
+    end
+
+    ds = @db[:tmp_dolls]
+    @db.transaction do
+      @db.immediate_constraints
+      ds.insert(:id=>1)
+      @db.defer_constraints
+      ds.insert(:id=>2, :x=>1, :y=>3)
+      proc{@db.immediate_constraints}.must_raise Sequel::ForeignKeyConstraintViolation
+    end
+    @db[:tmp_dolls].must_be_empty
+    
+    @db.transaction do
+      @db.immediate_constraints
+      @db.defer_constraints(:constraints=>:y_fk)
+      ds.insert(:id=>1, :x=>1, :y=>2)
+      proc{@db.immediate_constraints}.must_raise Sequel::ForeignKeyConstraintViolation
+    end
+    @db[:tmp_dolls].must_be_empty
+
+    @db.transaction do
+      @db.defer_constraints
+      ds.insert(:id=>1, :x=>1, :y=>2)
+      proc{@db.immediate_constraints(:constraints=>:y_fk)}.must_raise Sequel::ForeignKeyConstraintViolation
+    end
+    @db[:tmp_dolls].must_be_empty
+
+    @db.transaction do
+      @db.immediate_constraints
+      @db.defer_constraints(:constraints=>[:x_fk, :y_fk])
+      ds.insert(:id=>1, :x=>3, :y=>2)
+      ds.update(:id=>1, :x=>1, :y=>1)
+    end
+    @db[:tmp_dolls].count.must_equal 1
+  end
+
   it "should have #check_constraints method for getting check constraints" do
     @db.create_table(:tmp_dolls) do
       Integer :i

spec/core/mock_adapter_spec.rb

@@ -145,7 +145,7 @@ def foo
     rs.must_equal [{:a=>1}] * 2
   end
 
-  it "should be able to set an exception to raise by setting the :fetch option to an exception class " do
+  it "should be able to set an exception to raise by setting the :fetch option to an exception class" do
     db = Sequel.mock(:fetch=>ArgumentError)
     proc{db[:t].all}.must_raise(Sequel::DatabaseError)
     begin
@@ -212,7 +212,7 @@ def foo
     db[:b].delete.must_equal 1
   end
 
-  it "should be able to set an exception to raise by setting the :numrows option to an exception class " do
+  it "should be able to set an exception to raise by setting the :numrows option to an exception class" do
     db = Sequel.mock(:numrows=>ArgumentError)
     proc{db[:t].update(:a=>1)}.must_raise(Sequel::DatabaseError)
     begin
@@ -280,7 +280,7 @@ def foo
     db[:b].insert(:a=>1).must_equal 1
   end
 
-  it "should be able to set an exception to raise by setting the :autoid option to an exception class " do
+  it "should be able to set an exception to raise by setting the :autoid option to an exception class" do
     db = Sequel.mock(:autoid=>ArgumentError)
     proc{db[:t].insert(:a=>1)}.must_raise(Sequel::DatabaseError)
     begin
@@ -869,6 +869,26 @@ def @db.schema(x) [[:id, {:primary_key=>false, :auto_increment=>false}]] end
   it "should recognize 40P01 SQL state as a serialization failure" do
     @db.send(:database_specific_error_class_from_sqlstate, '40P01').must_equal Sequel::SerializationFailure
   end
+
+  it "should use correct SQL for defer_constraints and immediate_constraints" do
+    @db.defer_constraints
+    @db.sqls.must_equal ['SET CONSTRAINTS ALL DEFERRED']
+    @db.immediate_constraints
+    @db.sqls.must_equal ['SET CONSTRAINTS ALL IMMEDIATE']
+
+    @db.defer_constraints(:constraints=>:a)
+    @db.sqls.must_equal ['SET CONSTRAINTS "a" DEFERRED']
+    @db.immediate_constraints(:constraints=>[:a, :b])
+    @db.sqls.must_equal ['SET CONSTRAINTS "a", "b" IMMEDIATE']
+  end
+
+  it "should correctly handle defer_constraints and immediate_constraints :server option" do
+    db = Sequel.connect("mock://postgres", :servers=>{:test=>{}})
+    db.defer_constraints(:server=>:test)
+    db.sqls.must_equal ['SET CONSTRAINTS ALL DEFERRED -- test']
+    db.immediate_constraints(:server=>:test)
+    db.sqls.must_equal ['SET CONSTRAINTS ALL IMMEDIATE -- test']
+  end
 end
 
 describe "MySQL support" do