openstack/deb-python-kafka
Code Issues Proposed changes

Improve error handling and tests w.r.t. codecs

Browse Source 
Add function kafka.protocol.create_message_set() that takes a list of
payloads and a codec and returns a message set with the desired encoding.

Introduce kafka.common.UnsupportedCodecError, raised if an unknown codec
is specified.

Include a test for the new function.
This commit is contained in:
Patrick Lucas
2014-05-07 10:02:57 -07:00
parent 39796ec491
commit 805b52a34d
4 changed files with 91 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
kafka/common.py
View File

@@ -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:

29
kafka/producer.py
View File

@@ -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,

19
kafka/protocol.py
View File

@@ -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)

64
test/test_protocol.py
View File

@@ -1,23 +1,30 @@
import contextlib
from contextlib import contextmanager
import struct
import unittest2
import mock
from mock import sentinel
from kafka import KafkaClient
from kafka.common import (
OffsetRequest, OffsetCommitRequest, OffsetFetchRequest,
OffsetResponse, OffsetCommitResponse, OffsetFetchResponse,
ProduceRequest, FetchRequest, Message, ChecksumError,
ConsumerFetchSizeTooSmall, ProduceResponse, FetchResponse,
OffsetAndMessage, BrokerMetadata, PartitionMetadata,
TopicAndPartition, KafkaUnavailableError, ProtocolError,
LeaderUnavailableError, PartitionUnavailableError
ConsumerFetchSizeTooSmall, ProduceResponse, FetchResponse, OffsetAndMessage,
BrokerMetadata, PartitionMetadata, TopicAndPartition, KafkaUnavailableError,
ProtocolError, LeaderUnavailableError, PartitionUnavailableError,
UnsupportedCodecError
)
from kafka.codec import (
has_snappy, gzip_encode, gzip_decode,
snappy_encode, snappy_decode
)
import kafka.protocol
from kafka.protocol import (
create_gzip_message, create_message, create_snappy_message, KafkaProtocol,
ATTRIBUTE_CODEC_MASK, CODEC_GZIP, CODEC_SNAPPY
ATTRIBUTE_CODEC_MASK, CODEC_NONE, CODEC_GZIP, CODEC_SNAPPY, KafkaProtocol,
create_message, create_gzip_message, create_snappy_message,
create_message_set
)
class TestProtocol(unittest2.TestCase):
@@ -691,3 +698,48 @@ class TestProtocol(unittest2.TestCase):
OffsetFetchResponse(topic = 'topic1', partition = 2, offset = 4, error = 0, metadata = "meta"),
OffsetFetchResponse(topic = 'topic1', partition = 4, offset = 8, error = 0, metadata = "meta"),
]))
@contextmanager
def mock_create_message_fns(self):
patches = contextlib.nested(
mock.patch.object(kafka.protocol, "create_message",
return_value=sentinel.message),
mock.patch.object(kafka.protocol, "create_gzip_message",
return_value=sentinel.gzip_message),
mock.patch.object(kafka.protocol, "create_snappy_message",
return_value=sentinel.snappy_message),
)
with patches:
yield
def test_create_message_set(self):
messages = [1, 2, 3]
# Default codec is CODEC_NONE. Expect list of regular messages.
expect = [sentinel.message] * len(messages)
with self.mock_create_message_fns():
message_set = create_message_set(messages)
self.assertEqual(message_set, expect)
# CODEC_NONE: Expect list of regular messages.
expect = [sentinel.message] * len(messages)
with self.mock_create_message_fns():
message_set = create_message_set(messages, CODEC_NONE)
self.assertEqual(message_set, expect)
# CODEC_GZIP: Expect list of one gzip-encoded message.
expect = [sentinel.gzip_message]
with self.mock_create_message_fns():
message_set = create_message_set(messages, CODEC_GZIP)
self.assertEqual(message_set, expect)
# CODEC_SNAPPY: Expect list of one snappy-encoded message.
expect = [sentinel.snappy_message]
with self.mock_create_message_fns():
message_set = create_message_set(messages, CODEC_SNAPPY)
self.assertEqual(message_set, expect)
# Unknown codec should raise UnsupportedCodecError.
with self.assertRaises(UnsupportedCodecError):
create_message_set(messages, -1)