diff options
Diffstat (limited to 'kafka/protocol.py')
| -rw-r--r-- | kafka/protocol.py | 19 |
1 files changed, 18 insertions, 1 deletions
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) |