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)
