def get_topic_partition_with_error(cluster_config, error, fetch_unavailable_brokers=False): """Fetches the metadata from the cluster and returns the set of (topic, partition) tuples containing all the topic-partitions currently affected by the specified error. It also fetches unavailable-broker list if required.""" metadata = get_topic_partition_metadata(cluster_config.broker_list) affected_partitions = set() if fetch_unavailable_brokers: unavailable_brokers = set() with ZK(cluster_config) as zk: for partitions in metadata.values(): for partition_metadata in partitions.values(): if int(partition_metadata.error) == error: if fetch_unavailable_brokers: unavailable_brokers |= get_unavailable_brokers( zk, partition_metadata) affected_partitions.add((partition_metadata.topic, partition_metadata.partition)) if fetch_unavailable_brokers: return affected_partitions, unavailable_brokers else: return affected_partitions
def test_set_topic_config(self, mock_client): with mock.patch.object( ZK, 'set', autospec=True ) as mock_set: with ZK(self.cluster_config) as zk: zk.set_topic_config( "some_topic", {"version": 1, "config": {"cleanup.policy": "compact"}} ) mock_set.assert_called_once_with( zk, '/config/topics/some_topic', json.dumps({"version": 1, "config": {"cleanup.policy": "compact"}}) ) expected_create_call = mock.call( '/config/changes/config_change_', "some_topic", None, False, True, False ) assert mock_client.return_value.create.call_args_list == [expected_create_call]
def test_create(self, mock_client): with ZK(self.cluster_config) as zk: zk.create('/kafka/consumers/some_group/offsets') zk.create('/kafka/consumers/some_group/offsets', value='some_val', acl=None, ephemeral=True, sequence=True, makepath=True) mock_obj = mock.Mock() zk.create( '/kafka/consumers/some_group/offsets', value='some_val', acl=mock_obj, ) call_list = [ mock.call('/kafka/consumers/some_group/offsets', '', None, False, False, False), mock.call('/kafka/consumers/some_group/offsets', 'some_val', None, True, True, True), mock.call('/kafka/consumers/some_group/offsets', 'some_val', mock_obj, False, False, False), ] assert mock_client.return_value.create.call_args_list == call_list
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 groupid in zk.get_children("/consumers"): print( "Error: Offsets for Consumer Group ID {groupid} not found.".format( groupid=groupid ), file=sys.stderr, ) else: 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 apply_log_compaction(self, topics): self.log.info("Applying compaction settings on {} topics".format( len(topics))) compacted_topics = [] skipped_topics = [] missed_topics = [] cluster = get_config().cluster_config with ZK(cluster) as zk: for topic in topics: try: current_config = zk.get_topic_config(topic) if 'cleanup.policy' not in current_config['config']: # if we already have the config set or there was a # manual override we don't want to set again current_config['config']['cleanup.policy'] = 'compact' if not self.dry_run: zk.set_topic_config(topic=topic, value=current_config) compacted_topics.append(topic) else: skipped_topics.append(topic) except NoNodeError: missed_topics.append(topic) self.log_results(compacted_topics=compacted_topics, skipped_topics=skipped_topics, missed_topics=missed_topics)
def run(self, cluster_config, args): self.cluster_config = cluster_config self.args = args with ZK(self.cluster_config) as self.zk: broker_ids = get_broker_ids(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 broker_ids: terminate( status_code.OK, prepare_terminate_message( 'No brokers detected, nothing to check'), args.json, ) if args.first_broker_only and not is_first_broker( broker_ids, 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 test_delete_topic_partitions(self, mock_client): with mock.patch.object( ZK, 'delete', autospec=True ) as mock_delete: with ZK(self.cluster_config) as zk: zk.delete_topic_partitions( 'some_group', 'some_topic', [0, 1, 2] ) call_list = [ mock.call( zk, '/consumers/some_group/offsets/some_topic/0' ), mock.call( zk, '/consumers/some_group/offsets/some_topic/1' ), mock.call( zk, '/consumers/some_group/offsets/some_topic/2' ), ] assert mock_delete.call_args_list == call_list
def rename_group_with_storage_zookeeper(cls, old_groupid, new_groupid, topics_dict, cluster_config): with ZK(cluster_config) as zk: try: topics = zk.get_children( "/consumers/{groupid}/offsets".format(groupid=new_groupid)) except NoNodeError: # Consumer Group ID doesn't exist. pass else: preprocess_topics( old_groupid, list(topics_dict.keys()), new_groupid, topics, ) old_offsets = fetch_offsets(zk, old_groupid, topics_dict) create_offsets(zk, new_groupid, old_offsets) try: old_base_path = "/consumers/{groupid}".format( groupid=old_groupid, ) zk.delete(old_base_path, recursive=True) except: print( "Error: Unable to migrate all metadata in Zookeeper. " "Please re-run the command.", file=sys.stderr) raise
def run(cls, args, cluster_config): # Setup the Kafka client client = KafkaToolClient(cluster_config.broker_list) client.load_metadata_for_topics() topics_dict = cls.preprocess_args( args.groupid, args.topic, args.partitions, cluster_config, client ) with ZK(cluster_config) as zk: if args.storage == 'zookeeper': unsubscriber = ZookeeperUnsubscriber(zk) elif args.storage == 'kafka': unsubscriber = KafkaUnsubscriber(client) else: print( "Invalid storage option: {}".format(args.storage), file=sys.stderr, ) sys.exit(1) unsubscriber.unsubscribe_topic( args.groupid, args.topic, args.partitions, topics_dict, )
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: check_run_on_controller(self.zk, self.args) return self.run_command()
def step_impl5(context): cluster_config = get_cluster_config() with ZK(cluster_config) as zk: offsets = zk.get_group_offsets(context.group) new_offsets = zk.get_group_offsets(NEW_GROUP) assert context.topic in offsets assert new_offsets == offsets
def test_set_topic_config_kafka_9(self, mock_client): with mock.patch.object(ZK, 'set', autospec=True) as mock_set: with ZK(self.cluster_config) as zk: config = { "version": 1, "config": { "cleanup.policy": "compact" } } config_change = { "version": 1, "entity_type": "topics", "entity_name": "some_topic" } zk.set_topic_config("some_topic", config, (0, 9, 2)) serialized_config = dump_json(config) serialized_config_change = dump_json(config_change) mock_set.assert_called_once_with( zk, '/config/topics/some_topic', serialized_config, ) expected_create_call = mock.call( '/config/changes/config_change_', serialized_config_change, None, False, True, False) assert mock_client.return_value.create.call_args_list == [ expected_create_call ]
def test_set_broker_config_kafka_10(self, mock_client): with mock.patch.object(ZK, 'set', autospec=True) as mock_set: with ZK(self.cluster_config) as zk: config = { "version": 1, "config": { "leader.replication.throttled.rate": "42" } } config_change = {"entity_path": "brokers/0", "version": 2} zk.set_broker_config(0, config) serialized_config = dump_json(config) serialized_config_change = dump_json(config_change) mock_set.assert_called_once_with( zk, '/config/brokers/0', serialized_config, ) expected_create_call = mock.call( '/config/changes/config_change_', serialized_config_change, None, False, True, False) assert mock_client.return_value.create.call_args_list == [ expected_create_call ]
def test_get_brokers_names_only(self, mock_client): with ZK(self.cluster_config) as zk: zk.get_children = mock.Mock(return_value=[1, 2, 3], ) expected = {1: None, 2: None, 3: None} actual = zk.get_brokers(names_only=True) zk.get_children.assert_called_with("/brokers/ids") assert actual == expected
def test_get_brokers_empty_cluster(self, mock_client): with ZK(self.cluster_config) as zk: zk.get_children = mock.Mock(side_effect=NoNodeError()) actual_with_no_node_error = zk.get_brokers() expected_with_no_node_error = {} zk.get_children.assert_called_with("/brokers/ids") assert actual_with_no_node_error == expected_with_no_node_error
def step_impl3(context): cluster_config = get_cluster_config() with ZK(cluster_config) as zk: current_config = zk.get_topic_config(context.topic) current_config['config']['max.message.bytes'] = '1000' zk.set_topic_config(context.topic, value=current_config) time.sleep( 2) # sleeping for 2 seconds to ensure config is actually picked up
def test_get_topic_config(self, mock_client): with ZK(self.cluster_config) as zk: zk.zk.get = mock.Mock(return_value=( b'{"version": 1, "config": {"cleanup.policy": "compact"}}', "Random node info that doesn't matter")) actual = zk.get_topic_config("some_topic") expected = {"version": 1, "config": {"cleanup.policy": "compact"}} assert actual == expected
def get_broker_list(cluster_config): """Returns a dictionary of brokers in the form {id: host} :param cluster_config: the configuration of the cluster :type cluster_config: map """ with ZK(cluster_config) as zk: brokers = sorted(zk.get_brokers().items(), key=itemgetter(0)) return [(id, data['host']) for id, data in brokers]
def test_get_brokers_with_metadata_for_plaintext(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":[],"host":"broker"}', None)) expected = {1: {'host': 'broker'}} actual = zk.get_brokers() assert actual[1]['host'] == expected[1]['host']
def test_get_brokers_with_metadata_for_sasl(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":["PLAINTEXTSASL://broker:9093"],"host":null}', None)) expected = {1: {'host': 'broker'}} actual = zk.get_brokers() assert actual[1]['host'] == expected[1]['host']
def test_get_brokers_with_metadata(self, mock_client): with ZK(self.cluster_config) as zk: zk.get_children = mock.Mock(return_value=[1, 2, 3], ) zk.get_broker_metadata = mock.Mock(return_value='broker', ) expected = {1: 'broker', 2: 'broker', 3: 'broker'} actual = zk.get_brokers() zk.get_children.assert_called_with("/brokers/ids") calls = zk.get_broker_metadata.mock_calls zk.get_broker_metadata.assert_has_calls(calls) assert actual == expected
def get_zookeeper_groups(cls, cluster_config): '''Get the group_id of groups committed into Zookeeper.''' with ZK(cluster_config) as zk: try: return zk.get_children("/consumers") except NoNodeError: print( "Error: No consumers node found in zookeeper", file=sys.stderr, )
def test_get_nonexistent_topic_config(self, mock_client): """ Test getting configuration for topics that don't exist. """ with ZK(self.cluster_config) as zk: zk.zk.get = mock.Mock(side_effect=NoNodeError()) zk.get_topics = mock.Mock(return_value={}) with pytest.raises(NoNodeError): zk.get_topic_config("some_topic")
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 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 run(cls, args, cluster_config): # Setup the Kafka client client = KafkaToolClient(cluster_config.broker_list) client.load_metadata_for_topics() if args.topic and args.topics: print( "Error: Cannot specify --topic and --topics at the same time.", file=sys.stderr, ) sys.exit(1) if args.partitions and args.topics: print( "Error: Cannot use --partitions with --topics. Use --topic " "instead.", file=sys.stderr, ) sys.exit(1) # if topic is not None topics_dict will contain only info about that # topic, otherwise it will contain info about all topics for the group topics_dict = cls.preprocess_args( args.groupid, args.topic, args.partitions, cluster_config, client, storage=args.storage, ) topics = args.topics if args.topics else ( [args.topic] if args.topic else []) for topic in topics: if topic not in topics_dict: print( "Error: Consumer {groupid} is not subscribed to topic:" " {topic}.".format(groupid=args.groupid, topic=topic), file=sys.stderr, ) sys.exit(1) with ZK(cluster_config) as zk: if args.storage == 'zookeeper': unsubscriber = ZookeeperUnsubscriber(zk) else: unsubscriber = KafkaUnsubscriber(client) unsubscriber.unsubscribe_topics( args.groupid, topics, args.partitions, topics_dict, )
def test_get_broker_config(self, mock_client): with ZK(self.cluster_config) as zk: zk.zk.get = mock.Mock( return_value=( b'{"version": 1, "config": {"leader.replication.throttled.rate": "42"}}', "Random node info that doesn't matter" ) ) actual = zk.get_broker_config(0) expected = {"version": 1, "config": {"leader.replication.throttled.rate": "42"}} assert actual == expected
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 get_broker_list(cluster_config): """Returns a dictionary of brokers in the form {id: host} :param cluster_config: the configuration of the cluster :type cluster_config: kafka_utils.utils.config.ClusterConfig :returns: all the brokers in the cluster :rtype: map of (broker_id, host) pairs """ with ZK(cluster_config) as zk: brokers = sorted(list(zk.get_brokers().items()), key=itemgetter(0)) return [(int(id), data['host']) for id, data in brokers]
def test_get_topic_config_8(self, mock_client): """ Test getting configuration for topics created in Kafa prior to 0.9.0. """ with ZK(self.cluster_config) as zk: zk.zk.get = mock.Mock(side_effect=NoNodeError()) zk.get_topics = mock.Mock(return_value={"some_topic": {}}) actual = zk.get_topic_config("some_topic") expected = {"config": {}} assert actual == expected