Added optional argument key to create_message_set function

1 files changed, 4 insertions, 4 deletions

diff --git a/kafka/protocol.py b/kafka/protocol.py
index 266e963..13b973e 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)