def send_msgs(conn, event): thread_responses = [False] * num_requests_per_conn for i in range(num_requests_per_conn): qmsg = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) conn.send_msg(qmsg, cb=partial(cb, conn, event, thread_responses, i)) event.wait()
def test_multiple_connections(self): """ Test multiple connections with pipelined requests. """ conns = [self.klass.factory() for i in range(5)] events = [Event() for i in range(5)] query = "SELECT keyspace_name FROM system.schema_keyspaces LIMIT 1" def cb(event, conn, count, *args, **kwargs): count += 1 if count >= 10: conn.close() event.set() else: conn.send_msg( QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE), cb=partial(cb, event, conn, count)) for event, conn in zip(events, conns): conn.send_msg( QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE), cb=partial(cb, event, conn, 0)) for event in events: event.wait()
def send_msgs(all_responses, thread_responses): for i in range(num_requests_per_conn): qmsg = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) conn.send_msg(qmsg, cb=partial(cb, all_responses, thread_responses, i))
def set_keyspace_async(self, keyspace, callback): """ Use this in order to avoid deadlocking the event loop thread. When the operation completes, `callback` will be called with two arguments: this connection and an Exception if an error occurred, otherwise :const:`None`. """ if not keyspace or keyspace == self.keyspace: return query = QueryMessage(query='USE "%s"' % (keyspace, ), consistency_level=ConsistencyLevel.ONE) def process_result(result): if isinstance(result, ResultMessage): self.keyspace = keyspace callback(self, None) elif isinstance(result, InvalidRequestException): callback(self, result.to_exception()) else: callback( self, self.defunct( ConnectionException( "Problem while setting keyspace: %r" % (result, ), self.host))) self.send_msg(query, process_result, wait_for_id=True)
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 set_keyspace(self, keyspace): if not keyspace or keyspace == self.keyspace: return with self.lock: query = 'USE "%s"' % (keyspace, ) try: result = self.wait_for_response( QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE)) if isinstance(result, ResultMessage): self.keyspace = keyspace else: raise self.defunct( ConnectionException( "Problem while setting keyspace: %r" % (result, ), self.host)) except InvalidRequestException as ire: # the keyspace probably doesn't exist raise ire.to_exception() except Exception as exc: raise self.defunct( ConnectionException( "Problem while setting keyspace: %r" % (exc, ), self.host))
def cb(event, conn, count, *args, **kwargs): count += 1 if count >= 10: conn.close() event.set() else: conn.send_msg( QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE), cb=partial(cb, event, conn, count))
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 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_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() 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 test_single_connection_pipelined_requests(self): """ Test a single connection with pipelined requests. """ conn = self.klass.factory() query = "SELECT keyspace_name FROM system.schema_keyspaces LIMIT 1" responses = [False] * 100 event = Event() def cb(response_list, request_num, *args, **kwargs): response_list[request_num] = True if all(response_list): conn.close() event.set() for i in range(100): conn.send_msg( QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE), cb=partial(cb, responses, i)) event.wait()
def test_single_connection(self): """ Test a single connection with sequential requests. """ conn = self.klass.factory() query = "SELECT keyspace_name FROM system.schema_keyspaces LIMIT 1" event = Event() def cb(count, *args, **kwargs): count += 1 if count >= 10: conn.close() event.set() else: conn.send_msg( QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE), cb=partial(cb, count)) conn.send_msg( QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE), cb=partial(cb, 0)) event.wait()
def make_response_future(self, session): query = SimpleStatement("SELECT * FROM foo") message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE) return ResponseFuture(session, message, query)