Implement pause/resume support (closes confluentinc#120, @dangra)

dangra · edenhill · commit 1dd2e1e74a36 · 2018-01-16T10:02:40.000+01:00
* Implement pause/resume support. Closes confluentinc#120 * Add basic test case for pause/resume methods * fix typo in method documentation * do not leak c_parts on pause/resume
diff --git a/confluent_kafka/src/Consumer.c b/confluent_kafka/src/Consumer.c
@@ -649,6 +649,55 @@ static PyObject *Consumer_position (Handle *self, PyObject *args,
 	return plist;
 }
 
+static PyObject *Consumer_pause(Handle *self, PyObject *args,
+                    PyObject *kwargs) {
+
+    PyObject *plist;
+	rd_kafka_topic_partition_list_t *c_parts;
+    rd_kafka_resp_err_t err;
+    static char *kws[] = {"partitions", NULL};
+
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O", kws, &plist))
+        return NULL;
+
+    if (!(c_parts = py_to_c_parts(plist)))
+        return NULL;
+
+    err = rd_kafka_pause_partitions(self->rk, c_parts);
+    rd_kafka_topic_partition_list_destroy(c_parts);
+    if (err) {
+        cfl_PyErr_Format(err,
+                "Failed to pause partitions: %s",
+                rd_kafka_err2str(err));
+        return NULL;
+    }
+	Py_RETURN_NONE;
+}
+
+static PyObject *Consumer_resume (Handle *self, PyObject *args,
+                    PyObject *kwargs) {
+
+    PyObject *plist;
+	rd_kafka_topic_partition_list_t *c_parts;
+    rd_kafka_resp_err_t err;
+    static char *kws[] = {"partitions", NULL};
+
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O", kws, &plist))
+        return NULL;
+
+    if (!(c_parts = py_to_c_parts(plist)))
+        return NULL;
+
+    err = rd_kafka_resume_partitions(self->rk, c_parts);
+    rd_kafka_topic_partition_list_destroy(c_parts);
+    if (err) {
+        cfl_PyErr_Format(err,
+                "Failed to resume partitions: %s",
+                rd_kafka_err2str(err));
+        return NULL;
+    }
+	Py_RETURN_NONE;
+}
 
 static PyObject *Consumer_get_watermark_offsets (Handle *self, PyObject *args,
                                                  PyObject *kwargs) {
@@ -1063,6 +1112,30 @@ static PyMethodDef Consumer_methods[] = {
       "  :raises: RuntimeError if called on a closed consumer\n"
 	  "\n"
 	},
+	{ "pause", (PyCFunction)Consumer_pause,
+	  METH_VARARGS|METH_KEYWORDS,
+	  ".. py:function:: pause(partitions)\n"
+	  "\n"
+	  "  Pause consumption for the provided list of partitions.\n"
+	  "\n"
+	  "  :param list(TopicPartition) partitions: List of topic+partitions "
+	  "to pause.\n"
+	  "  :rtype: None\n"
+	  "  :raises: KafkaException\n"
+	  "\n"
+	},
+	{ "resume", (PyCFunction)Consumer_resume,
+	  METH_VARARGS|METH_KEYWORDS,
+	  ".. py:function:: resume(partitions)\n"
+	  "\n"
+	  "  Resume consumption for the provided list of partitions.\n"
+	  "\n"
+	  "  :param list(TopicPartition) partitions: List of topic+partitions "
+	  "to resume.\n"
+	  "  :rtype: None\n"
+	  "  :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"
diff --git a/tests/test_Consumer.py b/tests/test_Consumer.py
@@ -59,6 +59,12 @@ def dummy_assign_revoke(consumer, partitions):
     assignment = kc.assignment()
     assert partitions == assignment
 
+    # Pause partitions
+    kc.pause(partitions)
+
+    # Resume partitions
+    kc.resume(partitions)
+
     # Get cached watermarks, should all be invalid.
     lo, hi = kc.get_watermark_offsets(partitions[0], cached=True)
     assert lo == -1001 and hi == -1001
