Fix ResponseFuture._set_result crashes on connection error when used with PrepareMessage

Author: aboudreault

CHANGELOG.rst

@@ -11,7 +11,8 @@ Features
 Bug Fixes
 ---------
 * re-raising the CQLEngineException will fail on Python 3 (PYTHON-1166)
-* Connection fails to validate ssl certificate hostname when SSLContext.check_hostname is set
+* Connection fails to validate ssl certificate hostname when SSLContext.check_hostname is set (PYTHON-1186)
+* ResponseFuture._set_result crashes on connection error when used with PrepareMessage (PYTHON-1187)
 
 3.20.2
 ======

cassandra/cluster.py

@@ -4059,8 +4059,9 @@ def _set_result(self, host, connection, pool, response):
                     log.warning("Host %s error: %s.", host, response.summary)
                     if self._metrics is not None:
                         self._metrics.on_other_error()
+                    cl = getattr(self.message, 'consistency_level', None)
                     retry = retry_policy.on_request_error(
-                        self.query, self.message.consistency_level, error=response,
+                        self.query, cl, error=response,
                         retry_num=self._query_retries)
                 elif isinstance(response, PreparedQueryNotFound):
                     if self.prepared_statement:
@@ -4117,9 +4118,9 @@ def _set_result(self, host, connection, pool, response):
                     self._metrics.on_connection_error()
                 if not isinstance(response, ConnectionShutdown):
                     self._connection.defunct(response)
+                cl = getattr(self.message, 'consistency_level', None)
                 retry = self._retry_policy.on_request_error(
-                    self.query, self.message.consistency_level, error=response,
-                    retry_num=self._query_retries)
+                    self.query, cl, error=response, retry_num=self._query_retries)
                 self._handle_retry_decision(retry, response, host)
             elif isinstance(response, Exception):
                 if hasattr(response, 'to_exception'):

tests/unit/test_response_future.py

@@ -200,6 +200,28 @@ def test_unavailable_error_message(self):
         rf._set_result(None, None, None, result)
         self.assertRaises(Exception, rf.result)
 
+    def test_request_error_with_prepare_message(self):
+        session = self.make_session()
+        query = SimpleStatement("SELECT * FROM foobar")
+        retry_policy = Mock()
+        retry_policy.on_request_error.return_value = (RetryPolicy.RETHROW, None)
+        message = PrepareMessage(query=query)
+
+        rf = ResponseFuture(session, message, query, 1, retry_policy=retry_policy)
+        rf._query_retries = 1
+        rf.send_request()
+        result = Mock(spec=OverloadedErrorMessage)
+        result.to_exception.return_value = result
+        rf._set_result(None, None, None, result)
+        self.assertIsInstance(rf._final_exception, OverloadedErrorMessage)
+
+        rf = ResponseFuture(session, message, query, 1, retry_policy=retry_policy)
+        rf._query_retries = 1
+        rf.send_request()
+        result = Mock(spec=ConnectionException)
+        rf._set_result(None, None, None, result)
+        self.assertIsInstance(rf._final_exception, ConnectionException)
+
     def test_retry_policy_says_ignore(self):
         session = self.make_session()
         query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)")