Add transaction_connection_validator extension for retrying transacti…

…ons on new connection if there is a disconnect error when starting transaction

This is a less invasive approach than the connection_validator
extension, while still having mostly the same effect. It's limited
to recognizing disconnect errors when starting transactions, but
for many applications, those are the most important cases where
retrying is useful.

CHANGELOG

@@ -1,3 +1,7 @@
+=== master
+
+* Add transaction_connection_validator extension for retrying transactions on new connection if ther is a disconnect error when starting transaction (jeremyevans)
+
 === 5.76.0 (2024-01-01)
 
 * Improve performance and flexibility of regexp matching in sqlite adapter (paddor) (#2108)

doc/testing.rdoc

@@ -165,7 +165,7 @@ SEQUEL_ASYNC_THREAD_POOL_PREEMPT :: Use the async_thread_pool extension when run
 SEQUEL_CHECK_PENDING :: Try running all specs (note, can cause lockups for some adapters), and raise errors for skipped specs that don't fail
 SEQUEL_COLUMNS_INTROSPECTION :: Use the columns_introspection extension when running the specs
 SEQUEL_CONCURRENT_EAGER_LOADING :: Use the async_thread_pool extension and concurrent_eager_loading plugin when running the specs
-SEQUEL_CONNECTION_VALIDATOR :: Use the connection validator extension when running the specs
+SEQUEL_CONNECTION_VALIDATOR :: Use the connection_validator extension when running the adapter/integration specs
 SEQUEL_DUPLICATE_COLUMNS_HANDLER :: Use the duplicate columns handler extension with value given when running the specs
 SEQUEL_ERROR_SQL :: Use the error_sql extension when running the specs
 SEQUEL_FIBER_CONCURRENCY :: Use the fiber_concurrency extension when running the adapter and integration specs
@@ -186,4 +186,5 @@ SEQUEL_QUERY_PER_ASSOCIATION_DB_2_URL :: Run query-per-association integration t
 SEQUEL_QUERY_PER_ASSOCIATION_DB_3_URL :: Run query-per-association integration tests with multiple databases (all 4 must be set to run)
 SEQUEL_SPLIT_SYMBOLS :: Turn on symbol splitting when running the adapter and integration specs
 SEQUEL_SYNCHRONIZE_SQL :: Use the synchronize_sql extension when running the specs
+SEQUEL_TRANSACTION_CONNECTION_VALIDATOR :: Use the transaction_connection_validator extension when running the adapter/integration specs
 SEQUEL_TZINFO_VERSION :: Force the given tzinfo version when running the specs (e.g. '>=2')

lib/sequel/extensions/transaction_connection_validator.rb

@@ -0,0 +1,78 @@
+# frozen-string-literal: true
+#
+# The transaction_connection_validator extension automatically
+# retries a transaction on a connection if an disconnect error
+# is raised when sending the statement to begin a new
+# transaction, as long as the user has not already checked out
+# a connection.  This is safe to do because no other queries
+# have been issued on the connection, and no user-level code
+# is run before retrying.
+#
+# This approach to connection validation can be significantly
+# lower overhead than the connection_validator extension,
+# though it does not handle all cases handled by the
+# connection_validator extension.  However, it performs the
+# validation checks on every new transaction, so it will
+# automatically handle disconnected connections in some cases
+# where the connection_validator extension will not by default
+# (as the connection_validator extension only checks
+# connections if they have not been used in the last hour by
+# default).
+#
+# Related module: Sequel::TransactionConnectionValidator
+
+#
+module Sequel
+  module TransactionConnectionValidator
+    class DisconnectRetry < DatabaseDisconnectError
+      # The connection that raised the disconnect error
+      attr_accessor :connection
+
+      # The underlying disconnect error, in case it needs to be reraised.
+      attr_accessor :database_error
+    end
+
+    # Rescue disconnect errors raised when beginning a new transaction.  If there
+    # is a disconnnect error, it should be safe to retry the transaction using a
+    # new connection, as we haven't yielded control to the user yet.
+    def transaction(opts=OPTS)
+      super
+    rescue DisconnectRetry => e
+      if synchronize(opts[:server]){|conn| conn.equal?(e.connection)}
+        # If retrying would use the same connection, that means the
+        # connection was not removed from the pool, which means the caller has
+        # already checked out the connection, and retrying will not be successful.
+        # In this case, we can only reraise the exception.
+        raise e.database_error
+      end
+
+      num_retries ||= 0 
+      num_retries += 1
+      retry if num_retries < 5
+
+      raise e.database_error
+    end
+
+    private
+
+    # Reraise disconnect errors as DisconnectRetry so they can be retried.
+    def begin_new_transaction(conn, opts)
+      super
+    rescue Sequel::DatabaseDisconnectError, *database_error_classes => e
+      if e.is_a?(Sequel::DatabaseDisconnectError) || disconnect_error?(e, OPTS)
+        exception = DisconnectRetry.new(e.message)
+        exception.set_backtrace([])
+        exception.connection = conn
+        unless e.is_a?(Sequel::DatabaseError)
+          e = Sequel.convert_exception_class(e, database_error_class(e, OPTS))
+        end
+        exception.database_error = e
+        raise exception
+      end
+
+      raise
+    end
+  end
+
+  Database.register_extension(:transaction_connection_validator, TransactionConnectionValidator)
+end

spec/adapters/spec_helper.rb

@@ -40,6 +40,15 @@
 
 require_relative '../async_spec_helper'
 
+if ENV['SEQUEL_TRANSACTION_CONNECTION_VALIDATOR']
+  DB.extension(:transaction_connection_validator)
+end
+
+if ENV['SEQUEL_CONNECTION_VALIDATOR']
+  DB.extension(:connection_validator)
+  DB.pool.connection_validation_timeout = -1
+end
+
 DB.extension :pg_timestamptz if ENV['SEQUEL_PG_TIMESTAMPTZ']
 DB.extension :integer64 if ENV['SEQUEL_INTEGER64']
 DB.extension :index_caching if ENV['SEQUEL_INDEX_CACHING']

spec/extensions/transaction_connection_validator_spec.rb

@@ -0,0 +1,102 @@
+require_relative "spec_helper"
+
+describe "transaction_connection_validator extension" do
+  database_error = Class.new(StandardError)
+
+  before do
+    @db = Sequel.mock
+    @m = Module.new do
+      def post_execute(conn, sql); end
+      def disconnect_connection(conn)
+        @sqls << 'disconnect'
+      end
+      define_method(:database_error_classes) do
+        [database_error]
+      end
+      def disconnect_error?(e, opts)
+        e.message.include? 'disconnect'
+      end
+      def connect(server)
+        @sqls << 'connect'
+        super
+      end
+      def log_connection_execute(conn, sql)
+        res = super
+        post_execute(conn, sql)
+        res
+      end
+    end
+    @db.extend @m
+    @db.extension(:transaction_connection_validator)
+  end
+
+  it "should not affect transactions that do not raise exceptions" do
+    @db.transaction{}
+    @db.sqls.must_equal ['BEGIN', 'COMMIT']
+  end
+
+  it "should retry transactions for disconnects during BEGIN" do
+    conns = []
+    @db.define_singleton_method(:post_execute) do |conn, sql|
+      conns << conn
+      raise database_error, "disconnect error" if @sqls == ['BEGIN']
+    end
+    @db.transaction{}
+    @db.sqls.must_equal ['BEGIN', 'ROLLBACK', 'disconnect', 'connect', 'BEGIN', 'COMMIT']
+    conns.uniq.size.must_equal 2
+  end
+
+  it "should handle DatabaseDisconnectErrors as disconnects" do
+    conns = []
+    @db.define_singleton_method(:post_execute) do |conn, sql|
+      conns << conn
+      raise Sequel::DatabaseDisconnectError if @sqls == ['BEGIN']
+    end
+    @db.transaction{}
+    @db.sqls.must_equal ['BEGIN', 'ROLLBACK', 'disconnect', 'connect', 'BEGIN', 'COMMIT']
+    conns.uniq.size.must_equal 2
+  end
+
+  it "should not retry if a connection has already been checked out before calling transaction" do
+    conns = []
+    @db.define_singleton_method(:post_execute) do |conn, sql|
+      conns << conn
+      raise Sequel::DatabaseDisconnectError if @sqls == ['BEGIN']
+    end
+
+    c = nil
+    proc do
+      @db.synchronize do |c1|
+        c = c1
+        @db.transaction{}
+      end
+    end.must_raise(Sequel::DatabaseDisconnectError)
+    @db.sqls.must_equal ['BEGIN', 'ROLLBACK', 'disconnect']
+    conns.uniq.must_equal [c]
+  end
+
+  it "should not retry transaction setup more than 5 times" do
+    conns = []
+    @db.define_singleton_method(:post_execute) do |conn, sql|
+      conns << conn
+      raise database_error, "disconnect error"
+    end
+    proc do 
+      @db.transaction{}
+    end.must_raise(Sequel::DatabaseDisconnectError)
+    @db.sqls.must_equal(['BEGIN', 'ROLLBACK', 'disconnect', 'connect'] * 5)
+    conns.uniq.size.must_equal 5
+  end
+
+  it "should not retry on non-disconnect errors" do
+    conns = []
+    @db.define_singleton_method(:post_execute) do |conn, sql|
+      conns << conn
+      raise database_error, "normal error" if @sqls == ['BEGIN']
+    end
+
+    proc{@db.transaction{}}.must_raise(Sequel::DatabaseError)
+    @db.sqls.must_equal ['BEGIN', 'ROLLBACK']
+    conns.uniq.size.must_equal 1
+  end
+end

www/pages/plugins.html.erb

@@ -770,6 +770,10 @@
 <a class="a" href="rdoc-plugins/files/lib/sequel/extensions/sql_log_normalizer_rb.html">sql_log_normalizer </a>
 <span class="ul__span">Normalizes SQL before logging, helpful for analytics and sensitive data.</span>
 </li>
+<li class="ul__li ul__li--grid">
+<a class="a" href="rdoc-plugins/files/lib/sequel/extensions/transaction_connection_validator_rb.html">transaction_connection_validator </a>
+<span class="ul__span">Handle disconnect failures detected when starting a new transaction using a new connection transparently.</span>
+</li>
 </ul>
 
 <a name="extensions-dataset-sequel"></a>