def test_get_my_subscribed_partitions(self, _):
with mock.patch.object(
ZK,
'get_children',
autospec=True,
) as mock_children:
with ZK(self.cluster_config) as zk:
zk.get_my_subscribed_partitions(
'some_group',
'some_topic',
)
mock_children.assert_called_once_with(
zk,
'/consumers/some_group/offsets/some_topic',
)
def get_topics_for_group_from_zookeeper(cls, cluster_config, groupid,
fail_on_error):
topics = []
with ZK(cluster_config) as zk:
# Query zookeeper to get the list of topics that this consumer is
# subscribed to.
try:
topics = zk.get_my_subscribed_topics(groupid)
except NoNodeError:
if fail_on_error:
print("Error: Consumer Group ID {groupid} does not exist.".
format(groupid=groupid),
file=sys.stderr)
sys.exit(1)
return topics
def test_delete_topic(self, _):
with mock.patch.object(
ZK,
'delete',
autospec=True
) as mock_delete:
with ZK(self.cluster_config) as zk:
zk.delete_topic(
'some_group',
'some_topic',
)
mock_delete.assert_called_once_with(
zk,
'/consumers/some_group/offsets/some_topic',
True,
)
def test_get_brokers_with_metadata_for_ssl(self, mock_client):
with ZK(self.cluster_config) as zk:
zk.get_children = mock.Mock(return_value=[1], )
zk.get = mock.Mock(return_value=(
b'{"endpoints":["SSL://broker:9093"],"host":null}', None))
expected = {1: {'host': 'broker'}}
actual = zk.get_brokers()
assert actual[1]['host'] == expected[1]['host']
zk.get = mock.Mock(return_value=(
b'{"endpoints":["INTERNAL://broker:9093","EXTERNAL://broker:9093"],"host":null}',
None))
expected = {1: {'host': 'broker'}}
actual = zk.get_brokers()
assert actual[1]['host'] == expected[1]['host']
def test_get_topics(self, mock_client):
with ZK(self.cluster_config) as zk:
zk.zk.get = mock.Mock(return_value=((
b'{"version": "1", "partitions": {"0": [1, 0]}}',
MockGetTopics(31000),
)))
zk._fetch_partition_state = mock.Mock(return_value=((
b'{"version": "2"}',
MockGetTopics(32000),
)))
actual_with_fetch_state = zk.get_topics("some_topic")
expected_with_fetch_state = {
'some_topic': {
'ctime': 31.0,
'partitions': {
'0': {
'replicas': [1, 0],
'ctime': 32.0,
'version': '2',
},
},
'version': '1',
},
}
assert actual_with_fetch_state == expected_with_fetch_state
zk._fetch_partition_info = mock.Mock(
return_value=MockGetTopics(33000))
actual_without_fetch_state = zk.get_topics(
"some_topic", fetch_partition_state=False)
expected_without_fetch_state = {
'some_topic': {
'ctime': 31.0,
'partitions': {
'0': {
'replicas': [1, 0],
'ctime': 33.0,
},
},
'version': '1',
},
}
assert actual_without_fetch_state == expected_without_fetch_state
def run(self, cluster_config, rg_parser, args):
self.cluster_config = cluster_config
self.args = args
with ZK(self.cluster_config) as self.zk:
self.log.debug(
'Starting %s for cluster: %s and zookeeper: %s',
self.__class__.__name__,
self.cluster_config.name,
self.cluster_config.zookeeper,
)
brokers = self.zk.get_brokers()
assignment = self.zk.get_cluster_assignment()
ct = ClusterTopology(
assignment,
brokers,
rg_parser.get_replication_group,
)
self.run_command(ct)
def run():
opts = parse_opts()
if opts.verbose:
logging.basicConfig(level=logging.DEBUG)
else:
logging.basicConfig(level=logging.WARN)
if not validate_opts(opts):
sys.exit(1)
cluster_config = config.get_cluster_config(
opts.cluster_type,
opts.cluster_name,
opts.discovery_base_path,
)
with ZK(cluster_config) as zk:
brokers = zk.get_brokers(names_only=True)
print_throttles(zk, brokers)
if opts.read_only:
return
print("Applying new replication throttles")
if not opts.clear:
apply_throttles(
zk,
brokers,
opts.leader_throttle,
opts.follower_throttle,
)
else:
clear_throttles(zk, brokers)
print("New replication throttles applied.")
print_throttles(zk, brokers)
if not opts.clear:
print(
"NOTE: Do not forget to --clear throttles once the reassignment plan completes."
)
def run(
self,
cluster_config,
rg_parser,
partition_measurer,
cluster_balancer,
args,
):
"""Initialize cluster_config, args, and zk then call run_command."""
self.cluster_config = cluster_config
self.args = args
with ZK(self.cluster_config) as self.zk:
self.log.debug(
'Starting %s for cluster: %s and zookeeper: %s',
self.__class__.__name__,
self.cluster_config.name,
self.cluster_config.zookeeper,
)
brokers = self.zk.get_brokers()
assignment = self.zk.get_cluster_assignment()
pm = partition_measurer(
self.cluster_config,
brokers,
assignment,
args,
)
ct = ClusterTopology(
assignment,
brokers,
pm,
rg_parser.get_replication_group,
)
if len(ct.partitions) == 0:
self.log.info("The cluster is empty. No actions to perform.")
return
# Exit if there is an on-going reassignment
if self.is_reassignment_pending():
self.log.error('Previous reassignment pending.')
sys.exit(1)
self.run_command(ct, cluster_balancer(ct, args))
def test_delete(self, mock_client):
with ZK(self.cluster_config) as zk:
zk.delete(
'/kafka/consumers/some_group/offsets',
)
zk.delete(
'/kafka/consumers/some_group/offsets',
recursive=True
)
call_list = [
mock.call(
'/kafka/consumers/some_group/offsets',
recursive=False
),
mock.call(
'/kafka/consumers/some_group/offsets',
recursive=True
),
]
assert mock_client.return_value.delete.call_args_list == call_list
def copy_group_zk(cls, client, topics, source_group, destination_group,
cluster_config):
with ZK(cluster_config) as zk:
try:
topics_dest_group = zk.get_children(
"/consumers/{groupid}/offsets".format(
groupid=destination_group, ))
except NoNodeError:
# Consumer Group ID doesn't exist.
pass
else:
preprocess_topics(
source_group,
topics,
destination_group,
topics_dest_group,
)
# Fetch offsets
source_offsets = fetch_offsets(zk, source_group, topics)
create_offsets(zk, destination_group, source_offsets)
def test_set(self, mock_client):
with ZK(self.cluster_config) as zk:
zk.set(
'config/topics/some_topic',
'some_val'
)
zk.set(
'brokers/topics/some_topic',
'{"name": "some_topic", "more": "properties"}'
)
call_list = [
mock.call(
'config/topics/some_topic',
'some_val'
),
mock.call(
'brokers/topics/some_topic',
'{"name": "some_topic", "more": "properties"}'
)
]
assert mock_client.return_value.set.call_args_list == call_list
def get_broker_list(cluster_config, active_controller_for_last=False):
"""Returns a dictionary of brokers in the form {id: host}
:param cluster_config: the configuration of the cluster
:type cluster_config: map
:param active_controller_for_last: end with active controller
:type active_controller_for_last: bool
"""
with ZK(cluster_config) as zk:
brokers_dict = zk.get_brokers()
if active_controller_for_last:
active_controller_id = zk.get_json('/controller').get('brokerid')
active_controller = brokers_dict.pop(active_controller_id, None)
if active_controller is None:
print(
"ERROR: Unable to retrieve the active controller information"
)
sys.exit(1)
brokers = sorted(list(brokers_dict.items()), key=itemgetter(0))
if active_controller_for_last:
brokers.append((active_controller_id, active_controller))
return [(id, data['host']) for id, data in brokers]
def run(self, cluster_config, args):
self.cluster_config = cluster_config
self.args = args
with ZK(self.cluster_config) as self.zk:
if args.controller_only and not is_controller(
self.zk, args.broker_id):
terminate(
status_code.OK,
prepare_terminate_message(
'Broker {} is not the controller, nothing to check'.
format(args.broker_id), ),
args.json,
)
if args.first_broker_only and not is_first_broker(
self.zk, args.broker_id):
terminate(
status_code.OK,
prepare_terminate_message(
'Broker {} has not the lowest id, nothing to check'.
format(args.broker_id), ),
args.json,
)
return self.run_command()
def run(cls, args, cluster_config):
if args.source_groupid == args.dest_groupid:
print(
"Error: Source group ID and destination group ID are same.",
file=sys.stderr,
)
sys.exit(1)
# Setup the Kafka client
client = KafkaToolClient(cluster_config.broker_list)
client.load_metadata_for_topics()
source_topics = cls.preprocess_args(
args.source_groupid,
args.topic,
args.partitions,
cluster_config,
client,
)
with ZK(cluster_config) as zk:
try:
topics_dest_group = zk.get_children(
"/consumers/{groupid}/offsets".format(
groupid=args.dest_groupid, ))
except NoNodeError:
# Consumer Group ID doesn't exist.
pass
else:
preprocess_topics(
args.source_groupid,
source_topics.keys(),
args.dest_groupid,
topics_dest_group,
)
# Fetch offsets
source_offsets = fetch_offsets(zk, args.source_groupid,
source_topics)
create_offsets(zk, args.dest_groupid, source_offsets)
def step_impl4(context):
cluster_config = get_cluster_config()
with ZK(cluster_config) as zk:
offsets = zk.get_group_offsets(context.group)
assert offsets[context.topic]["0"] == context.msgs_produced
def delete_group_zk(cls, cluster_config, group):
with ZK(cluster_config) as zk:
zk.delete_group(group)
def step_impl4(context):
cluster_config = get_cluster_config()
with ZK(cluster_config) as zk:
offsets = zk.get_group_offsets(context.group)
assert context.topic not in offsets
def step_impl4(context):
cluster_config = get_cluster_config()
with ZK(cluster_config) as zk:
offsets = zk.get_group_offsets(context.group)
assert offsets[context.topic]["0"] == RESTORED_OFFSET
context.offsets_file.close()
def set_min_isr(topic, min_isr):
cluster_config = get_cluster_config()
with ZK(cluster_config) as zk:
config = zk.get_topic_config(topic)
config['config'] = {ISR_CONF_NAME: str(min_isr)}
zk.set_topic_config(topic, config)
def _kafka_topics(self):
with ZK(self.config.cluster_config) as zk:
return zk.get_topics(names_only=True, fetch_partition_state=False)
def step_impl3(context):
cluster_config = get_cluster_config()
with ZK(cluster_config) as zk:
offsets = zk.get_group_offsets(context.group)
assert offsets[context.topic]["0"] == SET_OFFSET