Allow open transactions to be restored upon DB reconnect

If we know that any open transactions had not been queried (and
raw_connection has not been used), then it's safe for us to pick up with
a new database connection mid-request.

Author: matthewd

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

@@ -330,8 +330,24 @@ def transaction_open?
         current_transaction.open?
       end
 
-      def reset_transaction # :nodoc:
+      def reset_transaction(restore: false) # :nodoc:
+        # Store the existing transaction state to the side
+        old_state = @transaction_manager if restore && @transaction_manager&.restorable?
+
         @transaction_manager = ConnectionAdapters::TransactionManager.new(self)
+
+        if block_given?
+          # Reconfigure the connection without any transaction state in the way
+          result = yield
+
+          # Now the connection's fully established, we can swap back
+          if old_state
+            @transaction_manager = old_state
+            @transaction_manager.restore_transactions
+          end
+
+          result
+        end
       end
 
       # Register a record with the current transaction so that its after_commit and after_rollback callbacks

activerecord/lib/active_record/connection_adapters/abstract/transaction.rb

@@ -140,6 +140,10 @@ def materialized?
         @materialized
       end
 
+      def restore!
+        @materialized = false
+      end
+
       def rollback_records
         return unless records
         ite = records.uniq(&:__id__)
@@ -348,6 +352,20 @@ def dirty_current_transaction
         current_transaction.dirty!
       end
 
+      def restore_transactions
+        return false unless restorable?
+
+        @stack.each(&:restore!)
+
+        materialize_transactions unless @lazy_transactions_enabled
+
+        true
+      end
+
+      def restorable?
+        @stack.none?(&:dirty?)
+      end
+
       def materialize_transactions
         return if @materializing_transactions
 

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -113,6 +113,8 @@ def initialize(connection, logger = nil, config = {}) # :nodoc:
 
         @default_timezone = self.class.validate_default_timezone(config[:default_timezone])
 
+        @raw_connection_dirty = false
+
         configure_connection
       end
 
@@ -555,10 +557,11 @@ def active?
       # new connection with the database. Implementors should call super
       # immediately after establishing the new connection (and while still
       # holding @lock).
-      def reconnect!
-        clear_cache!(new_connection: true)
-        reset_transaction
-        configure_connection
+      def reconnect!(restore_transactions: false)
+        reset_transaction(restore: restore_transactions) do
+          clear_cache!(new_connection: true)
+          configure_connection
+        end
       end
 
       # Disconnects from the database if already connected. Otherwise, this
@@ -637,6 +640,7 @@ def verify!
       # PostgreSQL's lo_* methods.
       def raw_connection
         disable_lazy_transactions!
+        @raw_connection_dirty = true
         @raw_connection
       end
 
@@ -789,6 +793,10 @@ def extract_limit(sql_type)
       EXTENDED_TYPE_MAPS = Concurrent::Map.new
 
       private
+        def reconnect_can_restore_state?
+          transaction_manager.restorable? && !@raw_connection_dirty
+        end
+
         def extended_type_map_key
           if @default_timezone
             { default_timezone: @default_timezone }

activerecord/lib/active_record/connection_adapters/mysql2_adapter.rb

@@ -123,9 +123,9 @@ def active?
         @raw_connection.ping
       end
 
-      def reconnect!
+      def reconnect!(restore_transactions: false)
         @lock.synchronize do
-          disconnect!
+          @raw_connection.close
           connect
           super
         end

activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb

@@ -322,7 +322,7 @@ def reload_type_map # :nodoc:
       end
 
       # Close then reopen the connection.
-      def reconnect!
+      def reconnect!(restore_transactions: false)
         @lock.synchronize do
           begin
             @raw_connection.reset

activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb

@@ -163,7 +163,7 @@ def active?
         !@raw_connection.closed?
       end
 
-      def reconnect!
+      def reconnect!(restore_transactions: false)
         @lock.synchronize do
           if active?
             @raw_connection.rollback rescue nil

activerecord/test/cases/adapter_test.rb

@@ -394,6 +394,23 @@ def setup
         assert_not raw_transaction_open?(@connection)
       end
 
+      test "materialized transaction state can be restored after a reconnect" do
+        @connection.begin_transaction
+        assert_predicate @connection, :transaction_open?
+        # +materialize_transactions+ currently automatically dirties the
+        # connection, which would make it unrestorable
+        @connection.transaction_manager.stub(:dirty_current_transaction, nil) do
+          @connection.materialize_transactions
+        end
+        assert raw_transaction_open?(@connection)
+        @connection.reconnect!(restore_transactions: true)
+        assert_predicate @connection, :transaction_open?
+        assert_not raw_transaction_open?(@connection)
+      ensure
+        @connection.reconnect!
+        assert_not_predicate @connection, :transaction_open?
+      end
+
       test "materialized transaction state is reset after a disconnect" do
         @connection.begin_transaction
         assert_predicate @connection, :transaction_open?
@@ -417,6 +434,21 @@ def setup
         assert_not raw_transaction_open?(@connection)
       end
 
+      test "unmaterialized transaction state can be restored after a reconnect" do
+        @connection.begin_transaction
+        assert_predicate @connection, :transaction_open?
+        assert_not raw_transaction_open?(@connection)
+        @connection.reconnect!(restore_transactions: true)
+        assert_predicate @connection, :transaction_open?
+        assert_not raw_transaction_open?(@connection)
+        @connection.materialize_transactions
+        assert raw_transaction_open?(@connection)
+      ensure
+        @connection.reconnect!
+        assert_not_predicate @connection, :transaction_open?
+        assert_not raw_transaction_open?(@connection)
+      end
+
       test "unmaterialized transaction state is reset after a disconnect" do
         @connection.begin_transaction
         assert_predicate @connection, :transaction_open?