Merge pull request rails#52529 from Shopify/mysql-begin-isolation-retry

Preserve isolation level when retrying MySQL transactions

Author: byroot

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

@@ -102,16 +102,16 @@ def select_rows(arel, name = nil, binds = [], async: false)
         select_all(arel, name, binds, async: async).then(&:rows)
       end
 
-      def query_value(sql, name = nil) # :nodoc:
-        single_value_from_rows(query(sql, name))
+      def query_value(...) # :nodoc:
+        single_value_from_rows(query(...))
       end
 
-      def query_values(sql, name = nil) # :nodoc:
-        query(sql, name).map(&:first)
+      def query_values(...) # :nodoc:
+        query(...).map(&:first)
       end
 
-      def query(sql, name = nil) # :nodoc:
-        internal_exec_query(sql, name).rows
+      def query(...) # :nodoc:
+        internal_exec_query(...).rows
       end
 
       # Determines whether the SQL statement is a write query.
@@ -591,9 +591,9 @@ def internal_execute(sql, name = "SQL", binds = [], prepare: false, async: false
           raw_execute(sql, name, binds, prepare: prepare, async: async, allow_retry: allow_retry, materialize_transactions: materialize_transactions, &block)
         end
 
-        def execute_batch(statements, name = nil)
+        def execute_batch(statements, name = nil, **kwargs)
           statements.each do |statement|
-            raw_execute(statement, name)
+            raw_execute(statement, name, **kwargs)
           end
         end
 

activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb

@@ -227,8 +227,12 @@ def begin_db_transaction # :nodoc:
       def begin_isolated_db_transaction(isolation) # :nodoc:
         # From MySQL manual: The [SET TRANSACTION] statement applies only to the next single transaction performed within the session.
         # So we don't need to implement #reset_isolation_level
-        internal_execute("SET TRANSACTION ISOLATION LEVEL #{transaction_isolation_levels.fetch(isolation)}", "TRANSACTION", allow_retry: true, materialize_transactions: false)
-        begin_db_transaction
+        execute_batch(
+          ["SET TRANSACTION ISOLATION LEVEL #{transaction_isolation_levels.fetch(isolation)}", "BEGIN"],
+          "TRANSACTION",
+          allow_retry: true,
+          materialize_transactions: false,
+        )
       end
 
       def commit_db_transaction # :nodoc:
@@ -565,7 +569,7 @@ def table_options(table_name) # :nodoc:
 
       # SHOW VARIABLES LIKE 'name'
       def show_variable(name)
-        query_value("SELECT @@#{name}", "SCHEMA")
+        query_value("SELECT @@#{name}", "SCHEMA", materialize_transactions: false, allow_retry: true)
       rescue ActiveRecord::StatementInvalid
         nil
       end

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

@@ -14,9 +14,9 @@ def select_all(*, **) # :nodoc:
         end
 
         private
-          def execute_batch(statements, name = nil)
+          def execute_batch(statements, name = nil, **kwargs)
             combine_multi_statements(statements).each do |statement|
-              raw_execute(statement, name, batch: true)
+              raw_execute(statement, name, batch: true, **kwargs)
             end
           end
 

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

@@ -192,8 +192,8 @@ def affected_rows(result)
             affected_rows
           end
 
-          def execute_batch(statements, name = nil)
-            raw_execute(combine_multi_statements(statements), name, batch: true)
+          def execute_batch(statements, name = nil, **kwargs)
+            raw_execute(combine_multi_statements(statements), name, batch: true, **kwargs)
           end
 
           def build_truncate_statements(table_names)

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

@@ -123,9 +123,9 @@ def affected_rows(result)
             @last_affected_rows
           end
 
-          def execute_batch(statements, name = nil)
+          def execute_batch(statements, name = nil, **kwargs)
             sql = combine_multi_statements(statements)
-            raw_execute(sql, name, batch: true)
+            raw_execute(sql, name, batch: true, **kwargs)
           end
 
           def build_truncate_statement(table_name)

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

@@ -58,9 +58,9 @@ def last_inserted_id(result)
             end
           end
 
-          def execute_batch(statements, name = nil)
+          def execute_batch(statements, name = nil, **kwargs)
             combine_multi_statements(statements).each do |statement|
-              raw_execute(statement, name, batch: true)
+              raw_execute(statement, name, batch: true, **kwargs)
             end
           end
       end

activerecord/test/cases/adapters/abstract_mysql_adapter/transaction_test.rb

@@ -149,5 +149,56 @@ class Sample < ActiveRecord::Base
       end
       assert_kind_of ActiveRecord::QueryAborted, error
     end
+
+    test "reconnect preserves isolation level" do
+      pool = Sample.connection_pool
+      @connection = Sample.lease_connection
+
+      Sample.transaction do
+        @connection.materialize_transactions
+        # Double check that the INSERT isn't seen with default isolation level
+        assert_no_difference -> { Sample.count } do
+          Thread.new do
+            Sample.create!(value: 1)
+          end.join
+        end
+      end
+
+      Sample.transaction(isolation: :read_committed) do
+        @connection.materialize_transactions
+        # Double check that the INSERT is seen with :read_committed
+        assert_difference -> { Sample.count }, +1 do
+          Thread.new do
+            Sample.create!(value: 1)
+          end.join
+        end
+      end
+
+      first_begin_failed = false
+      @connection.singleton_class.define_method(:perform_query) do |raw_connection, sql, *args, **kwargs|
+        # Simulates the first BEGIN attempt failing
+        if sql.include?("BEGIN") && !first_begin_failed
+          first_begin_failed = true
+          raise ActiveRecord::ConnectionFailed, "Simulated failure"
+        end
+        super(raw_connection, sql, *args, **kwargs)
+      end
+
+      Sample.transaction(isolation: :read_committed) do
+        @connection.materialize_transactions
+        # Indirectly check that the retried transaction is READ COMMITTED
+        assert_difference -> { Sample.count }, +1 do
+          Thread.new do
+            Sample.create!(value: 1)
+          end.join
+        end
+      end
+
+      assert first_begin_failed
+    ensure
+      pool.remove(@connection) # We're monkey patching the instance, we shouldn't re-use it
+      @connection&.disconnect!
+      Sample.release_connection
+    end
   end
 end