def test_wrap_round_robin(self):
cluster = Mock(spec=Cluster)
cluster.metadata = Mock(spec=Metadata)
hosts = [Host(str(i), SimpleConvictionPolicy) for i in range(4)]
for host in hosts:
host.set_up()
def get_replicas(keyspace, packed_key):
index = struct.unpack('>i', packed_key)[0]
return list(islice(cycle(hosts), index, index + 2))
cluster.metadata.get_replicas.side_effect = get_replicas
policy = TokenAwarePolicy(RoundRobinPolicy())
policy.populate(cluster, hosts)
for i in range(4):
query = Statement(routing_key=struct.pack('>i', i), keyspace='keyspace_name')
qplan = list(policy.make_query_plan(None, query))
replicas = get_replicas(None, struct.pack('>i', i))
other = set(h for h in hosts if h not in replicas)
self.assertEqual(replicas, qplan[:2])
self.assertEqual(other, set(qplan[2:]))
# Should use the secondary policy
for i in range(4):
qplan = list(policy.make_query_plan())
self.assertEqual(set(qplan), set(hosts))
def _assert_shuffle(self, patched_shuffle, keyspace, routing_key):
hosts = [Host(str(i), SimpleConvictionPolicy) for i in range(4)]
for host in hosts:
host.set_up()
cluster = Mock(spec=Cluster)
cluster.metadata = Mock(spec=Metadata)
replicas = hosts[2:]
cluster.metadata.get_replicas.return_value = replicas
child_policy = Mock()
child_policy.make_query_plan.return_value = hosts
child_policy.distance.return_value = HostDistance.LOCAL
policy = TokenAwarePolicy(child_policy, shuffle_replicas=True)
policy.populate(cluster, hosts)
cluster.metadata.get_replicas.reset_mock()
child_policy.make_query_plan.reset_mock()
query = Statement(routing_key=routing_key)
qplan = list(policy.make_query_plan(keyspace, query))
if keyspace is None or routing_key is None:
self.assertEqual(hosts, qplan)
self.assertEqual(cluster.metadata.get_replicas.call_count, 0)
child_policy.make_query_plan.assert_called_once_with(keyspace, query)
self.assertEqual(patched_shuffle.call_count, 0)
else:
self.assertEqual(set(replicas), set(qplan[:2]))
self.assertEqual(hosts[:2], qplan[2:])
child_policy.make_query_plan.assert_called_once_with(keyspace, query)
self.assertEqual(patched_shuffle.call_count, 1)
def test_rfthree_tokenaware_none_down(self):
keyspace = 'test_rfthree_tokenaware_none_down'
cluster = TestCluster(
execution_profiles={
EXEC_PROFILE_DEFAULT:
ExecutionProfile(TokenAwarePolicy(RoundRobinPolicy()))
})
session = cluster.connect(wait_for_all_pools=True)
wait_for_up(cluster, 1)
wait_for_up(cluster, 2)
create_schema(cluster, session, keyspace, replication_factor=3)
self._insert(session, keyspace, count=1)
self._query(session, keyspace, count=1)
self.coordinator_stats.assert_query_count_equals(self, 1, 0)
self.coordinator_stats.assert_query_count_equals(self, 2, 1)
self.coordinator_stats.assert_query_count_equals(self, 3, 0)
self.coordinator_stats.reset_counts()
self._assert_writes_succeed(session, keyspace,
SINGLE_DC_CONSISTENCY_LEVELS)
self._assert_reads_succeed(session,
keyspace,
SINGLE_DC_CONSISTENCY_LEVELS -
set([ConsistencyLevel.ANY]),
expected_reader=2)
cluster.shutdown()
def get_cursor(existing=None, hosts=None, new=False):
global session
if existing:
return existing
if session and not new:
return session
if hosts is None:
hosts = os.environ.get('KEYVALUE_HOSTS', '').split(',')
auth = None
try:
with open('keys/cassandra-config.json') as f:
auth = PlainTextAuthProvider(**json.load(f))
except Exception as e:
log.warning('Could not get Cassandra credentials from file.')
log.warning(e)
cluster = Cluster(
hosts,
load_balancing_policy=TokenAwarePolicy(RoundRobinPolicy()),
auth_provider=auth)
session = cluster.connect('default')
session.row_factory = named_tuple_factory
return session
def test_rfthree_tokenaware_downgradingcl(self):
keyspace = 'test_rfthree_tokenaware_downgradingcl'
cluster = Cluster(
load_balancing_policy=TokenAwarePolicy(RoundRobinPolicy()),
default_retry_policy=DowngradingConsistencyRetryPolicy(),
protocol_version=PROTOCOL_VERSION)
self.rfthree_downgradingcl(cluster, keyspace, False)
def test_rfthree_tokenaware_none_down(self):
keyspace = 'test_rfthree_tokenaware_none_down'
cluster = Cluster(load_balancing_policy=TokenAwarePolicy(
RoundRobinPolicy()),
protocol_version=PROTOCOL_VERSION)
session = cluster.connect(wait_for_all_pools=True)
wait_for_up(cluster, 1)
wait_for_up(cluster, 2)
create_schema(cluster, session, keyspace, replication_factor=3)
self._insert(session, keyspace, count=1)
self._query(session, keyspace, count=1)
self.coordinator_stats.assert_query_count_equals(self, 1, 0)
self.coordinator_stats.assert_query_count_equals(self, 2, 1)
self.coordinator_stats.assert_query_count_equals(self, 3, 0)
self.coordinator_stats.reset_counts()
self._assert_writes_succeed(session, keyspace,
SINGLE_DC_CONSISTENCY_LEVELS)
self._assert_reads_succeed(session,
keyspace,
SINGLE_DC_CONSISTENCY_LEVELS -
set([ConsistencyLevel.ANY]),
expected_reader=2)
cluster.shutdown()
def init(args):
global solr_connection
solr_connection = SolrConnection(args.solr)
global solr_collection
solr_collection = solr_connection[args.collection]
global SOLR_UNIQUE_KEY
SOLR_UNIQUE_KEY = args.solrIdField
dc_policy = RoundRobinPolicy()
token_policy = TokenAwarePolicy(dc_policy)
if args.cassandraUsername and args.cassandraPassword:
auth_provider = PlainTextAuthProvider(username=args.cassandraUsername, password=args.cassandraPassword)
else:
auth_provider = None
global cassandra_cluster
cassandra_cluster = Cluster(contact_points=args.cassandra, port=args.cassandraPort,
protocol_version=int(args.cassandraProtocolVersion),
load_balancing_policy=token_policy,
auth_provider=auth_provider)
global cassandra_session
cassandra_session = cassandra_cluster.connect(keyspace=args.cassandraKeyspace)
global cassandra_table
cassandra_table = args.cassandraTable
def init(self, config):
log = logging.getLogger(__name__)
log.info("*** STARTING DOMS INITIALIZATION ***")
domsconfig = ConfigParser.RawConfigParser()
domsconfig.readfp(pkg_resources.resource_stream(
__name__, "domsconfig.ini"),
filename='domsconfig.ini')
cassHost = domsconfig.get("cassandra", "host")
cassPort = domsconfig.get("cassandra", "port")
cassKeyspace = domsconfig.get("cassandra", "keyspace")
cassDatacenter = domsconfig.get("cassandra", "local_datacenter")
cassVersion = int(domsconfig.get("cassandra", "protocol_version"))
log.info("Cassandra Host(s): %s" % (cassHost))
log.info("Cassandra Keyspace: %s" % (cassKeyspace))
log.info("Cassandra Datacenter: %s" % (cassDatacenter))
log.info("Cassandra Protocol Version: %s" % (cassVersion))
dc_policy = DCAwareRoundRobinPolicy(cassDatacenter)
token_policy = TokenAwarePolicy(dc_policy)
with Cluster([host for host in cassHost.split(',')],
port=cassPort,
load_balancing_policy=token_policy,
protocol_version=cassVersion) as cluster:
session = cluster.connect()
self.createKeyspace(session, cassKeyspace)
self.createTables(session)
def _test_downgrading_cl(self, keyspace, rf, accepted):
cluster = TestCluster(
execution_profiles={
EXEC_PROFILE_DEFAULT:
ExecutionProfile(TokenAwarePolicy(RoundRobinPolicy()),
DowngradingConsistencyRetryPolicy())
})
session = cluster.connect(wait_for_all_pools=True)
create_schema(cluster, session, keyspace, replication_factor=rf)
self._insert(session, keyspace, 1)
self._query(session, keyspace, 1)
self.coordinator_stats.assert_query_count_equals(self, 1, 0)
self.coordinator_stats.assert_query_count_equals(self, 2, 1)
self.coordinator_stats.assert_query_count_equals(self, 3, 0)
try:
force_stop(2)
wait_for_down(cluster, 2)
self._assert_writes_succeed(session, keyspace, accepted)
self._assert_reads_succeed(session, keyspace,
accepted - set([ConsistencyLevel.ANY]))
self._assert_writes_fail(session, keyspace,
SINGLE_DC_CONSISTENCY_LEVELS - accepted)
self._assert_reads_fail(session, keyspace,
SINGLE_DC_CONSISTENCY_LEVELS - accepted)
finally:
start(2)
wait_for_up(cluster, 2)
cluster.shutdown()
def test_token_aware_with_rf_2(self, use_prepared=False):
use_singledc()
keyspace = 'test_token_aware_with_rf_2'
cluster = Cluster(load_balancing_policy=TokenAwarePolicy(
RoundRobinPolicy()),
protocol_version=PROTOCOL_VERSION)
session = cluster.connect()
wait_for_up(cluster, 1, wait=False)
wait_for_up(cluster, 2, wait=False)
wait_for_up(cluster, 3)
create_schema(session, keyspace, replication_factor=2)
self._insert(session, keyspace)
self._query(session, keyspace)
self.coordinator_stats.assert_query_count_equals(self, 1, 0)
self.coordinator_stats.assert_query_count_equals(self, 2, 12)
self.coordinator_stats.assert_query_count_equals(self, 3, 0)
self.coordinator_stats.reset_counts()
stop(2)
wait_for_down(cluster, 2, wait=True)
self._query(session, keyspace)
self.coordinator_stats.assert_query_count_equals(self, 1, 0)
self.coordinator_stats.assert_query_count_equals(self, 2, 0)
self.coordinator_stats.assert_query_count_equals(self, 3, 12)
def test_token_aware_composite_key(self):
use_singledc()
keyspace = 'test_token_aware_composite_key'
table = 'composite'
cluster = Cluster(load_balancing_policy=TokenAwarePolicy(
RoundRobinPolicy()),
protocol_version=PROTOCOL_VERSION)
session = cluster.connect()
wait_for_up(cluster, 1, wait=False)
wait_for_up(cluster, 2, wait=False)
wait_for_up(cluster, 3)
create_schema(session, keyspace, replication_factor=2)
session.execute('CREATE TABLE %s ('
'k1 int, '
'k2 int, '
'i int, '
'PRIMARY KEY ((k1, k2)))' % table)
prepared = session.prepare('INSERT INTO %s '
'(k1, k2, i) '
'VALUES '
'(?, ?, ?)' % table)
session.execute(prepared.bind((1, 2, 3)))
results = session.execute('SELECT * FROM %s WHERE k1 = 1 AND k2 = 2' %
table)
self.assertTrue(len(results) == 1)
self.assertTrue(results[0].i)
def _test_tokenaware_one_node_down(self, keyspace, rf, accepted):
cluster = Cluster(load_balancing_policy=TokenAwarePolicy(
RoundRobinPolicy()),
protocol_version=PROTOCOL_VERSION)
session = cluster.connect()
wait_for_up(cluster, 1, wait=False)
wait_for_up(cluster, 2)
create_schema(session, keyspace, replication_factor=rf)
self._insert(session, keyspace, count=1)
self._query(session, keyspace, count=1)
self.coordinator_stats.assert_query_count_equals(self, 1, 0)
self.coordinator_stats.assert_query_count_equals(self, 2, 1)
self.coordinator_stats.assert_query_count_equals(self, 3, 0)
try:
force_stop(2)
wait_for_down(cluster, 2)
self._assert_writes_succeed(session, keyspace, accepted)
self._assert_reads_succeed(session, keyspace,
accepted - set([ConsistencyLevel.ANY]))
self._assert_writes_fail(session, keyspace,
SINGLE_DC_CONSISTENCY_LEVELS - accepted)
self._assert_reads_fail(session, keyspace,
SINGLE_DC_CONSISTENCY_LEVELS - accepted)
finally:
start(2)
wait_for_up(cluster, 2)
def test_token_aware_with_rf_2(self, use_prepared=False):
use_singledc()
keyspace = 'test_token_aware_with_rf_2'
cluster, session = self._cluster_session_with_lbp(
TokenAwarePolicy(RoundRobinPolicy()))
self._wait_for_nodes_up(range(1, 4), cluster)
create_schema(cluster, session, keyspace, replication_factor=2)
self._insert(session, keyspace)
self._query(session, keyspace)
self.coordinator_stats.assert_query_count_equals(self, 1, 0)
self.coordinator_stats.assert_query_count_equals(self, 2, 12)
self.coordinator_stats.assert_query_count_equals(self, 3, 0)
self.coordinator_stats.reset_counts()
stop(2)
self._wait_for_nodes_down([2], cluster)
self._query(session, keyspace)
self.coordinator_stats.assert_query_count_equals(self, 1, 0)
self.coordinator_stats.assert_query_count_equals(self, 2, 0)
self.coordinator_stats.assert_query_count_equals(self, 3, 12)
cluster.shutdown()
def test_statement_keyspace(self):
hosts = [Host(str(i), SimpleConvictionPolicy) for i in range(4)]
for host in hosts:
host.set_up()
cluster = Mock(spec=Cluster)
cluster.metadata = Mock(spec=Metadata)
replicas = hosts[2:]
cluster.metadata.get_replicas.return_value = replicas
child_policy = Mock()
child_policy.make_query_plan.return_value = hosts
child_policy.distance.return_value = HostDistance.LOCAL
policy = TokenAwarePolicy(child_policy)
policy.populate(cluster, hosts)
# no keyspace, child policy is called
keyspace = None
routing_key = 'routing_key'
query = Statement(routing_key=routing_key)
qplan = list(policy.make_query_plan(keyspace, query))
self.assertEqual(hosts, qplan)
self.assertEqual(cluster.metadata.get_replicas.call_count, 0)
child_policy.make_query_plan.assert_called_once_with(keyspace, query)
# working keyspace, no statement
cluster.metadata.get_replicas.reset_mock()
keyspace = 'working_keyspace'
routing_key = 'routing_key'
query = Statement(routing_key=routing_key)
qplan = list(policy.make_query_plan(keyspace, query))
self.assertEqual(replicas + hosts[:2], qplan)
cluster.metadata.get_replicas.assert_called_with(keyspace, routing_key)
# statement keyspace, no working
cluster.metadata.get_replicas.reset_mock()
working_keyspace = None
statement_keyspace = 'statement_keyspace'
routing_key = 'routing_key'
query = Statement(routing_key=routing_key, keyspace=statement_keyspace)
qplan = list(policy.make_query_plan(working_keyspace, query))
self.assertEqual(replicas + hosts[:2], qplan)
cluster.metadata.get_replicas.assert_called_with(
statement_keyspace, routing_key)
# both keyspaces set, statement keyspace used for routing
cluster.metadata.get_replicas.reset_mock()
working_keyspace = 'working_keyspace'
statement_keyspace = 'statement_keyspace'
routing_key = 'routing_key'
query = Statement(routing_key=routing_key, keyspace=statement_keyspace)
qplan = list(policy.make_query_plan(working_keyspace, query))
self.assertEqual(replicas + hosts[:2], qplan)
cluster.metadata.get_replicas.assert_called_with(
statement_keyspace, routing_key)
def _set_up_shuffle_test(self, keyspace, replication_factor):
use_singledc()
cluster, session = self._cluster_session_with_lbp(
TokenAwarePolicy(RoundRobinPolicy(), shuffle_replicas=True)
)
self._wait_for_nodes_up(range(1, 4), cluster)
create_schema(cluster, session, keyspace, replication_factor=replication_factor)
return cluster, session
def __init__(self):
# Apache Cassandra connection
parser = configloader.ConfigParser()
parser.read_ini_file()
print(parser.getKey('DEFAULT', 'ip'))
list_of_ip = ['192.168.56.101']
self.cluster = Cluster(list_of_ip,
load_balancing_policy=TokenAwarePolicy(
RoundRobinPolicy()))
def test_rfthree_tokenaware_downgradingcl(self):
keyspace = 'test_rfthree_tokenaware_downgradingcl'
with TestCluster(
execution_profiles={
EXEC_PROFILE_DEFAULT:
ExecutionProfile(TokenAwarePolicy(RoundRobinPolicy()),
DowngradingConsistencyRetryPolicy())
}) as cluster:
self.rfthree_downgradingcl(cluster, keyspace, False)
def test_token_aware_policy(self):
self.assertEqual(
insights_registry.serialize(TokenAwarePolicy(child_policy=LoadBalancingPolicy())),
{'namespace': 'cassandra.policies',
'options': {'child_policy': {'namespace': 'cassandra.policies',
'options': {},
'type': 'LoadBalancingPolicy'},
'shuffle_replicas': False},
'type': 'TokenAwarePolicy'}
)
def test_token_aware_with_local_table(self):
use_singledc()
cluster, session = self._cluster_session_with_lbp(TokenAwarePolicy(RoundRobinPolicy()))
self.addCleanup(cluster.shutdown)
self._wait_for_nodes_up(range(1, 4), cluster)
p = session.prepare("SELECT * FROM system.local WHERE key=?")
# this would blow up prior to 61b4fad
r = session.execute(p, ('local',))
self.assertEqual(r[0].key, 'local')
def test_wrap_dc_aware(self):
cluster = Mock(spec=Cluster)
cluster.metadata = Mock(spec=Metadata)
hosts = [Host(str(i), SimpleConvictionPolicy) for i in range(4)]
for host in hosts:
host.set_up()
for h in hosts[:2]:
h.set_location_info("dc1", "rack1")
for h in hosts[2:]:
h.set_location_info("dc2", "rack1")
def get_replicas(keyspace, packed_key):
index = struct.unpack('>i', packed_key)[0]
# return one node from each DC
if index % 2 == 0:
return [hosts[0], hosts[2]]
else:
return [hosts[1], hosts[3]]
cluster.metadata.get_replicas.side_effect = get_replicas
policy = TokenAwarePolicy(
DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=1))
policy.populate(cluster, hosts)
for i in range(4):
query = Statement(routing_key=struct.pack('>i', i),
keyspace='keyspace_name')
qplan = list(policy.make_query_plan(None, query))
replicas = get_replicas(None, struct.pack('>i', i))
# first should be the only local replica
self.assertIn(qplan[0], replicas)
self.assertEqual(qplan[0].datacenter, "dc1")
# then the local non-replica
self.assertNotIn(qplan[1], replicas)
self.assertEqual(qplan[1].datacenter, "dc1")
# then one of the remotes (used_hosts_per_remote_dc is 1, so we
# shouldn't see two remotes)
self.assertEqual(qplan[2].datacenter, "dc2")
self.assertEqual(3, len(qplan))
def token_aware(self, keyspace, use_prepared=False):
use_singledc()
cluster = Cluster(
load_balancing_policy=TokenAwarePolicy(RoundRobinPolicy()))
session = cluster.connect()
wait_for_up(cluster, 1, wait=False)
wait_for_up(cluster, 2, wait=False)
wait_for_up(cluster, 3)
create_schema(session, keyspace, replication_factor=1)
self._insert(session, keyspace)
self._query(session, keyspace, use_prepared=use_prepared)
self.coordinator_stats.assert_query_count_equals(self, 1, 0)
self.coordinator_stats.assert_query_count_equals(self, 2, 12)
self.coordinator_stats.assert_query_count_equals(self, 3, 0)
self.coordinator_stats.reset_counts()
self._query(session, keyspace, use_prepared=use_prepared)
self.coordinator_stats.assert_query_count_equals(self, 1, 0)
self.coordinator_stats.assert_query_count_equals(self, 2, 12)
self.coordinator_stats.assert_query_count_equals(self, 3, 0)
self.coordinator_stats.reset_counts()
force_stop(2)
wait_for_down(cluster, 2, wait=True)
try:
self._query(session, keyspace, use_prepared=use_prepared)
self.fail()
except Unavailable as e:
self.assertEqual(e.consistency, 1)
self.assertEqual(e.required_replicas, 1)
self.assertEqual(e.alive_replicas, 0)
self.coordinator_stats.reset_counts()
start(2)
wait_for_up(cluster, 2, wait=True)
self._query(session, keyspace, use_prepared=use_prepared)
self.coordinator_stats.assert_query_count_equals(self, 1, 0)
self.coordinator_stats.assert_query_count_equals(self, 2, 12)
self.coordinator_stats.assert_query_count_equals(self, 3, 0)
self.coordinator_stats.reset_counts()
decommission(2)
wait_for_down(cluster, 2, wait=True)
self._query(session, keyspace, use_prepared=use_prepared)
self.coordinator_stats.assert_query_count_equals(self, 1, 6)
self.coordinator_stats.assert_query_count_equals(self, 2, 0)
self.coordinator_stats.assert_query_count_equals(self, 3, 6)
def test_wrap_dc_aware(self):
cluster = Mock(spec=Cluster)
cluster.metadata = Mock(spec=Metadata)
hosts = [Host(str(i), SimpleConvictionPolicy) for i in range(4)]
for host in hosts:
host.set_up()
for h in hosts[:2]:
h.set_location_info("dc1", "rack1")
for h in hosts[2:]:
h.set_location_info("dc2", "rack1")
def get_replicas(keyspace, packed_key):
index = struct.unpack('>i', packed_key)[0]
# return one node from each DC
if index % 2 == 0:
return [hosts[0], hosts[2]]
else:
return [hosts[1], hosts[3]]
cluster.metadata.get_replicas.side_effect = get_replicas
policy = TokenAwarePolicy(DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=1))
policy.populate(cluster, hosts)
for i in range(4):
query = Statement(routing_key=struct.pack('>i', i), keyspace='keyspace_name')
qplan = list(policy.make_query_plan(None, query))
replicas = get_replicas(None, struct.pack('>i', i))
# first should be the only local replica
self.assertIn(qplan[0], replicas)
self.assertEqual(qplan[0].datacenter, "dc1")
# then the local non-replica
self.assertNotIn(qplan[1], replicas)
self.assertEqual(qplan[1].datacenter, "dc1")
# then one of the remotes (used_hosts_per_remote_dc is 1, so we
# shouldn't see two remotes)
self.assertEqual(qplan[2].datacenter, "dc2")
self.assertEqual(3, len(qplan))
def get_lb_policy(policy_name: str, policy_args: Dict[str, Any]) -> Policy:
"""
Creates load balancing policy.
:param policy_name: Name of the policy to use.
:type policy_name: str
:param policy_args: Parameters for the policy.
:type policy_args: Dict
"""
if policy_name == 'DCAwareRoundRobinPolicy':
local_dc = policy_args.get('local_dc', '')
used_hosts_per_remote_dc = int(
policy_args.get('used_hosts_per_remote_dc', 0))
return DCAwareRoundRobinPolicy(local_dc, used_hosts_per_remote_dc)
if policy_name == 'WhiteListRoundRobinPolicy':
hosts = policy_args.get('hosts')
if not hosts:
raise Exception(
'Hosts must be specified for WhiteListRoundRobinPolicy')
return WhiteListRoundRobinPolicy(hosts)
if policy_name == 'TokenAwarePolicy':
allowed_child_policies = (
'RoundRobinPolicy',
'DCAwareRoundRobinPolicy',
'WhiteListRoundRobinPolicy',
)
child_policy_name = policy_args.get('child_load_balancing_policy',
'RoundRobinPolicy')
child_policy_args = policy_args.get(
'child_load_balancing_policy_args', {})
if child_policy_name not in allowed_child_policies:
return TokenAwarePolicy(RoundRobinPolicy())
else:
child_policy = CassandraHook.get_lb_policy(
child_policy_name, child_policy_args)
return TokenAwarePolicy(child_policy)
# Fallback to default RoundRobinPolicy
return RoundRobinPolicy()
def main():
# Apache Kafka connection
consumer = Consumer(conf)
# Apache Cassandra connection
list_of_ip = (['192.168.56.101', '192.168.56.102', '192.168.56.103'])
cluster = Cluster(list_of_ip,
load_balancing_policy=TokenAwarePolicy(
RoundRobinPolicy()))
session = cluster.connect()
session.set_keyspace('thirdeye_test')
connection.set_session(session)
basic_consume_loop(consumer, ["airline_raw"], connection)
def test_token_aware_with_local_table(self):
use_singledc()
cluster = Cluster(load_balancing_policy=TokenAwarePolicy(
RoundRobinPolicy()),
protocol_version=PROTOCOL_VERSION)
session = cluster.connect()
p = session.prepare("SELECT * FROM system.local WHERE key=?")
# this would blow up prior to 61b4fad
r = session.execute(p, ('local', ))
self.assertEqual(len(r), 1)
self.assertEqual(r[0].key, 'local')
def init(self, config):
log = logging.getLogger(__name__)
log.info("*** STARTING DOMS INITIALIZATION ***")
domsconfig = configparser.SafeConfigParser()
domsconfig.read(DomsInitializer._get_config_files('domsconfig.ini'))
domsconfig = self.override_config(domsconfig, config)
cassHost = domsconfig.get("cassandra", "host")
cassPort = domsconfig.get("cassandra", "port")
cassUsername = domsconfig.get("cassandra", "username")
cassPassword = domsconfig.get("cassandra", "password")
cassKeyspace = domsconfig.get("cassandra", "keyspace")
cassDatacenter = domsconfig.get("cassandra", "local_datacenter")
cassVersion = int(domsconfig.get("cassandra", "protocol_version"))
cassPolicy = domsconfig.get("cassandra", "dc_policy")
try:
cassCreateKeyspaceGranted = domsconfig.get(
"cassandra", "create_keyspace_granted")
except configparser.NoOptionError:
cassCreateKeyspaceGranted = "True"
log.info("Cassandra Host(s): %s" % (cassHost))
log.info("Cassandra Keyspace: %s" % (cassKeyspace))
log.info("Cassandra Datacenter: %s" % (cassDatacenter))
log.info("Cassandra Protocol Version: %s" % (cassVersion))
log.info("Cassandra DC Policy: %s" % (cassPolicy))
if cassPolicy == 'DCAwareRoundRobinPolicy':
dc_policy = DCAwareRoundRobinPolicy(cassDatacenter)
token_policy = TokenAwarePolicy(dc_policy)
elif cassPolicy == 'WhiteListRoundRobinPolicy':
token_policy = WhiteListRoundRobinPolicy([cassHost])
if cassUsername and cassPassword:
auth_provider = PlainTextAuthProvider(username=cassUsername,
password=cassPassword)
else:
auth_provider = None
with Cluster([host for host in cassHost.split(',')],
port=int(cassPort),
load_balancing_policy=token_policy,
protocol_version=cassVersion,
auth_provider=auth_provider) as cluster:
session = cluster.connect()
if cassCreateKeyspaceGranted in ["True", "true"]:
self.createKeyspace(session, cassKeyspace)
else:
session.set_keyspace(cassKeyspace)
self.createTables(session)
def cassandra_cluster():
from cassandra.cluster import Cluster
from cassandra.policies import DCAwareRoundRobinPolicy, TokenAwarePolicy, RetryPolicy
cassandra_hosts = ['10.0.0.251', '10.0.0.250', '10.0.0.249']
try:
return Cluster(contact_points=cassandra_hosts,
load_balancing_policy=TokenAwarePolicy(
DCAwareRoundRobinPolicy(local_dc='Cassandra')),
default_retry_policy=RetryPolicy())
except Error as er:
print er
def get_lb_policy(policy_name, policy_args):
policies = {
'RoundRobinPolicy': RoundRobinPolicy,
'DCAwareRoundRobinPolicy': DCAwareRoundRobinPolicy,
'WhiteListRoundRobinPolicy': WhiteListRoundRobinPolicy,
'TokenAwarePolicy': TokenAwarePolicy,
}
if not policies.get(policy_name) or policy_name == 'RoundRobinPolicy':
return RoundRobinPolicy()
if policy_name == 'DCAwareRoundRobinPolicy':
local_dc = policy_args.get('local_dc', '')
used_hosts_per_remote_dc = int(
policy_args.get('used_hosts_per_remote_dc', 0))
return DCAwareRoundRobinPolicy(local_dc, used_hosts_per_remote_dc)
if policy_name == 'WhiteListRoundRobinPolicy':
hosts = policy_args.get('hosts')
if not hosts:
raise Exception(
'Hosts must be specified for WhiteListRoundRobinPolicy')
return WhiteListRoundRobinPolicy(hosts)
if policy_name == 'TokenAwarePolicy':
allowed_child_policies = (
'RoundRobinPolicy',
'DCAwareRoundRobinPolicy',
'WhiteListRoundRobinPolicy',
)
child_policy_name = policy_args.get('child_load_balancing_policy',
'RoundRobinPolicy')
child_policy_args = policy_args.get(
'child_load_balancing_policy_args', {})
if child_policy_name not in allowed_child_policies:
return TokenAwarePolicy(RoundRobinPolicy())
else:
child_policy = CassandraHook.get_lb_policy(
child_policy_name, child_policy_args)
return TokenAwarePolicy(child_policy)
def test_wrap_round_robin(self):
cluster = Mock(spec=Cluster)
cluster.metadata = Mock(spec=Metadata)
hosts = [Host(str(i), SimpleConvictionPolicy) for i in range(4)]
def get_replicas(keyspace, packed_key):
index = struct.unpack('>i', packed_key)[0]
return list(islice(cycle(hosts), index, index + 2))
cluster.metadata.get_replicas.side_effect = get_replicas
policy = TokenAwarePolicy(RoundRobinPolicy())
policy.populate(cluster, hosts)
for i in range(4):
query = Statement(routing_key=struct.pack('>i', i))
qplan = list(policy.make_query_plan(None, query))
replicas = get_replicas(None, struct.pack('>i', i))
other = set(h for h in hosts if h not in replicas)
self.assertEquals(replicas, qplan[:2])
self.assertEquals(other, set(qplan[2:]))
# Should use the secondary policy
for i in range(4):
qplan = list(policy.make_query_plan())
self.assertEquals(set(qplan), set(hosts))
def test_wrap_round_robin(self):
cluster = Mock(spec=Cluster)
cluster.metadata = Mock(spec=Metadata)
hosts = [Host(str(i), SimpleConvictionPolicy) for i in range(4)]
def get_replicas(packed_key):
index = struct.unpack(">i", packed_key)[0]
return list(islice(cycle(hosts), index, index + 2))
cluster.metadata.get_replicas.side_effect = get_replicas
policy = TokenAwarePolicy(RoundRobinPolicy())
policy.populate(cluster, hosts)
for i in range(4):
query = Query(routing_key=struct.pack(">i", i))
qplan = list(policy.make_query_plan(query))
replicas = get_replicas(struct.pack(">i", i))
other = set(h for h in hosts if h not in replicas)
self.assertEquals(replicas, qplan[:2])
self.assertEquals(other, set(qplan[2:]))
def test_statement_keyspace(self):
hosts = [Host(str(i), SimpleConvictionPolicy) for i in range(4)]
for host in hosts:
host.set_up()
cluster = Mock(spec=Cluster)
cluster.metadata = Mock(spec=Metadata)
replicas = hosts[2:]
cluster.metadata.get_replicas.return_value = replicas
child_policy = Mock()
child_policy.make_query_plan.return_value = hosts
child_policy.distance.return_value = HostDistance.LOCAL
policy = TokenAwarePolicy(child_policy)
policy.populate(cluster, hosts)
# no keyspace, child policy is called
keyspace = None
routing_key = 'routing_key'
query = Statement(routing_key=routing_key)
qplan = list(policy.make_query_plan(keyspace, query))
self.assertEqual(hosts, qplan)
self.assertEqual(cluster.metadata.get_replicas.call_count, 0)
child_policy.make_query_plan.assert_called_once_with(keyspace, query)
# working keyspace, no statement
cluster.metadata.get_replicas.reset_mock()
keyspace = 'working_keyspace'
routing_key = 'routing_key'
query = Statement(routing_key=routing_key)
qplan = list(policy.make_query_plan(keyspace, query))
self.assertEqual(replicas + hosts[:2], qplan)
cluster.metadata.get_replicas.assert_called_with(keyspace, routing_key)
# statement keyspace, no working
cluster.metadata.get_replicas.reset_mock()
working_keyspace = None
statement_keyspace = 'statement_keyspace'
routing_key = 'routing_key'
query = Statement(routing_key=routing_key, keyspace=statement_keyspace)
qplan = list(policy.make_query_plan(working_keyspace, query))
self.assertEqual(replicas + hosts[:2], qplan)
cluster.metadata.get_replicas.assert_called_with(statement_keyspace, routing_key)
# both keyspaces set, statement keyspace used for routing
cluster.metadata.get_replicas.reset_mock()
working_keyspace = 'working_keyspace'
statement_keyspace = 'statement_keyspace'
routing_key = 'routing_key'
query = Statement(routing_key=routing_key, keyspace=statement_keyspace)
qplan = list(policy.make_query_plan(working_keyspace, query))
self.assertEqual(replicas + hosts[:2], qplan)
cluster.metadata.get_replicas.assert_called_with(statement_keyspace, routing_key)
def main():
while 1:
log.info("\n\n\n\n Starting iteration\n")
cluster = Cluster(protocol_version=PROTOCOL_VERSION, connection_class=connection_class)
cluster.connect(wait_for_all_pools=True)
hosts = cluster.metadata.all_hosts()
address = hosts[0].address
node_to_stop = int(address.split('.')[-1:][0])
cluster.shutdown()
log.info("node_to_stop: {0}".format(node_to_stop))
contact_point = '127.0.0.{0}'.format(get_node_not_x(node_to_stop))
cluster = Cluster(contact_points=[contact_point], protocol_version=PROTOCOL_VERSION, connection_class=connection_class)
cluster.connect(wait_for_all_pools=True)
try:
force_stop(node_to_stop)
wait_for_down(cluster, node_to_stop)
# Attempt a query against that node. It should complete
cluster2 = Cluster(contact_points=all_contact_points, protocol_version=PROTOCOL_VERSION, connection_class=connection_class)
session2 = cluster2.connect()
session2.execute("SELECT * FROM system.local")
except Exception as e:
traceback.print_exc()
finally:
cluster2.shutdown()
log.info("\n\nDone, with cluster2, starting node {0}\n\n".format(node_to_stop))
start(node_to_stop)
wait_for_up(cluster, node_to_stop)
cluster.shutdown()
log.info("\n\nCluster is shutdown and node {0} should be up\n\n".format(node_to_stop))
cluster3 = Cluster(
load_balancing_policy=TokenAwarePolicy(RoundRobinPolicy()),
default_retry_policy=DowngradingConsistencyRetryPolicy(),
protocol_version=PROTOCOL_VERSION,
connection_class=connection_class)
session = cluster3.connect(wait_for_all_pools=True)
for _ in range(20):
hosts = cluster3.metadata.all_hosts()
log.info("cluster.metadata.all_hosts(): {0}".format(hosts))
if len(hosts) == 3 and (map(lambda x : x.is_up, hosts) == [True] * 3):
log.info("Found all the required nodes")
break
time.sleep(1)
else:
raise Exception("Some node never came back")
cluster3.shutdown()
time.sleep(4)
def connect_to_cassandra(self,
host: str,
port: Union[int, str] = 9042,
alias: str = None,
keyspace: str = None,
username: str = None,
password: str = '') -> Session:
"""
Connect to Apache Cassandra cluster.
AllowAllAuthenticator and PasswordAuthenticator are supported as authentication backend.
This setting should be in configuration file cassandra.yaml:
by default:
| authenticator: AllowAllAuthenticator
or for password authentification:
| authenticator: PasswordAuthenticator
*Args:*\n
_host_ - IP address or host name of a cluster node;\n
_port_ - connection port;\n
_alias_ - connection alias;\n
_keyspace_ - the name of the keyspace that the UDT is defined in;\n
_username_ - username to connect to cassandra
_password_ - password for username
*Returns:*\n
Index of current connection.
*Example:*\n
| Connect To Cassandra | 192.168.1.108 | 9042 | alias=cluster1 |
"""
logger.info('Connecting using : host={0}, port={1}, alias={2}'.format(
host, port, alias))
try:
auth_provider = PlainTextAuthProvider(
username=username, password=password) if username else None
cluster = Cluster([host],
port=int(port),
auth_provider=auth_provider,
load_balancing_policy=TokenAwarePolicy(
DCAwareRoundRobinPolicy()))
session = cluster.connect()
if keyspace is not None:
session.set_keyspace(keyspace)
self._connection = session
return self._cache.register(self._connection, alias)
except Exception as e:
raise Exception('Connect to Cassandra error: {0}'.format(e))
def __init__(self,
num_classes,
aug,
table,
label_col,
data_col,
id_col,
username,
cass_pass,
cassandra_ips,
thread_par=32,
port=9042):
self.aug = aug
self.num_classes = num_classes
self.label_col = label_col
self.data_col = data_col
self.id_col = id_col
self.finished_event = threading.Event()
self.lock = threading.Lock()
self.thread_par = thread_par
self.tot = None
self.cow = 0
self.onair = 0
self.errors = []
self.feats = []
self.labels = []
self.perm = []
self.bb = None
## multi-label when num_classes is small
self.multi_label = (num_classes <= _max_multilabs)
## cassandra parameters
prof_dict = ExecutionProfile(load_balancing_policy=TokenAwarePolicy(
DCAwareRoundRobinPolicy()),
row_factory=cassandra.query.dict_factory)
profs = {'dict': prof_dict}
auth_prov = PlainTextAuthProvider(username=username,
password=cass_pass)
self.cluster = Cluster(cassandra_ips,
execution_profiles=profs,
protocol_version=4,
auth_provider=auth_prov,
port=port)
self.cluster.connect_timeout = 10 #seconds
self.sess = self.cluster.connect()
self.table = table
query = f"SELECT {self.label_col}, {self.data_col} \
FROM {self.table} WHERE {self.id_col}=?"
self.prep = self.sess.prepare(query)
def __open(self):
if self.__cass_dc_policy == 'DCAwareRoundRobinPolicy':
dc_policy = DCAwareRoundRobinPolicy(self.__cass_local_DC)
elif self.__cass_dc_policy == 'WhiteListRoundRobinPolicy':
dc_policy = WhiteListRoundRobinPolicy([self.__cass_url])
if self.__cass_username and self.__cass_password:
auth_provider = PlainTextAuthProvider(username=self.__cass_username, password=self.__cass_password)
else:
auth_provider = None
token_policy = TokenAwarePolicy(dc_policy)
connection.setup([host for host in self.__cass_url.split(',')], self.__cass_keyspace,
protocol_version=self.__cass_protocol_version, load_balancing_policy=token_policy,
port=self.__cass_port,
auth_provider=auth_provider)
def test_status_updates(self):
"""
Same test as DCAwareRoundRobinPolicyTest.test_status_updates()
"""
hosts = [Host(i, SimpleConvictionPolicy) for i in range(4)]
for h in hosts[:2]:
h.set_location_info("dc1", "rack1")
for h in hosts[2:]:
h.set_location_info("dc2", "rack1")
policy = TokenAwarePolicy(DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=1))
policy.populate(self.FakeCluster(), hosts)
policy.on_down(hosts[0])
policy.on_remove(hosts[2])
new_local_host = Host(4, SimpleConvictionPolicy)
new_local_host.set_location_info("dc1", "rack1")
policy.on_up(new_local_host)
new_remote_host = Host(5, SimpleConvictionPolicy)
new_remote_host.set_location_info("dc9000", "rack1")
policy.on_add(new_remote_host)
# we now have two local hosts and two remote hosts in separate dcs
qplan = list(policy.make_query_plan())
self.assertEqual(set(qplan[:2]), set([hosts[1], new_local_host]))
self.assertEqual(set(qplan[2:]), set([hosts[3], new_remote_host]))
# since we have hosts in dc9000, the distance shouldn't be IGNORED
self.assertEqual(policy.distance(new_remote_host), HostDistance.REMOTE)
policy.on_down(new_local_host)
policy.on_down(hosts[1])
qplan = list(policy.make_query_plan())
self.assertEqual(set(qplan), set([hosts[3], new_remote_host]))
policy.on_down(new_remote_host)
policy.on_down(hosts[3])
qplan = list(policy.make_query_plan())
self.assertEqual(qplan, [])
def test_get_distance(self):
"""
Same test as DCAwareRoundRobinPolicyTest.test_get_distance()
Except a FakeCluster is needed for the metadata variable and
policy.child_policy is needed to change child policy settings
"""
policy = TokenAwarePolicy(DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=0))
host = Host("ip1", SimpleConvictionPolicy)
host.set_location_info("dc1", "rack1")
policy.populate(self.FakeCluster(), [host])
self.assertEqual(policy.distance(host), HostDistance.LOCAL)
# used_hosts_per_remote_dc is set to 0, so ignore it
remote_host = Host("ip2", SimpleConvictionPolicy)
remote_host.set_location_info("dc2", "rack1")
self.assertEqual(policy.distance(remote_host), HostDistance.IGNORED)
# dc2 isn't registered in the policy's live_hosts dict
policy._child_policy.used_hosts_per_remote_dc = 1
self.assertEqual(policy.distance(remote_host), HostDistance.IGNORED)
# make sure the policy has both dcs registered
policy.populate(self.FakeCluster(), [host, remote_host])
self.assertEqual(policy.distance(remote_host), HostDistance.REMOTE)
# since used_hosts_per_remote_dc is set to 1, only the first
# remote host in dc2 will be REMOTE, the rest are IGNORED
second_remote_host = Host("ip3", SimpleConvictionPolicy)
second_remote_host.set_location_info("dc2", "rack1")
policy.populate(self.FakeCluster(), [host, remote_host, second_remote_host])
distances = set([policy.distance(remote_host), policy.distance(second_remote_host)])
self.assertEqual(distances, set([HostDistance.REMOTE, HostDistance.IGNORED]))