def test_multiple_errbacks(self): session = self.make_session() pool = session._pools.get.return_value connection = Mock(spec=Connection) pool.borrow_connection.return_value = (connection, 1) query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") query.retry_policy = Mock() query.retry_policy.on_unavailable.return_value = (RetryPolicy.RETHROW, None) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query) rf.send_request() callback = Mock() arg = "positional" kwargs = {'one': 1, 'two': 2} rf.add_errback(callback, arg, **kwargs) callback2 = Mock() arg2 = "another" kwargs2 = {'three': 3, 'four': 4} rf.add_errback(callback2, arg2, **kwargs2) expected_exception = Unavailable("message", 1, 2, 3) result = Mock(spec=UnavailableErrorMessage, info={'something': 'here'}) result.to_exception.return_value = expected_exception rf._set_result(result) self.assertRaises(Exception, rf.result) callback.assert_called_once_with(expected_exception, arg, **kwargs) callback2.assert_called_once_with(expected_exception, arg2, **kwargs2)
def test_all_pools_shutdown(self): session = self.make_basic_session() session._load_balancer.make_query_plan.return_value = ['ip1', 'ip2'] session._pools.get.return_value.is_shutdown = True rf = ResponseFuture(session, Mock(), Mock()) rf.send_request() self.assertRaises(NoHostAvailable, rf.result)
def test_multiple_errbacks(self): session = self.make_session() pool = session._pools.get.return_value connection = Mock(spec=Connection) pool.borrow_connection.return_value = (connection, 1) query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") query.retry_policy = Mock() query.retry_policy.on_unavailable.return_value = (RetryPolicy.RETHROW, None) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1) rf.send_request() callback = Mock() arg = "positional" kwargs = {'one': 1, 'two': 2} rf.add_errback(callback, arg, **kwargs) callback2 = Mock() arg2 = "another" kwargs2 = {'three': 3, 'four': 4} rf.add_errback(callback2, arg2, **kwargs2) expected_exception = Unavailable("message", 1, 2, 3) result = Mock(spec=UnavailableErrorMessage, info={'something': 'here'}) result.to_exception.return_value = expected_exception rf._set_result(result) self.assertRaises(Exception, rf.result) callback.assert_called_once_with(expected_exception, arg, **kwargs) callback2.assert_called_once_with(expected_exception, arg2, **kwargs2)
def test_write_timeout_error_message(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1) rf.send_request() result = Mock(spec=WriteTimeoutErrorMessage, info={"write_type": 1, "required_responses":2, "received_responses":1, "consistency": 1}) rf._set_result(None, None, None, result) self.assertRaises(Exception, rf.result)
def test_unavailable_error_message(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") query.retry_policy = Mock() query.retry_policy.on_unavailable.return_value = (RetryPolicy.RETHROW, None) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query) rf.send_request() result = Mock(spec=UnavailableErrorMessage, info={}) rf._set_result(result) self.assertRaises(Exception, rf.result)
def test_read_timeout_error_message(self): session = self.make_session() query = SimpleStatement("SELECT * FROM foo") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1) rf.send_request() result = Mock(spec=ReadTimeoutErrorMessage, info={"data_retrieved": "", "required_responses":2, "received_responses":1, "consistency": 1}) rf._set_result(None, None, None, result) self.assertRaises(Exception, rf.result)
def test_read_timeout_error_message(self): session = self.make_session() query = SimpleStatement("SELECT * FROM foo") query.retry_policy = Mock() query.retry_policy.on_read_timeout.return_value = (RetryPolicy.RETHROW, None) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query) rf.send_request() result = Mock(spec=ReadTimeoutErrorMessage, info={}) rf._set_result(result) self.assertRaises(Exception, rf.result)
def test_retry_policy_says_retry(self): session = self.make_session() pool = session._pools.get.return_value query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") query.retry_policy = Mock() query.retry_policy.on_unavailable.return_value = (RetryPolicy.RETRY, ConsistencyLevel.ONE) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.QUORUM) rf = ResponseFuture(session, message, query) rf.send_request() rf.session._pools.get.assert_called_once_with('ip1') pool.borrow_connection.assert_called_once_with(timeout=ANY) connection = pool.borrow_connection.return_value connection.send_msg.assert_called_once_with(rf.message, cb=ANY) result = Mock(spec=UnavailableErrorMessage, info={}) rf._set_result(result) session.submit.assert_called_once_with(rf._retry_task, True) self.assertEqual(1, rf._query_retries) # simulate the executor running this rf._retry_task(True) # it should try again with the same host since this was # an UnavailableException rf.session._pools.get.assert_called_with('ip1') pool.borrow_connection.assert_called_with(timeout=ANY) connection = pool.borrow_connection.return_value connection.send_msg.assert_called_with(rf.message, cb=ANY)
def test_retry_policy_says_retry(self): session = self.make_session() pool = session._pools.get.return_value query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.QUORUM) connection = Mock(spec=Connection) pool.borrow_connection.return_value = (connection, 1) retry_policy = Mock() retry_policy.on_unavailable.return_value = (RetryPolicy.RETRY, ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1, retry_policy=retry_policy) rf.send_request() rf.session._pools.get.assert_called_once_with('ip1') pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY) connection.send_msg.assert_called_once_with( rf.message, 1, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[]) result = Mock(spec=UnavailableErrorMessage, info={}) host = Mock() rf._set_result(host, None, None, result) session.submit.assert_called_once_with(rf._retry_task, True, host) self.assertEqual(1, rf._query_retries) connection = Mock(spec=Connection) pool.borrow_connection.return_value = (connection, 2) # simulate the executor running this rf._retry_task(True, host) # it should try again with the same host since this was # an UnavailableException rf.session._pools.get.assert_called_with(host) pool.borrow_connection.assert_called_with(timeout=ANY, routing_key=ANY) connection.send_msg.assert_called_with( rf.message, 2, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[])
def test_errback(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") query.retry_policy = Mock() query.retry_policy.on_unavailable.return_value = (RetryPolicy.RETHROW, None) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query) rf.send_request() rf.add_errback(self.assertIsInstance, Exception) result = Mock(spec=UnavailableErrorMessage, info={}) rf._set_result(result) self.assertRaises(Exception, rf.result) # this should get called immediately now that the error is set rf.add_errback(self.assertIsInstance, Exception)
def _wrap_future(f: ResponseFuture) -> asyncio.Future: """Wrap a cassandra Future into an asyncio.Future object. Args: f: future to wrap Returns: And asyncio.Future object which can be awaited. """ loop = asyncio.get_event_loop() aio_future = loop.create_future() def on_result(result): loop.call_soon_threadsafe(aio_future.set_result, result) def on_error(exception, *_): loop.call_soon_threadsafe(aio_future.set_exception, exception) f.add_callback(on_result) f.add_errback(on_error) return aio_future
def test_errback(self): session = self.make_session() pool = session._pools.get.return_value connection = Mock(spec=Connection) pool.borrow_connection.return_value = (connection, 1) query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1) rf._query_retries = 1 rf.send_request() rf.add_errback(self.assertIsInstance, Exception) result = Mock(spec=UnavailableErrorMessage, info={"required_replicas":2, "alive_replicas": 1, "consistency": 1}) result.to_exception.return_value = Exception() rf._set_result(None, None, None, result) self.assertRaises(Exception, rf.result) # this should get called immediately now that the error is set rf.add_errback(self.assertIsInstance, Exception)
def test_retry_policy_says_ignore(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") query.retry_policy = Mock() query.retry_policy.on_unavailable.return_value = (RetryPolicy.IGNORE, None) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query) rf.send_request() result = Mock(spec=UnavailableErrorMessage, info={}) rf._set_result(result) self.assertEqual(None, rf.result())
def wrap_future( response_future: ResponseFuture, callback_fn: Callable[..., None], callback_args: Any, errback_fn: Callable[..., None], errback_args: Any, ) -> ResponseFuture: """Patch ResponseFuture.result() to wait for callback or errback to complete. The callback_fn and errback_fn are given special treatment: they must be complete before a result will be returned from ResponseFuture.result(). They are not given precedence over other callbacks or errbacks, so if another callback triggers the response from the service (and the server span is closed) the special callback might not complete. The special callback is added first and callbacks are executed in order, so generally the special callback should finish before any other callbacks. This fixes a race condition where the server span can complete before the callback has closed out the child span. """ response_future._callback_event = Event() response_future.add_callback(callback_fn, callback_args, response_future._callback_event) response_future.add_errback(errback_fn, errback_args, response_future._callback_event) def wait_for_callbacks_result(self: ResponseFuture) -> Any: exc = None try: result = ResponseFuture.result(self) except Exception as e: exc = e # wait for either _on_execute_complete or _on_execute_failed to run wait_result = self._callback_event.wait(timeout=0.01) if not wait_result: logger.warning( "Cassandra metrics callback took too long. Some metrics may be lost." ) if exc: raise exc # pylint: disable=E0702 return result # call __get__ to turn wait_for_callbacks_result into a bound method bound_method = wait_for_callbacks_result.__get__( # type: ignore response_future, ResponseFuture) # patch the ResponseFuture instance response_future.result = bound_method return response_future
def test_errback(self): session = self.make_session() pool = session._pools.get.return_value connection = Mock(spec=Connection) pool.borrow_connection.return_value = (connection, 1) query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") query.retry_policy = Mock() query.retry_policy.on_unavailable.return_value = (RetryPolicy.RETHROW, None) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1) rf.send_request() rf.add_errback(self.assertIsInstance, Exception) result = Mock(spec=UnavailableErrorMessage, info={}) rf._set_result(result) self.assertRaises(Exception, rf.result) # this should get called immediately now that the error is set rf.add_errback(self.assertIsInstance, Exception)
def wait_for_callbacks_result(self: ResponseFuture) -> Any: exc = None try: result = ResponseFuture.result(self) except Exception as e: exc = e # wait for either _on_execute_complete or _on_execute_failed to run wait_result = self._callback_event.wait(timeout=0.01) if not wait_result: logger.warning("Cassandra metrics callback took too long. Some metrics may be lost.") if exc: raise exc # pylint: disable=E0702 return result
def test_retry_policy_says_ignore(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) retry_policy = Mock() retry_policy.on_unavailable.return_value = (RetryPolicy.IGNORE, None) rf = ResponseFuture(session, message, query, 1, retry_policy=retry_policy) rf.send_request() result = Mock(spec=UnavailableErrorMessage, info={}) rf._set_result(None, None, None, result) self.assertFalse(rf.result())
def test_write_timeout_error_message(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") query.retry_policy = Mock() query.retry_policy.on_write_timeout.return_value = (RetryPolicy.RETHROW, None) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query) rf.send_request() result = Mock(spec=WriteTimeoutErrorMessage, info={}) rf._set_result(result) self.assertRaises(Exception, rf.result)
def test_read_timeout_error_message(self): session = self.make_session() query = SimpleStatement("SELECT * FROM foo") query.retry_policy = Mock() query.retry_policy.on_read_timeout.return_value = (RetryPolicy.RETHROW, None) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1) rf.send_request() result = Mock(spec=ReadTimeoutErrorMessage, info={}) rf._set_result(None, None, None, result) self.assertRaises(Exception, rf.result)
def test_unavailable_error_message(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") query.retry_policy = Mock() query.retry_policy.on_unavailable.return_value = (RetryPolicy.RETHROW, None) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1) rf._query_retries = 1 rf.send_request() result = Mock(spec=UnavailableErrorMessage, info={ "required_replicas": 2, "alive_replicas": 1, "consistency": 1 }) rf._set_result(None, None, None, result) self.assertRaises(Exception, rf.result)
def get_async_result(self, future: ResponseFuture) -> ResultSet: """ Get the result of asynchronous CQL statement. *Args:*\n _future_ - object, returned as a result of keyword [#Execute Async Cql | Execute Async Cql] *Returns:*\n Result of asynchronous CQL statement. *Example:*\n | ${obj}= | Execute Async Cql | SELECT * FROM system.schema_keyspaces; | | Sleep | 5 | | ${result}= | Get Async Result | ${obj} | | Log | ${result[1].keyspace_name} | =>\n system """ try: result = future.result() except Exception as e: raise Exception('Operation failed: {0}'.format(e)) return result
def test_write_timeout_error_message(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1) rf.send_request() result = Mock(spec=WriteTimeoutErrorMessage, info={ "write_type": 1, "required_responses": 2, "received_responses": 1, "consistency": 1 }) rf._set_result(None, None, None, result) self.assertRaises(Exception, rf.result)
def test_read_timeout_error_message(self): session = self.make_session() query = SimpleStatement("SELECT * FROM foo") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1) rf.send_request() result = Mock(spec=ReadTimeoutErrorMessage, info={ "data_retrieved": "", "required_responses": 2, "received_responses": 1, "consistency": 1 }) rf._set_result(None, None, None, result) self.assertRaises(Exception, rf.result)
def test_heartbeat_defunct_deadlock(self): """ Heartbeat defuncts all connections and clears request queues. Response future times out and even if it has been removed from request queue, timeout exception must be thrown. Otherwise event loop will deadlock on eventual ResponseFuture.result() call. PYTHON-1044 """ connection = MagicMock(spec=Connection) connection._requests = {} 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("SELECT * FROM foo") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1) rf.send_request() # Simulate Connection.error_all_requests() after heartbeat defuncts connection._requests = {} # Simulate ResponseFuture timing out rf._on_timeout() self.assertRaisesRegexp(OperationTimedOut, "Connection defunct by heartbeat", rf.result)
def test_add_callbacks(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") query.retry_policy = Mock() query.retry_policy.on_unavailable.return_value = (RetryPolicy.RETHROW, None) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) # test errback rf = ResponseFuture(session, message, query) rf.send_request() rf.add_callbacks( callback=self.assertEquals, callback_args=([{'col': 'val'}],), errback=self.assertIsInstance, errback_args=(Exception,)) result = Mock(spec=UnavailableErrorMessage, info={}) rf._set_result(result) self.assertRaises(Exception, rf.result) # test callback rf = ResponseFuture(session, message, query) rf.send_request() rf.add_callbacks( callback=self.assertEquals, callback_args=([{'col': 'val'}],), errback=self.assertIsInstance, errback_args=(Exception,)) response = Mock(spec=ResultMessage, kind=ResultMessage.KIND_ROWS, results=[{'col': 'val'}]) rf._set_result(response) self.assertEqual(rf.result(), [{'col': 'val'}])
def test_errback(self): session = self.make_session() pool = session._pools.get.return_value connection = Mock(spec=Connection) pool.borrow_connection.return_value = (connection, 1) query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) rf = ResponseFuture(session, message, query, 1) rf._query_retries = 1 rf.send_request() rf.add_errback(self.assertIsInstance, Exception) result = Mock(spec=UnavailableErrorMessage, info={ "required_replicas": 2, "alive_replicas": 1, "consistency": 1 }) result.to_exception.return_value = Exception() rf._set_result(None, None, None, result) self.assertRaises(Exception, rf.result) # this should get called immediately now that the error is set rf.add_errback(self.assertIsInstance, Exception)
def _send_query_message(self, session, timeout, **kwargs): query = "SELECT * FROM test3rf.test" message = QueryMessage(query=query, **kwargs) future = ResponseFuture(session, message, query=None, timeout=timeout) future.send_request() return future
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_add_callbacks(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") query.retry_policy = Mock() query.retry_policy.on_unavailable.return_value = (RetryPolicy.RETHROW, None) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) # test errback rf = ResponseFuture(session, message, query) rf.send_request() rf.add_callbacks(callback=self.assertEquals, callback_args=([{ 'col': 'val' }], ), errback=self.assertIsInstance, errback_args=(Exception, )) result = Mock(spec=UnavailableErrorMessage, info={}) rf._set_result(result) self.assertRaises(Exception, rf.result) # test callback rf = ResponseFuture(session, message, query) rf.send_request() rf.add_callbacks(callback=self.assertEquals, callback_args=([{ 'col': 'val' }], ), errback=self.assertIsInstance, errback_args=(Exception, )) response = Mock(spec=ResultMessage, kind=ResultMessage.KIND_ROWS, results=[{ 'col': 'val' }]) rf._set_result(response) self.assertEqual(rf.result(), [{'col': 'val'}])
def test_add_callbacks(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") query.retry_policy = Mock() query.retry_policy.on_unavailable.return_value = (RetryPolicy.RETHROW, None) message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) # test errback rf = ResponseFuture(session, message, query, 1) rf.send_request() rf.add_callbacks( callback=self.assertEqual, callback_args=([{'col': 'val'}],), errback=self.assertIsInstance, errback_args=(Exception,)) result = Mock(spec=UnavailableErrorMessage, info={}) rf._set_result(result) self.assertRaises(Exception, rf.result) # test callback rf = ResponseFuture(session, message, query, 1) rf.send_request() callback = Mock() expected_result = [{'col': 'val'}] arg = "positional" kwargs = {'one': 1, 'two': 2} rf.add_callbacks( callback=callback, callback_args=(arg,), callback_kwargs=kwargs, errback=self.assertIsInstance, errback_args=(Exception,)) rf._set_result(self.make_mock_response(expected_result)) self.assertEqual(rf.result(), expected_result) callback.assert_called_once_with(expected_result, arg, **kwargs)
def make_response_future(self, session): query = SimpleStatement("SELECT * FROM foo") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) return ResponseFuture(session, message, query, 1)
def test_add_callbacks(self): session = self.make_session() query = SimpleStatement("INSERT INFO foo (a, b) VALUES (1, 2)") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) # test errback rf = ResponseFuture(session, message, query, 1) rf._query_retries = 1 rf.send_request() rf.add_callbacks(callback=self.assertEqual, callback_args=([{ 'col': 'val' }], ), errback=self.assertIsInstance, errback_args=(Exception, )) result = Mock(spec=UnavailableErrorMessage, info={ "required_replicas": 2, "alive_replicas": 1, "consistency": 1 }) result.to_exception.return_value = Exception() rf._set_result(None, None, None, result) self.assertRaises(Exception, rf.result) # test callback rf = ResponseFuture(session, message, query, 1) rf.send_request() callback = Mock() expected_result = [{'col': 'val'}] arg = "positional" kwargs = {'one': 1, 'two': 2} rf.add_callbacks(callback=callback, callback_args=(arg, ), callback_kwargs=kwargs, errback=self.assertIsInstance, errback_args=(Exception, )) rf._set_result(None, None, None, self.make_mock_response(expected_result)) self.assertEqual(rf.result(), expected_result) callback.assert_called_once_with(expected_result, arg, **kwargs)