def test_multiple_connections(self):
"""
Test multiple connections with pipelined requests.
"""
conns = [self.get_connection() 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),
request_id=count,
cb=partial(cb, event, conn, count))
for event, conn in zip(events, conns):
conn.send_msg(
QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE),
request_id=0,
cb=partial(cb, event, conn, 0))
for event in events:
event.wait()
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)
with conn.lock:
request_id = conn.get_request_id()
conn.send_msg(qmsg, request_id, cb=partial(cb, conn, event, thread_responses, i))
event.wait()
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)
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)
connection.send_msg.assert_called_with(
rf.message,
2,
cb=ANY,
encoder=ProtocolHandler.encode_message,
decoder=ProtocolHandler.decode_message,
result_metadata=[])
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),
request_id=count,
cb=partial(cb, event, conn, count))
def send_msgs(all_responses, thread_responses):
for i in range(num_requests_per_conn):
qmsg = QueryMessage(query=query,
consistency_level=ConsistencyLevel.ONE)
with conn.lock:
request_id = conn.get_request_id()
conn.send_msg(qmsg,
request_id,
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`.
This method will always increment :attr:`.in_flight` attribute, even if
it doesn't need to make a request, just to maintain an
":attr:`.in_flight` is incremented" invariant.
"""
# Here we increment in_flight unconditionally, whether we need to issue
# a request or not. This is bad, but allows callers -- specifically
# _set_keyspace_for_all_conns -- to assume that we increment
# self.in_flight during this call. This allows the passed callback to
# safely call HostConnection{Pool,}.return_connection on this
# Connection.
#
# We use a busy wait on the lock here because:
# - we'll only spin if the connection is at max capacity, which is very
# unlikely for a set_keyspace call
# - it allows us to avoid signaling a condition every time a request completes
while True:
with self.lock:
if self.in_flight < self.max_request_id:
self.in_flight += 1
break
time.sleep(0.001)
if not keyspace or keyspace == self.keyspace:
callback(self, None)
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)))
# We've incremented self.in_flight above, so we "have permission" to
# acquire a new request id
request_id = self.get_request_id()
self.send_msg(query, request_id, process_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.send_request()
result = Mock(spec=UnavailableErrorMessage, info={})
rf._set_result(None, None, None, result)
self.assertRaises(Exception, rf.result)
def test_continuous_paging(self):
"""
Test to check continuous paging throws an Exception if it's not supported and the correct valuesa
are written to the buffer if the option is enabled.
@since DSE 2.0b3 GRAPH 1.0b1
@jira_ticket PYTHON-694
@expected_result the values are correctly written
@test_category connection
"""
max_pages = 4
max_pages_per_second = 3
continuous_paging_options = ContinuousPagingOptions(
max_pages=max_pages, max_pages_per_second=max_pages_per_second)
message = QueryMessage(
"a", 3, continuous_paging_options=continuous_paging_options)
io = Mock()
for version in [
version for version in ProtocolVersion.SUPPORTED_VERSIONS
if version != ProtocolVersion.DSE_V1
]:
self.assertRaises(UnsupportedOperation, message.send_body, io,
version)
io.reset_mock()
message.send_body(io, ProtocolVersion.DSE_V1)
# continuous paging adds two write calls to the buffer
self.assertEqual(len(io.write.mock_calls), 6)
# Check that the appropriate flag is set to True
self.assertEqual(
uint32_unpack(io.write.mock_calls[3][1][0])
& _WITH_SERIAL_CONSISTENCY_FLAG, 0)
self.assertEqual(
uint32_unpack(io.write.mock_calls[3][1][0]) & _PAGE_SIZE_FLAG, 0)
self.assertEqual(
uint32_unpack(io.write.mock_calls[3][1][0])
& _WITH_PAGING_STATE_FLAG, 0)
self.assertEqual(
uint32_unpack(io.write.mock_calls[3][1][0]) & _PAGING_OPTIONS_FLAG,
_PAGING_OPTIONS_FLAG)
# Test max_pages and max_pages_per_second are correctly written
self.assertEqual(uint32_unpack(io.write.mock_calls[4][1][0]),
max_pages)
self.assertEqual(uint32_unpack(io.write.mock_calls[5][1][0]),
max_pages_per_second)
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_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_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(None, None, None, result)
self.assertRaises(Exception, rf.result)
# test callback
rf = ResponseFuture(session, message, query, 1)
rf.send_request()
callback = Mock()
expected_result = (object(), object())
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[0], expected_result[1]))
self.assertEqual(rf.result()[0], expected_result)
callback.assert_called_once_with([expected_result], arg, **kwargs)
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:
callback(self, None)
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)))
request_id = None
# we use a busy wait on the lock here because:
# - we'll only spin if the connection is at max capacity, which is very
# unlikely for a set_keyspace call
# - it allows us to avoid signaling a condition every time a request completes
while True:
with self.lock:
if self.in_flight < self.max_request_id:
request_id = self.get_request_id()
self.in_flight += 1
break
time.sleep(0.001)
self.send_msg(query, request_id, process_result)
def test_query_message(self):
"""
Test to check the appropriate calls are made
@since 3.9
@jira_ticket PYTHON-713
@expected_result the values are correctly written
@test_category connection
"""
message = QueryMessage("a", 3)
io = Mock()
message.send_body(io, 4)
self._check_calls(io, [(b'\x00\x00\x00\x01', ), (b'a', ),
(b'\x00\x03', ), (b'\x00', )])
io.reset_mock()
message.send_body(io, 5)
self._check_calls(io, [(b'\x00\x00\x00\x01', ), (b'a', ),
(b'\x00\x03', ), (b'\x00\x00\x00\x00', )])
def test_single_connection_pipelined_requests(self):
"""
Test a single connection with pipelined requests.
"""
conn = self.get_connection()
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),
request_id=i,
cb=partial(cb, responses, i))
event.wait()
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)")
message = QueryMessage(query=query,
consistency_level=ConsistencyLevel.ONE)
retry_policy = Mock()
retry_policy.on_unavailable.return_value = (RetryPolicy.RETHROW, None)
rf = ResponseFuture(session,
message,
query,
1,
retry_policy=retry_policy)
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(None, None, None, 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 set_keyspace_blocking(self, keyspace):
if not keyspace or keyspace == self.keyspace:
return
query = QueryMessage(query='USE "%s"' % (keyspace, ),
consistency_level=ConsistencyLevel.ONE)
try:
result = self.wait_for_response(query)
except InvalidRequestException as ire:
# the keyspace probably doesn't exist
raise ire.to_exception()
except Exception as exc:
conn_exc = ConnectionException(
"Problem while setting keyspace: %r" % (exc, ), self.host)
self.defunct(conn_exc)
raise conn_exc
if isinstance(result, ResultMessage):
self.keyspace = keyspace
else:
conn_exc = ConnectionException(
"Problem while setting keyspace: %r" % (result, ), self.host)
self.defunct(conn_exc)
raise conn_exc
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(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_single_connection(self):
"""
Test a single connection with sequential requests.
"""
conn = self.get_connection()
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),
request_id=0,
cb=partial(cb, count))
conn.send_msg(
QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE),
request_id=0,
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, 1)
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