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.

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