def setUp(self):
self.factory_patcher = mock.patch(
'silverberg.thrift_client._ThriftClientFactory')
self.factory = self.factory_patcher.start()
def _create_factory(client_class, connection_lost):
self.connection_lost = connection_lost
self.factory.side_effect = _create_factory
self.addCleanup(self.factory_patcher.stop)
self.endpoint = mock.Mock()
self.twisted_transport = mock.Mock()
self.thrift_cassandra_client = mock.Mock(Client)
def _connect(factory):
self.connect_d = Deferred()
wrapper = mock.Mock()
wrapper.transport = self.twisted_transport
wrapper.wrapped.client = self.thrift_cassandra_client
return self.connect_d.addCallback(lambda _: wrapper)
self.endpoint.connect.side_effect = _connect
self.client = OnDemandThriftClient(self.endpoint, mock.Mock())
def __init__(self, cass_endpoint, keyspace, user=None, password=None,
disconnect_on_cancel=False):
self._client = OnDemandThriftClient(cass_endpoint, Cassandra.Client)
self._keyspace = keyspace
self._user = user
self._password = password
self._disconnect_on_cancel = disconnect_on_cancel
def setUp(self):
self.factory_patcher = mock.patch('silverberg.thrift_client._ThriftClientFactory')
self.factory = self.factory_patcher.start()
def _create_factory(client_class, connection_lost):
self.connection_lost = connection_lost
self.factory.side_effect = _create_factory
self.addCleanup(self.factory_patcher.stop)
self.endpoint = mock.Mock()
self.twisted_transport = mock.Mock()
self.thrift_cassandra_client = mock.Mock(Client)
def _connect(factory):
self.connect_d = Deferred()
wrapper = mock.Mock()
wrapper.transport = self.twisted_transport
wrapper.wrapped.client = self.thrift_cassandra_client
return self.connect_d.addCallback(lambda _: wrapper)
self.endpoint.connect.side_effect = _connect
self.client = OnDemandThriftClient(self.endpoint, mock.Mock())
def __init__(self, cass_endpoint, keyspace, user=None, password=None,
disconnect_on_cancel=False):
self._client = OnDemandThriftClient(cass_endpoint, Cassandra.Client)
self._keyspace = keyspace
self._user = user
self._password = password
self._disconnect_on_cancel = disconnect_on_cancel
def __init__(self, cass_endpoint, keyspace, user=None, password=None):
self._client = OnDemandThriftClient(cass_endpoint, Cassandra.Client)
self._keyspace = keyspace
self._user = user
self._password = password
class CQLClient(object):
"""
Cassandra CQL Client object.
Instantiate it and it will on-demand create a connection to the Cassandra
cluster.
:param cass_endpoint: A twisted Endpoint
:type cass_endpoint: twisted.internet.interfaces.IStreamClientEndpoint
:param keyspace: A keyspace to connect to
:type keyspace: str.
:param user: Username to connect with.
:type user: str.
:param password: Password to use in conjunction with Username.
:type password: str.
Upon connecting, the client will authenticate (if paramaters are provided)
and obtain the keyspace definition so that it can de-serialize properly.
n.b. Cassandra presently doesn't have any real support for password
authentication in the mainline as the simple access control options
are disabled; you probably need to secure your Cassandra server using
different methods and the password code isn't heavily tested.
"""
def __init__(self, cass_endpoint, keyspace, user=None, password=None):
self._client = OnDemandThriftClient(cass_endpoint, Cassandra.Client)
self._keyspace = keyspace
self._user = user
self._password = password
def _set_keyspace(self, client):
d = client.set_keyspace(self._keyspace)
return d.addCallback(lambda _: client)
def _login(self, client):
creds = {'user': self._user, 'password': self._password}
authreq = ttypes.AuthenticationRequest(creds)
d = client.login(authreq)
d.addCallback(lambda _: client)
return d
def _connection(self):
def _handshake(client):
d = succeed(client)
if self._user and self._password:
d.addCallback(self._login)
d.addCallback(self._set_keyspace)
return d
ds = self._client.connection(_handshake)
return ds
def disconnect(self):
"""
Disconnect from the cassandra cluster. Cassandara and Silverberg do
not require the connection to be closed before exiting. However, this
method may be useful if resources are constrained, or for testing
purposes if using or injecting :class:`TestCQLClient: is impossible.
:return: a :class:`Deferred` that fires with None when disconnected.
"""
return self._client.disconnect()
def describe_version(self):
"""
Query the Cassandra server for the version.
:returns: string -- the version tag
"""
def _vers(client):
return client.describe_version()
d = self._connection()
d.addCallback(_vers)
return d
def _unmarshal_result(self, schema, raw_rows, _unmarshallers):
rows = []
def _unmarshal_val(vtype, val):
if val is None:
return
# Differentiate between regular and collection types
# Collection types look like 'ListType(SomeCassandraType)',
# so try split into two parts and check if we can marshal them
types = str(vtype).replace(")", "").split("(")
# Regular type
if len(types) == 1 and vtype in _unmarshallers:
return _unmarshallers[vtype](val)
# Collection
elif len(types) == 2 and types[0] in _unmarshallers and types[1] in _unmarshallers:
return _unmarshallers[types[0]](types[1], val)
# XXX: We do not currently implement the full range of types.
# So we can not unmarshal all types in which case we should just
# return the raw bytes.
return val
for raw_row in raw_rows:
row = {}
for raw_col in raw_row.columns:
specific = schema.value_types[raw_col.name]
row[raw_col.name] = _unmarshal_val(specific, raw_col.value)
rows.append(row)
return rows
def execute(self, query, args, consistency):
"""
Execute a CQL query against the server.
:param query: The CQL query to execute
:type query: str.
:param args: The arguments to substitute
:type args: dict.
:param consistency: The consistency level
:type consistency: ConsistencyLevel
In order to avoid unpleasant issues of CQL injection
(Hey, just because there's no SQL doesn't mean that Little
Bobby Tables won't mess things up for you like in XKCD #327)
you probably want to use argument substitution instead of
concatting strings together to build a query.
Thus, like the official CQL driver for non-Twisted python
that comes with the Cassandra distro, we do variable substitution.
Example::
d = client.execute("UPDATE :table SET 'fff' = :val WHERE "
"KEY = :key",{"val":1234, "key": "fff", "table": "blah"})
:returns: A Deferred that fires with either None, an int, or an
iterable of `{'column': value, ...}` dictionaries, depending
on the CQL query. e.g. a UPDATE would return None,
whereas a SELECT would return an int or some rows
Example output::
[{"fff": 1222}]
"""
prep_query = prepare(query, args)
def _execute(client):
return client.execute_cql3_query(prep_query,
ttypes.Compression.NONE, consistency)
def _proc_results(result):
if result.type == ttypes.CqlResultType.ROWS:
return self._unmarshal_result(result.schema, result.rows,
unmarshallers)
elif result.type == ttypes.CqlResultType.INT:
return result.num
else:
return None
d = self._connection()
d.addCallback(_execute)
d.addCallback(_proc_results)
return d
class OnDemandThriftClientTests(BaseTestCase):
def setUp(self):
self.factory_patcher = mock.patch('silverberg.thrift_client._ThriftClientFactory')
self.factory = self.factory_patcher.start()
def _create_factory(client_class, connection_lost):
self.connection_lost = connection_lost
self.factory.side_effect = _create_factory
self.addCleanup(self.factory_patcher.stop)
self.endpoint = mock.Mock()
self.twisted_transport = mock.Mock()
self.thrift_cassandra_client = mock.Mock(Client)
def _connect(factory):
self.connect_d = Deferred()
wrapper = mock.Mock()
wrapper.transport = self.twisted_transport
wrapper.wrapped.client = self.thrift_cassandra_client
return self.connect_d.addCallback(lambda _: wrapper)
self.endpoint.connect.side_effect = _connect
self.client = OnDemandThriftClient(self.endpoint, mock.Mock())
def test_initial_connect(self):
d = self.client.connection()
self.connect_d.callback(None)
self.assertEqual(self.assertFired(d), self.thrift_cassandra_client)
def test_initial_handshake_non_deferred(self):
def _mock_handshake(client):
return None
m = mock.MagicMock(side_effect=_mock_handshake)
d = self.client.connection(m)
self.connect_d.callback(None)
self.assertEqual(self.assertFired(d), self.thrift_cassandra_client)
m.assert_called_once()
def test_initial_handshake_deferred(self):
def _mock_handshake_d(client):
return succeed(None)
m = mock.MagicMock(side_effect=_mock_handshake_d)
d = self.client.connection(m)
self.connect_d.callback(None)
self.assertEqual(self.assertFired(d), self.thrift_cassandra_client)
m.assert_called_once()
def test_connected(self):
d1 = self.client.connection()
self.connect_d.callback(None)
d2 = self.client.connection()
self.assertNotIdentical(d1, d2)
self.assertEqual(self.assertFired(d2), self.thrift_cassandra_client)
def test_connect_while_connecting(self):
d1 = self.client.connection()
d2 = self.client.connection()
self.assertNotIdentical(d1, d2)
self.connect_d.callback(None)
self.assertEqual(self.assertFired(d1), self.thrift_cassandra_client)
self.assertEqual(self.assertFired(d2), self.thrift_cassandra_client)
def test_connect_failed(self):
d = self.client.connection()
self.connect_d.errback(_TestConnectError())
self.assertFailed(d, _TestConnectError)
def test_connection_lost_cleanly(self):
d = self.client.connection()
self.connect_d.callback(None)
self.assertFired(d)
self.connection_lost(Failure(_TestConnectionDone()))
@mock.patch('silverberg.thrift_client.log')
def test_connection_lost_uncleanly(self, mock_log):
self.twisted_transport.getPeer.return_value = 'peer addr'
d = self.client.connection()
self.connect_d.callback(None)
self.assertFired(d)
reason = Failure(_TestConnectionLost())
self.connection_lost(reason)
mock_log.err.assert_called_once_with(
reason, "Lost current connection to 'peer addr', reconnecting on demand.",
system='OnDemandThriftClient', node='peer addr')
def test_disconnect(self):
d1 = self.client.connection()
self.connect_d.callback(None)
self.assertFired(d1)
d2 = self.client.disconnect()
self.twisted_transport.loseConnection.assert_called_once()
self.connection_lost(Failure(_TestConnectionDone()))
self.assertFired(d2)
def test_connect_while_disconnecting(self):
d1 = self.client.connection()
self.connect_d.callback(None)
self.assertFired(d1)
d2 = self.client.disconnect()
d3 = self.client.connection()
self.connection_lost(Failure(_TestConnectionDone()))
self.assertFired(d2)
self.assertFailed(d3, ClientDisconnecting)
def test_disconnect_while_connecting(self):
d1 = self.client.connection()
d2 = self.client.disconnect()
self.assertFailed(d2, ClientConnecting)
self.assertEqual(
len(self.twisted_transport.loseConnection.mock_calls), 0,
"loseConnection should not be called since not connected")
self.connect_d.callback(None)
self.assertFired(d1)
def test_disconnect_while_disconnecting(self):
self.client.connection()
self.connect_d.callback(None)
d1 = self.client.disconnect()
self.twisted_transport.reset_mock()
d2 = self.client.disconnect()
self.connection_lost(Failure(_TestConnectionDone()))
self.assertEqual(
len(self.twisted_transport.loseConnection.mock_calls), 0,
"loseConnection should not be called since not connected")
self.assertFired(d1)
self.assertFired(d2)
def test_disconnect_while_not_connected(self):
d1 = self.client.disconnect()
self.assertIdentical(self.assertFired(d1), None)
self.assertEqual(
len(self.twisted_transport.loseConnection.mock_calls), 0,
"loseConnection should not be called since not connected")
class OnDemandThriftClientTests(BaseTestCase):
def setUp(self):
self.factory_patcher = mock.patch(
'silverberg.thrift_client._ThriftClientFactory')
self.factory = self.factory_patcher.start()
def _create_factory(client_class, connection_lost):
self.connection_lost = connection_lost
self.factory.side_effect = _create_factory
self.addCleanup(self.factory_patcher.stop)
self.endpoint = mock.Mock()
self.twisted_transport = mock.Mock()
self.thrift_cassandra_client = mock.Mock(Client)
def _connect(factory):
self.connect_d = Deferred()
wrapper = mock.Mock()
wrapper.transport = self.twisted_transport
wrapper.wrapped.client = self.thrift_cassandra_client
return self.connect_d.addCallback(lambda _: wrapper)
self.endpoint.connect.side_effect = _connect
self.client = OnDemandThriftClient(self.endpoint, mock.Mock())
def test_initial_connect(self):
d = self.client.connection()
self.connect_d.callback(None)
self.assertEqual(self.assertFired(d), self.thrift_cassandra_client)
def test_initial_handshake_non_deferred(self):
def _mock_handshake(client):
return None
m = mock.MagicMock(side_effect=_mock_handshake)
d = self.client.connection(m)
self.connect_d.callback(None)
self.assertEqual(self.assertFired(d), self.thrift_cassandra_client)
m.assert_called_once()
def test_initial_handshake_deferred(self):
def _mock_handshake_d(client):
return succeed(None)
m = mock.MagicMock(side_effect=_mock_handshake_d)
d = self.client.connection(m)
self.connect_d.callback(None)
self.assertEqual(self.assertFired(d), self.thrift_cassandra_client)
m.assert_called_once()
def test_connected(self):
d1 = self.client.connection()
self.connect_d.callback(None)
d2 = self.client.connection()
self.assertNotIdentical(d1, d2)
self.assertEqual(self.assertFired(d2), self.thrift_cassandra_client)
def test_connect_while_connecting(self):
d1 = self.client.connection()
d2 = self.client.connection()
self.assertNotIdentical(d1, d2)
self.connect_d.callback(None)
self.assertEqual(self.assertFired(d1), self.thrift_cassandra_client)
self.assertEqual(self.assertFired(d2), self.thrift_cassandra_client)
def test_connect_failed(self):
d = self.client.connection()
self.connect_d.errback(_TestConnectError())
self.assertFailed(d, _TestConnectError)
def test_connection_lost_cleanly(self):
d = self.client.connection()
self.connect_d.callback(None)
self.assertFired(d)
self.connection_lost(Failure(_TestConnectionDone()))
@mock.patch('silverberg.thrift_client.log')
def test_connection_lost_uncleanly(self, mock_log):
self.twisted_transport.getPeer.return_value = 'peer addr'
d = self.client.connection()
self.connect_d.callback(None)
self.assertFired(d)
reason = Failure(_TestConnectionLost())
self.connection_lost(reason)
mock_log.err.assert_called_once_with(
reason,
"Lost current connection to 'peer addr', reconnecting on demand.",
system='OnDemandThriftClient',
node='peer addr')
def test_disconnect(self):
d1 = self.client.connection()
self.connect_d.callback(None)
self.assertFired(d1)
d2 = self.client.disconnect()
self.twisted_transport.loseConnection.assert_called_once()
self.connection_lost(Failure(_TestConnectionDone()))
self.assertFired(d2)
def test_connect_while_disconnecting(self):
d1 = self.client.connection()
self.connect_d.callback(None)
self.assertFired(d1)
d2 = self.client.disconnect()
d3 = self.client.connection()
self.connection_lost(Failure(_TestConnectionDone()))
self.assertFired(d2)
self.assertFailed(d3, ClientDisconnecting)
def test_disconnect_while_connecting(self):
d1 = self.client.connection()
d2 = self.client.disconnect()
self.assertFailed(d2, ClientConnecting)
self.assertEqual(
len(self.twisted_transport.loseConnection.mock_calls), 0,
"loseConnection should not be called since not connected")
self.connect_d.callback(None)
self.assertFired(d1)
def test_disconnect_while_disconnecting(self):
self.client.connection()
self.connect_d.callback(None)
d1 = self.client.disconnect()
self.twisted_transport.reset_mock()
d2 = self.client.disconnect()
self.connection_lost(Failure(_TestConnectionDone()))
self.assertEqual(
len(self.twisted_transport.loseConnection.mock_calls), 0,
"loseConnection should not be called since not connected")
self.assertFired(d1)
self.assertFired(d2)
def test_disconnect_while_not_connected(self):
d1 = self.client.disconnect()
self.assertIdentical(self.assertFired(d1), None)
self.assertEqual(
len(self.twisted_transport.loseConnection.mock_calls), 0,
"loseConnection should not be called since not connected")
