def test_has_metadata_for_topic(self, protocol, conn): conn.recv.return_value = 'response' # anything but None brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [ TopicMetadata('topic_still_creating', NO_LEADER, []), TopicMetadata('topic_doesnt_exist', UNKNOWN_TOPIC_OR_PARTITION, []), TopicMetadata('topic_noleaders', NO_ERROR, [ PartitionMetadata('topic_noleaders', 0, -1, [], [], NO_LEADER), PartitionMetadata('topic_noleaders', 1, -1, [], [], NO_LEADER), ]), ] protocol.decode_metadata_response.return_value = MetadataResponse( brokers, topics) client = KafkaClient(hosts=['broker_1:4567']) # Topics with no partitions return False self.assertFalse(client.has_metadata_for_topic('topic_still_creating')) self.assertFalse(client.has_metadata_for_topic('topic_doesnt_exist')) # Topic with partition metadata, but no leaders return True self.assertTrue(client.has_metadata_for_topic('topic_noleaders'))
def test_has_metadata_for_topic(self, protocol, conn): conn.recv.return_value = 'response' # anything but None brokers = [ BrokerMetadata(0, 'broker_1', 4567), BrokerMetadata(1, 'broker_2', 5678) ] topics = [ TopicMetadata(b'topic_still_creating', NO_LEADER, []), TopicMetadata(b'topic_doesnt_exist', UNKNOWN_TOPIC_OR_PARTITION, []), TopicMetadata(b'topic_noleaders', NO_ERROR, [ PartitionMetadata(b'topic_noleaders', 0, -1, [], [], NO_LEADER), PartitionMetadata(b'topic_noleaders', 1, -1, [], [], NO_LEADER), ]), ] protocol.decode_metadata_response.return_value = MetadataResponse(brokers, topics) client = KafkaClient(hosts=['broker_1:4567']) # Topics with no partitions return False self.assertFalse(client.has_metadata_for_topic('topic_still_creating')) self.assertFalse(client.has_metadata_for_topic('topic_doesnt_exist')) # Topic with partition metadata, but no leaders return True self.assertTrue(client.has_metadata_for_topic('topic_noleaders'))
class KafkaHelper: """ Utility class to interact with Kafka Brokers Internally uses kafka-python library """ def __init__(self): # TODO: Make kafka broker list configurable try: self.kafka = KafkaClient(kw_settings.KW_KAFKA_BROKER_LIST) except: print 'Error - connecting to Kafka broker : ' + kw_settings.KW_KAFKA_BROKER_LIST self.kafka = None self.retry_count = 5 self.retry_interval_in_ms = 5000 def close(self): if self.kafka: self.kafka.close() def _ensure_kafka_topic_exists(self, topic): result = False for i in range(self.retry_count): try: self.kafka.ensure_topic_exists(topic) result = True break except: print 'Warning - Unable to create kafka topic : ' + topic print traceback.print_exc() time.sleep(self.retry_interval_in_ms / 1000) return result def upload_file_to_kafka(self, topic, file_path, **kwargs): """ Utility function to upload contents of file to a given kafka topic :param topic: Kafka topic to which the file will be uploaded :param file_path: Absolute path of the file to be uploaded :param kwargs: append - If True, then file content will be uploaded to existing topic. If topic is not present then new one will be created. If false, and topic is not present then new topic is created. If topic is already present then error is returned. Default, async=False :return: True if content was uploaded else false """ append = kwargs.get('append', False) result = False producer = None try: if not append: # Check if topic is already present if self.kafka.has_metadata_for_topic(topic): print 'Error - Kafka topic : ' + topic + ' already present and append is : ' + str(append) return False # In case of append is True and topic already present/not present # and append is False and topic already not present if self._ensure_kafka_topic_exists(topic): producer = SimpleProducer(self.kafka, batch_send=True, batch_send_every_n=20) with open(file_path, 'rU') as fh: for line in fh: producer.send_messages(topic, line.strip()) result = True except: print 'Error - uploading file : ' + file_path + ' to topic : ' + topic finally: if producer: producer.stop() return result