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
