diff options
Diffstat (limited to 'kafka')
| -rw-r--r-- | kafka/common.py | 6 | ||||
| -rw-r--r-- | kafka/producer.py | 29 | ||||
| -rw-r--r-- | kafka/protocol.py | 19 |
3 files changed, 33 insertions, 21 deletions
diff --git a/kafka/common.py b/kafka/common.py index d515532..209754b 100644 --- a/kafka/common.py +++ b/kafka/common.py @@ -170,6 +170,11 @@ class ConsumerNoMoreData(KafkaError): class ProtocolError(KafkaError): pass + +class UnsupportedCodecError(KafkaError): + pass + + kafka_errors = { -1 : UnknownError, 1 : OffsetOutOfRangeError, @@ -187,6 +192,7 @@ kafka_errors = { 13 : StaleLeaderEpochCodeError, } + def check_error(response): error = kafka_errors.get(response.error) if error: diff --git a/kafka/producer.py b/kafka/producer.py index 9ecb341..8e40be5 100644 --- a/kafka/producer.py +++ b/kafka/producer.py @@ -9,12 +9,11 @@ from collections import defaultdict from itertools import cycle from multiprocessing import Queue, Process -from kafka.common import ProduceRequest, TopicAndPartition -from kafka.partitioner import HashedPartitioner -from kafka.protocol import ( - CODEC_NONE, CODEC_GZIP, CODEC_SNAPPY, ALL_CODECS, - create_message, create_gzip_message, create_snappy_message, +from kafka.common import ( + ProduceRequest, TopicAndPartition, UnsupportedCodecError ) +from kafka.partitioner import HashedPartitioner +from kafka.protocol import CODEC_NONE, ALL_CODECS, create_message_set log = logging.getLogger("kafka") @@ -66,13 +65,7 @@ def _send_upstream(queue, client, codec, batch_time, batch_size, # Send collected requests upstream reqs = [] for topic_partition, msg in msgset.items(): - if codec == CODEC_GZIP: - messages = [create_gzip_message(msg)] - elif codec == CODEC_SNAPPY: - messages = [create_snappy_message(msg)] - else: - messages = [create_message(m) for m in msg] - + messages = create_message_set(msg, codec) req = ProduceRequest(topic_partition.topic, topic_partition.partition, messages) @@ -132,7 +125,9 @@ class Producer(object): if codec is None: codec = CODEC_NONE - assert codec in ALL_CODECS + elif codec not in ALL_CODECS: + raise UnsupportedCodecError("Codec 0x%02x unsupported" % codec) + self.codec = codec if self.async: @@ -159,13 +154,7 @@ class Producer(object): self.queue.put((TopicAndPartition(topic, partition), m)) resp = [] else: - if self.codec == CODEC_GZIP: - messages = [create_gzip_message(msg)] - elif self.codec == CODEC_SNAPPY: - messages = [create_snappy_message(msg)] - else: - messages = [create_message(m) for m in msg] - + messages = create_message_set(msg, self.codec) req = ProduceRequest(topic, partition, messages) try: resp = self.client.send_produce_request([req], acks=self.req_acks, diff --git a/kafka/protocol.py b/kafka/protocol.py index 730ae6b..58661c7 100644 --- a/kafka/protocol.py +++ b/kafka/protocol.py @@ -9,7 +9,8 @@ from kafka.common import ( BrokerMetadata, PartitionMetadata, Message, OffsetAndMessage, ProduceResponse, FetchResponse, OffsetResponse, OffsetCommitResponse, OffsetFetchResponse, ProtocolError, - BufferUnderflowError, ChecksumError, ConsumerFetchSizeTooSmall + BufferUnderflowError, ChecksumError, ConsumerFetchSizeTooSmall, + UnsupportedCodecError ) from kafka.util import ( read_short_string, read_int_string, relative_unpack, @@ -568,3 +569,19 @@ def create_snappy_message(payloads, key=None): codec = ATTRIBUTE_CODEC_MASK & CODEC_SNAPPY return Message(0, 0x00 | codec, key, snapped) + + +def create_message_set(messages, codec=CODEC_NONE): + """Create a message set using the given codec. + + If codec is CODEC_NONE, return a list of raw Kafka messages. Otherwise, + return a list containing a single codec-encoded message. + """ + if codec == CODEC_NONE: + return [create_message(m) for m in messages] + elif codec == CODEC_GZIP: + return [create_gzip_message(messages)] + elif codec == CODEC_SNAPPY: + return [create_snappy_message(messages)] + else: + raise UnsupportedCodecError("Codec 0x%02x unsupported" % codec) |