Merge pull request rails#49811 from composerinteralia/verify-only-when-connection-exercised

Prevent marking broken connections as verified

Author: matthewd

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -977,7 +977,11 @@ def reconnect_can_restore_state?
         # If +allow_retry+ is true, a connection-related exception will
         # cause an automatic reconnect and re-run of the block, up to
         # the connection's configured +connection_retries+ setting
-        # and the configured +retry_deadline+ limit.
+        # and the configured +retry_deadline+ limit. (Note that when
+        # +allow_retry+ is true, it's possible to return without having marked
+        # the connection as verified. If the block is guaranteed to exercise the
+        # connection, consider calling `verified!` to avoid needless
+        # verification queries in subsequent calls.)
         #
         # If +materialize_transactions+ is false, the block will be run without
         # ensuring virtual transactions have been materialized in the DB
@@ -1028,9 +1032,7 @@ def with_raw_connection(allow_retry: false, materialize_transactions: true)
             end
 
             begin
-              result = yield @raw_connection
-              @verified = true
-              result
+              yield @raw_connection
             rescue => original_exception
               translated_exception = translate_exception_class(original_exception, nil, nil)
               invalidate_transaction(translated_exception)
@@ -1065,6 +1067,13 @@ def with_raw_connection(allow_retry: false, materialize_transactions: true)
           end
         end
 
+        # Mark the connection as verified. Call this inside a
+        # `with_raw_connection` block only when the block is guaranteed to
+        # exercise the raw connection.
+        def verified!
+          @verified = true
+        end
+
         def retryable_connection_error?(exception)
           exception.is_a?(ConnectionNotEstablished) || exception.is_a?(ConnectionFailed)
         end

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

@@ -102,6 +102,7 @@ def raw_execute(sql, name, async: false, allow_retry: false, materialize_transac
               with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn|
                 sync_timezone_changes(conn)
                 result = conn.query(sql)
+                verified!
                 handle_warnings(sql)
                 result
               end
@@ -130,6 +131,8 @@ def exec_stmt_and_free(sql, name, binds, cache_stmt: false, async: false)
                   result = ActiveSupport::Dependencies.interlock.permit_concurrent_loads do
                     stmt.execute(*type_casted_binds)
                   end
+                  verified!
+                  result
                 rescue ::Mysql2::Error => e
                   if cache_stmt
                     @statements.delete(sql)

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

@@ -16,7 +16,9 @@ def query(sql, name = nil) # :nodoc:
 
           log(sql, name) do
             with_raw_connection do |conn|
-              conn.async_exec(sql).map_types!(@type_map_for_results).values
+              result = conn.async_exec(sql).map_types!(@type_map_for_results).values
+              verified!
+              result
             end
           end
         end
@@ -51,6 +53,7 @@ def raw_execute(sql, name, async: false, allow_retry: false, materialize_transac
           log(sql, name, async: async) do
             with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn|
               result = conn.async_exec(sql)
+              verified!
               handle_warnings(result)
               result
             end

activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb

@@ -893,7 +893,9 @@ def exec_no_cache(sql, name, binds, async:, allow_retry:, materialize_transactio
           type_casted_binds = type_casted_binds(binds)
           log(sql, name, binds, type_casted_binds, async: async) do
             with_raw_connection do |conn|
-              conn.exec_params(sql, type_casted_binds)
+              result = conn.exec_params(sql, type_casted_binds)
+              verified!
+              result
             end
           end
         end
@@ -908,7 +910,9 @@ def exec_cache(sql, name, binds, async:, allow_retry:, materialize_transactions:
             type_casted_binds = type_casted_binds(binds)
 
             log(sql, name, binds, type_casted_binds, stmt_key, async: async) do
-              conn.exec_prepared(stmt_key, type_casted_binds)
+              result = conn.exec_prepared(stmt_key, type_casted_binds)
+              verified!
+              result
             end
           end
         rescue ActiveRecord::StatementInvalid => e

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

@@ -50,6 +50,7 @@ def internal_exec_query(sql, name = nil, binds = [], prepare: false, async: fals
                 stmt.bind_params(type_casted_binds)
                 records = stmt.to_a
               end
+              verified!
 
               build_result(columns: cols, rows: records)
             end
@@ -76,7 +77,9 @@ def begin_isolated_db_transaction(isolation) # :nodoc:
         def begin_db_transaction # :nodoc:
           log("begin transaction", "TRANSACTION") do
             with_raw_connection(allow_retry: true, materialize_transactions: false) do |conn|
-              conn.transaction
+              result = conn.transaction
+              verified!
+              result
             end
           end
         end
@@ -112,7 +115,9 @@ def high_precision_current_timestamp
           def raw_execute(sql, name, async: false, allow_retry: false, materialize_transactions: false)
             log(sql, name, async: async) do
               with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn|
-                conn.execute(sql)
+                result = conn.execute(sql)
+                verified!
+                result
               end
             end
           end
@@ -133,7 +138,9 @@ def execute_batch(statements, name = nil)
 
             log(sql, name) do
               with_raw_connection do |conn|
-                conn.execute_batch2(sql)
+                result = conn.execute_batch2(sql)
+                verified!
+                result
               end
             end
           end

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

@@ -47,6 +47,7 @@ def raw_execute(sql, name, async: false, allow_retry: false, materialize_transac
               with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn|
                 sync_timezone_changes(conn)
                 result = conn.query(sql)
+                verified!
                 handle_warnings(sql)
                 result
               end

activerecord/test/cases/adapter_test.rb

@@ -612,6 +612,26 @@ def teardown
         assert_predicate @connection, :active?
       end
 
+      test "quoting a string on a 'clean' failed connection will not prevent reconnecting" do
+        remote_disconnect @connection
+
+        @connection.clean! # this simulates a fresh checkout from the pool
+
+        # Clean did not verify / fix the connection
+        assert_not_predicate @connection, :active?
+
+        # Quote string will not verify a broken connection (although it may
+        # reconnect in some cases)
+        Post.connection.quote_string("")
+
+        # Because the connection hasn't been verified since checkout,
+        # and the query cannot safely be retried, the connection will be
+        # verified before querying.
+        Post.delete_all
+
+        assert_predicate @connection, :active?
+      end
+
       test "querying after a failed query restores and succeeds" do
         Post.first # Connection verified (and prepared statement pool populated if enabled)