Added Consumer.get_watermark_offsets() (#31)

Author: edenhill

confluent_kafka/src/Consumer.c

@@ -381,6 +381,54 @@ static PyObject *Consumer_position (Handle *self, PyObject *args,
 }
 
 
+static PyObject *Consumer_get_watermark_offsets (Handle *self, PyObject *args,
+						 PyObject *kwargs) {
+
+	TopicPartition *tp;
+	rd_kafka_resp_err_t err;
+	double tmout = -1.0f;
+	int cached = 0;
+	int64_t low = RD_KAFKA_OFFSET_INVALID, high = RD_KAFKA_OFFSET_INVALID;
+	static char *kws[] = { "partition", "timeout", "cached", NULL };
+	PyObject *rlist;
+
+	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|db", kws,
+					 (PyObject **)&tp, &tmout, &cached))
+		return NULL;
+
+
+	if (PyObject_Type((PyObject *)tp) != (PyObject *)&TopicPartitionType) {
+		PyErr_Format(PyExc_TypeError,
+			     "expected %s", TopicPartitionType.tp_name);
+		return NULL;
+	}
+
+	if (cached) {
+		err = rd_kafka_get_watermark_offsets(self->rk,
+						     tp->topic, tp->partition,
+						     &low, &high);
+	} else {
+		err = rd_kafka_query_watermark_offsets(self->rk,
+						       tp->topic, tp->partition,
+						       &low, &high,
+						       tmout >= 0 ? (int)(tmout * 1000.0f) : -1);
+	}
+
+	if (err) {
+		cfl_PyErr_Format(err,
+				 "Failed to get watermark offsets: %s",
+				 rd_kafka_err2str(err));
+		return NULL;
+	}
+
+	rlist = PyList_New(2);
+	PyList_SetItem(rlist, 0, PyLong_FromLongLong(low));
+	PyList_SetItem(rlist, 1, PyLong_FromLongLong(high));
+
+	return rlist;
+}
+
+
 
 static PyObject *Consumer_poll (Handle *self, PyObject *args,
 				    PyObject *kwargs) {
@@ -541,6 +589,22 @@ static PyMethodDef Consumer_methods[] = {
 	  "  :raises: KafkaException\n"
 	  "\n"
 	},
+	{ "get_watermark_offsets", (PyCFunction)Consumer_get_watermark_offsets,
+	  METH_VARARGS|METH_KEYWORDS,
+	  ".. py:function:: get_watermark_offsets(partition, [timeout=None], [cached=False])\n"
+	  "\n"
+	  "  Retrieve low and high offsets for partition.\n"
+	  "\n"
+	  "  :param TopicPartition partition: Topic+partition to return offsets for."
+	  "  :param float timeout: Request timeout (when cached=False).\n"
+	  "  :param bool cached: Instead of querying the broker used cached information. "
+	  "The low offset is updated periodically (if statistics.interval.ms is set) while "
+	  "the high offset is updated on each message fetched from the broker for this partition."
+	  "  :returns: List of [low,high] on success or None on timeout.\n"
+	  "  :rtype: list(int)\n"
+	  "  :raises: KafkaException\n"
+	  "\n"
+	},
 	{ "close", (PyCFunction)Consumer_close, METH_NOARGS,
 	  "\n"
 	  "  Close down and terminate the Kafka Consumer.\n"

confluent_kafka/src/confluent_kafka.c

@@ -510,15 +510,6 @@ PyObject *Message_new0 (const rd_kafka_message_t *rkm) {
  *
  *
  ****************************************************************************/
-typedef struct {
-	PyObject_HEAD
-	char *topic;
-	int   partition;
-	int64_t offset;
-	PyObject *error;
-} TopicPartition;
-
-
 static int TopicPartition_clear (TopicPartition *self) {
 	if (self->topic) {
 		free(self->topic);
@@ -600,8 +591,6 @@ static PyObject *TopicPartition_str0 (TopicPartition *self) {
 }
 
 
-static PyTypeObject TopicPartitionType;
-
 static PyObject *
 TopicPartition_richcompare (TopicPartition *self, PyObject *o2,
 			    int op) {
@@ -658,7 +647,7 @@ static long TopicPartition_hash (TopicPartition *self) {
 }
 
 
-static PyTypeObject TopicPartitionType = {
+PyTypeObject TopicPartitionType = {
 	PyVarObject_HEAD_INIT(NULL, 0)
 	"cimpl.TopicPartition",         /*tp_name*/
 	sizeof(TopicPartition),       /*tp_basicsize*/

confluent_kafka/src/confluent_kafka.h

@@ -181,6 +181,26 @@ void CallState_resume (CallState *cs);
 void CallState_crash (CallState *cs);
 
 
+/****************************************************************************
+ *
+ *
+ * TopicPartition
+ *
+ *
+ *
+ *
+ ****************************************************************************/
+typedef struct {
+	PyObject_HEAD
+	char *topic;
+	int   partition;
+	int64_t offset;
+	PyObject *error;
+} TopicPartition;
+
+extern PyTypeObject TopicPartitionType;
+
+
 /****************************************************************************
  *
  *

examples/integration_test.py

@@ -258,6 +258,18 @@ def verify_consumer():
             print('max_msgcnt %d reached' % msgcnt)
             break
 
+    # Get current assignment
+    assignment = c.assignment()
+
+    # Get cached watermark offsets
+    # Since we're not making use of statistics the low offset is not known so ignore it.
+    lo,hi = c.get_watermark_offsets(assignment[0], cached=True)
+    print('Cached offsets for %s: %d - %d' % (assignment[0], lo, hi))
+
+    # Query broker for offsets
+    lo,hi = c.get_watermark_offsets(assignment[0], timeout=1.0)
+    print('Queried offsets for %s: %d - %d' % (assignment[0], lo, hi))
+
 
     # Close consumer
     c.close()

tests/test_Consumer.py

@@ -43,6 +43,16 @@ def dummy_assign_revoke (consumer, partitions):
     assignment = kc.assignment()
     assert partitions == assignment
 
+    # Get cached watermarks, should all be invalid.
+    lo, hi = kc.get_watermark_offsets(partitions[0], cached=True)
+    assert lo == -1001 and hi == -1001
+
+    # Query broker for watermarks, should raise an exception.
+    try:
+        lo, hi = kc.get_watermark_offsets(partitions[0], timeout=0.5, cached=False)
+    except KafkaException as e:
+        assert e.args[0].code() in (KafkaError._TIMED_OUT, KafkaError._WAIT_COORD)
+
     kc.unassign()
 
     kc.commit(async=True)