cluster: report a timeout waiting for schema agreement, not a defunct connection

cassandra/cluster.py

@@ -7403,8 +7403,14 @@ def _on_timeout(self, _attempts=0):
             # We should still raise OperationTimedOut to reject the future so that the main event thread will not
             # wait for it endlessly
             except KeyError:
-                key = "Connection defunct by heartbeat"
-                errors = {key: "Client request timeout. See Session.execute[_async](timeout)"}
+                if not self.is_schema_agreed:
+                    # The request was also removed when its response arrived:
+                    # a schema change result, which now waits for schema
+                    # agreement. Report that instead of a defunct connection.
+                    errors = self._schema_agreement_timeout_errors()
+                else:
+                    key = "Connection defunct by heartbeat"
+                    errors = {key: "Client request timeout. See Session.execute[_async](timeout)"}
                 self._set_final_exception(OperationTimedOut(errors, self._current_host,
                                                             timeout=self.timeout,
                                                             in_flight=connection.in_flight))
@@ -7436,14 +7442,17 @@ def _on_timeout(self, _attempts=0):
                     key = str(self._current_host.endpoint)
                 errors = {key: "Client request timeout. See Session.execute[_async](timeout)"}
             else:
-                connection = self.session.cluster.control_connection._connection
-                host = str(connection.endpoint) if connection else 'unknown'
-                errors = {host: "Request timed out while waiting for schema agreement. See Session.execute[_async](timeout) and Cluster.max_schema_agreement_wait."}
+                errors = self._schema_agreement_timeout_errors()
 
         self._set_final_exception(OperationTimedOut(errors, self._current_host,
                                                     timeout=self.timeout,
                                                     in_flight=conn_in_flight))
 
+    def _schema_agreement_timeout_errors(self):
+        connection = self.session.cluster.control_connection._connection
+        host = str(connection.endpoint) if connection else 'unknown'
+        return {host: "Request timed out while waiting for schema agreement. See Session.execute[_async](timeout) and Cluster.max_schema_agreement_wait."}
+
     def _on_speculative_execute(self):
         self._timer = None
         if not self._event.is_set():

tests/unit/test_response_future.py

@@ -749,6 +749,55 @@ def test_heartbeat_defunct_deadlock(self):
         assert exc_info.value.timeout == 1
         assert exc_info.value.in_flight == 5
 
+    def test_timeout_while_waiting_for_schema_agreement(self):
+        """
+        A schema change result removes the request from the connection's request
+        queue, like a defunct connection does, and then waits for schema agreement.
+        A timeout during that wait must be reported as a timeout while waiting for
+        schema agreement - the server did apply the schema change - not as a
+        connection defunct by heartbeat.
+        """
+
+        connection = MagicMock(spec=Connection)
+        connection._requests = {}
+        connection.in_flight = 5
+        connection.orphaned_request_ids = set()
+        connection.features.tablets_routing_v2 = False
+        connection.features.tablets_routing_v1 = False
+
+        pool = Mock()
+        pool.is_shutdown = False
+        pool.borrow_connection.return_value = [connection, 1]
+
+        session = self.make_basic_session()
+        session.cluster._default_load_balancing_policy.make_query_plan.return_value = [Mock(), Mock()]
+        session._pools.get.return_value = pool
+
+        query = SimpleStatement("DROP KEYSPACE ks")
+        message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE)
+
+        rf = ResponseFuture(session, message, query, 1)
+        rf.send_request()
+
+        # Simulate the arrival of the schema change result: the connection has
+        # removed the request from its queue, and the result waits for schema
+        # agreement in session.submit(), which this mock session never runs.
+        connection._requests = {}
+        result = Mock(spec=ResultMessage,
+                      kind=RESULT_KIND_SCHEMA_CHANGE,
+                      schema_change_event={'target_type': SchemaTargetType.KEYSPACE,
+                                           'change_type': SchemaChangeType.DROPPED,
+                                           'keyspace': "ks"})
+        rf._set_result(None, connection, pool, result)
+        session.submit.assert_called_once()
+
+        # Simulate ResponseFuture timing out
+        rf._on_timeout()
+        with pytest.raises(OperationTimedOut, match="waiting for schema agreement") as exc_info:
+            rf.result()
+        assert "Connection defunct by heartbeat" not in str(exc_info.value)
+        assert exc_info.value.timeout == 1
+
     def test_read_timeout_error_message(self):
         session = self.make_session()
         query = SimpleStatement("SELECT * FROM foo")