Message header support (#299, @johnistan)

* first stab at header support

* Initial support for message headers

* Fix temp testing changes

* Fix temp testing changes

* Add method doc

* add error handling and remove debug statements

* responding to PR comments and fixed up tests

* responding to PR comments and fixed up tests

* Responding to requests

* Fix c_headers destroy and added tests for set_headers

* use new detach header api

* moves detach headers to consumer code to keep callback consistent. Message headers are not supported in delivery callbacks. Adds documentation to report this limitation

* fix consumer.consume header support

* adds ifdef gaurds for older versions of librdkafka

* update rd_kafka_producev to handle timestamp support but not header support

* cleanup pep8 error

* responded to PR comments

Author: johnistan

confluent_kafka/src/Consumer.c

@@ -885,6 +885,11 @@ static PyObject *Consumer_poll (Handle *self, PyObject *args,
                 Py_RETURN_NONE;
 
         msgobj = Message_new0(self, rkm);
+#ifdef RD_KAFKA_V_HEADERS
+        // Have to deatch headers outside Message_new0 because it declares the
+        // rk message as a const
+        rd_kafka_message_detach_headers(rkm, &((Message *)msgobj)->c_headers);
+#endif
         rd_kafka_message_destroy(rkm);
 
         return msgobj;
@@ -946,6 +951,11 @@ static PyObject *Consumer_consume (Handle *self, PyObject *args,
 
         for (i = 0; i < n; i++) {
                 PyObject *msgobj = Message_new0(self, rkmessages[i]);
+#ifdef RD_KAFKA_V_HEADERS
+                // Have to deatch headers outside Message_new0 because it declares the
+                // rk message as a const
+                rd_kafka_message_detach_headers(rkmessages[i], &((Message *)msgobj)->c_headers);
+#endif
                 PyList_SET_ITEM(msglist, i, msgobj);
                 rd_kafka_message_destroy(rkmessages[i]);
         }

confluent_kafka/src/Producer.c

@@ -258,7 +258,11 @@ Producer_producev (Handle *self,
                    const char *topic, int32_t partition,
                    const void *value, size_t value_len,
                    const void *key, size_t key_len,
-                   void *opaque, int64_t timestamp) {
+                   void *opaque, int64_t timestamp
+#if RD_KAFKA_V_HEADERS
+                   ,rd_kafka_headers_t *headers
+#endif
+                   ) {
 
         return rd_kafka_producev(self->rk,
                                  RD_KAFKA_V_MSGFLAGS(RD_KAFKA_MSG_F_COPY),
@@ -268,6 +272,9 @@ Producer_producev (Handle *self,
                                  RD_KAFKA_V_VALUE((void *)value,
                                                   (size_t)value_len),
                                  RD_KAFKA_V_TIMESTAMP(timestamp),
+#if RD_KAFKA_V_HEADERS
+                                 RD_KAFKA_V_HEADERS(headers),
+#endif
                                  RD_KAFKA_V_OPAQUE(opaque),
                                  RD_KAFKA_V_END);
 }
@@ -302,27 +309,32 @@ static PyObject *Producer_produce (Handle *self, PyObject *args,
 	const char *topic, *value = NULL, *key = NULL;
 	int value_len = 0, key_len = 0;
 	int partition = RD_KAFKA_PARTITION_UA;
-	PyObject *dr_cb = NULL, *dr_cb2 = NULL, *partitioner_cb = NULL;
+	PyObject *headers = NULL, *dr_cb = NULL, *dr_cb2 = NULL, *partitioner_cb = NULL;
         long long timestamp = 0;
         rd_kafka_resp_err_t err;
 	struct Producer_msgstate *msgstate;
+#ifdef RD_KAFKA_V_HEADERS
+    rd_kafka_headers_t *rd_headers = NULL;
+#endif
+
 	static char *kws[] = { "topic",
 			       "value",
 			       "key",
 			       "partition",
 			       "callback",
 			       "on_delivery", /* Alias */
 			       "partitioner",
-                               "timestamp",
+                   "timestamp",
+                   "headers",
 			       NULL };
 
 	if (!PyArg_ParseTupleAndKeywords(args, kwargs,
-					 "s|z#z#iOOOL"
+					 "s|z#z#iOOOLO"
                                          , kws,
 					 &topic, &value, &value_len,
 					 &key, &key_len, &partition,
 					 &dr_cb, &dr_cb2, &partitioner_cb,
-                                         &timestamp))
+                     &timestamp, &headers))
 		return NULL;
 
 #if !HAVE_PRODUCEV
@@ -337,6 +349,24 @@ static PyObject *Producer_produce (Handle *self, PyObject *args,
         }
 #endif
 
+#ifndef RD_KAFKA_V_HEADERS
+    if (headers) {
+            PyErr_Format(PyExc_NotImplementedError,
+                         "Producer message headers requires "
+                         "confluent-kafka-python built for librdkafka "
+                         "version >=v0.11.4 (librdkafka runtime 0x%x, "
+                         "buildtime 0x%x)",
+                         rd_kafka_version(), RD_KAFKA_VERSION);
+            return NULL;
+    }
+#else
+    if (headers) {
+        if(!(rd_headers = py_headers_to_c(headers)))
+            return NULL;
+    }
+#endif
+
+
 	if (dr_cb2 && !dr_cb) /* Alias */
 		dr_cb = dr_cb2;
 
@@ -345,6 +375,7 @@ static PyObject *Producer_produce (Handle *self, PyObject *args,
 	if (!partitioner_cb || partitioner_cb == Py_None)
 		partitioner_cb = self->u.Producer.partitioner_cb;
 
+
 	/* Create msgstate if necessary, may return NULL if no callbacks
 	 * are wanted. */
 	msgstate = Producer_msgstate_new(self, dr_cb, partitioner_cb);
@@ -354,13 +385,16 @@ static PyObject *Producer_produce (Handle *self, PyObject *args,
         err = Producer_producev(self, topic, partition,
                                 value, value_len,
                                 key, key_len,
-                                msgstate, timestamp);
+                                msgstate, timestamp
+#if RD_KAFKA_V_HEADERS
+                                ,rd_headers
+#endif
+                                );
 #else
         err = Producer_produce0(self, topic, partition,
                                 value, value_len,
                                 key, key_len,
                                 msgstate);
-
 #endif
 
         if (err) {
@@ -458,6 +492,10 @@ static PyMethodDef Producer_methods[] = {
 	  "``callback`` (alias ``on_delivery``) argument to pass a function "
 	  "(or lambda) that will be called from :py:func:`poll()` when the "
 	  "message has been successfully delivered or permanently fails delivery.\n"
+      "\n"
+      "  Currently message headers are not supported on the message returned to the "
+      "callback. The ``msg.headers()`` will return None even if the original message "
+      "had headers set.\n"
 	  "\n"
 	  "  :param str topic: Topic to produce message to\n"
 	  "  :param str|bytes value: Message payload\n"

confluent_kafka/src/confluent_kafka.c

@@ -333,6 +333,34 @@ static PyObject *Message_timestamp (Message *self, PyObject *ignore) {
 			     self->timestamp);
 }
 
+static PyObject *Message_headers (Message *self, PyObject *ignore) {
+#ifdef RD_KAFKA_V_HEADERS
+	if (self->headers) {
+        Py_INCREF(self->headers);
+		return self->headers;
+    } else if (self->c_headers) {
+        self->headers = c_headers_to_py(self->c_headers);
+        rd_kafka_headers_destroy(self->c_headers);
+        self->c_headers = NULL;
+        Py_INCREF(self->headers);
+        return self->headers;
+	} else {
+		Py_RETURN_NONE;
+    }
+#else
+		Py_RETURN_NONE;
+#endif
+}
+
+static PyObject *Message_set_headers (Message *self, PyObject *new_headers) {
+   if (self->headers)
+        Py_DECREF(self->headers);
+   self->headers = new_headers;
+   Py_INCREF(self->headers);
+
+   Py_RETURN_NONE;
+}
+
 static PyObject *Message_set_value (Message *self, PyObject *new_val) {
    if (self->value)
         Py_DECREF(self->value);
@@ -409,6 +437,21 @@ static PyMethodDef Message_methods[] = {
 	  "  :rtype: (int, int)\n"
 	  "\n"
 	},
+	{ "headers", (PyCFunction)Message_headers, METH_NOARGS,
+      "  Retrieve the headers set on a message. Each header is a key value"
+      "pair. Please note that header keys are ordered and can repeat.\n"
+      "\n"
+	  "  :returns: list of two-tuples, one (key, value) pair for each header.\n"
+	  "  :rtype: [(str, bytes),...] or None.\n"
+	  "\n"
+	},
+	{ "set_headers", (PyCFunction)Message_set_headers, METH_O,
+	  "  Set the field 'Message.headers' with new value.\n"
+	  "  :param: object value: Message.headers.\n"
+	  "  :returns: None.\n"
+	  "  :rtype: None\n"
+	  "\n"
+	},
 	{ "set_value", (PyCFunction)Message_set_value, METH_O,
 	  "  Set the field 'Message.value' with new value.\n"
 	  "  :param: object value: Message.value.\n"
@@ -443,6 +486,16 @@ static int Message_clear (Message *self) {
 		Py_DECREF(self->error);
 		self->error = NULL;
 	}
+	if (self->headers) {
+		Py_DECREF(self->headers);
+		self->headers = NULL;
+	}
+#ifdef RD_KAFKA_V_HEADERS
+    if (self->c_headers){
+        rd_kafka_headers_destroy(self->c_headers);
+        self->c_headers = NULL;
+    }
+#endif
 	return 0;
 }
 
@@ -464,6 +517,8 @@ static int Message_traverse (Message *self,
 		Py_VISIT(self->key);
 	if (self->error)
 		Py_VISIT(self->error);
+	if (self->headers)
+		Py_VISIT(self->headers);
 	return 0;
 }
 
@@ -883,6 +938,81 @@ rd_kafka_topic_partition_list_t *py_to_c_parts (PyObject *plist) {
 	return c_parts;
 }
 
+#ifdef RD_KAFKA_V_HEADERS
+/**
+ * @brief Convert Python list[(header_key, header_value),...]) to C rd_kafka_topic_partition_list_t.
+ *
+ * @returns The new Python list[(header_key, header_value),...] object.
+ */
+rd_kafka_headers_t *py_headers_to_c (PyObject *headers_plist) {
+    int i, len;
+    rd_kafka_headers_t *rd_headers = NULL;
+    rd_kafka_resp_err_t err;
+    const char *header_key, *header_value = NULL;
+    int header_key_len = 0, header_value_len = 0;
+
+    len = PyList_Size(headers_plist);
+    rd_headers = rd_kafka_headers_new(len);
+
+    for (i = 0; i < len; i++) {
+
+        if(!PyArg_ParseTuple(PyList_GET_ITEM(headers_plist, i), "s#z#", &header_key,
+                &header_key_len, &header_value, &header_value_len)){
+            rd_kafka_headers_destroy(rd_headers);
+            PyErr_SetString(PyExc_TypeError,
+                    "Headers are expected to be a tuple of (key, value)");
+            return NULL;
+        }
+
+        err = rd_kafka_header_add(rd_headers, header_key, header_key_len, header_value, header_value_len);
+        if (err) {
+            rd_kafka_headers_destroy(rd_headers);
+            cfl_PyErr_Format(err,
+                     "Unable to create message headers: %s",
+                     rd_kafka_err2str(err));
+            return NULL;
+        }
+    }
+    return rd_headers;
+}
+
+/**
+ * @brief Convert rd_kafka_headers_t to Python list[(header_key, header_value),...])
+ *
+ * @returns The new C headers on success or NULL on error.
+ */
+PyObject *c_headers_to_py (rd_kafka_headers_t *headers) {
+    size_t idx = 0;
+    size_t header_size = 0;
+    const char *header_key;
+    const void *header_value;
+    size_t header_value_size;
+    PyObject *header_list;
+
+    header_size = rd_kafka_header_cnt(headers); 
+    header_list = PyList_New(header_size);
+
+    while (!rd_kafka_header_get_all(headers, idx++,
+                                     &header_key, &header_value, &header_value_size)) {
+            // Create one (key, value) tuple for each header
+            PyObject *header_tuple = PyTuple_New(2);
+            PyTuple_SetItem(header_tuple, 0,
+                cfl_PyUnistr(_FromString(header_key))
+            );
+
+            if (header_value) {
+                    PyTuple_SetItem(header_tuple, 1,
+                        cfl_PyBin(_FromStringAndSize(header_value, header_value_size))
+                    );
+            } else {
+                PyTuple_SetItem(header_tuple, 1, Py_None);
+            }
+        PyList_SET_ITEM(header_list, idx-1, header_tuple);
+    }
+
+    return header_list;
+}
+#endif
 
 /****************************************************************************
  *

confluent_kafka/src/confluent_kafka.h

@@ -267,7 +267,10 @@ rd_kafka_conf_t *common_conf_setup (rd_kafka_type_t ktype,
 PyObject *c_parts_to_py (const rd_kafka_topic_partition_list_t *c_parts);
 rd_kafka_topic_partition_list_t *py_to_c_parts (PyObject *plist);
 
-
+#ifdef RD_KAFKA_V_HEADERS
+rd_kafka_headers_t *py_headers_to_c (PyObject *headers_plist);
+PyObject *c_headers_to_py (rd_kafka_headers_t *headers);
+#endif
 /****************************************************************************
  *
  *
@@ -286,6 +289,10 @@ typedef struct {
 	PyObject *topic;
 	PyObject *value;
 	PyObject *key;
+	PyObject *headers;
+#ifdef RD_KAFKA_V_HEADERS
+	rd_kafka_headers_t *c_headers;
+#endif
 	PyObject *error;
 	int32_t partition;
 	int64_t offset;

docs/index.rst

@@ -101,6 +101,9 @@ The Python bindings also provide some additional configuration properties:
   delivery result (success or failure).
   This property may also be set per-message by passing ``callback=callable``
   (or ``on_delivery=callable``) to the confluent_kafka.Producer.produce() function.
+  Currently message headers are not supported on the message returned to the
+  callback. The ``msg.headers()`` will return None even if the original message
+  had headers set.
 
 * ``on_commit(kafka.KafkaError, list(kafka.TopicPartition))`` (**Consumer**): Callback used to indicate success or failure
   of commit requests.

examples/integration_test.py

@@ -117,11 +117,14 @@ def verify_producer():
     p = confluent_kafka.Producer(**conf)
     print('producer at %s' % p)
 
+    headers = [('foo1', 'bar'), ('foo1', 'bar2'), ('foo2', b'1')]
+
     # Produce some messages
-    p.produce(topic, 'Hello Python!')
+    p.produce(topic, 'Hello Python!', headers=headers)
+    p.produce(topic, key='Just a key and headers', headers=headers)
     p.produce(topic, key='Just a key')
     p.produce(topic, partition=1, value='Strictly for partition 1',
-              key='mykey')
+              key='mykey', headers=headers)
 
     # Produce more messages, now with delivery report callbacks in various forms.
     mydr = MyTestDr()
@@ -479,9 +482,16 @@ def print_wmark(consumer, parts):
                 break
 
         tstype, timestamp = msg.timestamp()
-        print('%s[%d]@%d: key=%s, value=%s, tstype=%d, timestamp=%s' %
+        headers = msg.headers()
+        if headers:
+            example_header = headers
+
+        msg.set_headers([('foo', 'bar')])
+        assert msg.headers() == [('foo', 'bar')]
+
+        print('%s[%d]@%d: key=%s, value=%s, tstype=%d, timestamp=%s headers=%s' %
               (msg.topic(), msg.partition(), msg.offset(),
-               msg.key(), msg.value(), tstype, timestamp))
+               msg.key(), msg.value(), tstype, timestamp, headers))
 
         if first_msg is None:
             first_msg = msg
@@ -511,6 +521,9 @@ def print_wmark(consumer, parts):
             print('max_msgcnt %d reached' % msgcnt)
             break
 
+    assert example_header, "We should have received at least one header"
+    assert example_header == [(u'foo1', 'bar'), (u'foo1', 'bar2'), (u'foo2', '1')]
+
     # Get current assignment
     assignment = c.assignment()