diff --git a/kafka/protocol.py b/kafka/protocol.py index 266e9638e..13b973e03 100644 --- a/kafka/protocol.py +++ b/kafka/protocol.py @@ -597,17 +597,17 @@ def create_snappy_message(payloads, key=None): return Message(0, 0x00 | codec, key, snapped) -def create_message_set(messages, codec=CODEC_NONE): +def create_message_set(messages, codec=CODEC_NONE, key=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] + return [create_message(m, key) for m in messages] elif codec == CODEC_GZIP: - return [create_gzip_message(messages)] + return [create_gzip_message(messages, key)] elif codec == CODEC_SNAPPY: - return [create_snappy_message(messages)] + return [create_snappy_message(messages, key)] else: raise UnsupportedCodecError("Codec 0x%02x unsupported" % codec)