Add Consumer.seek()

edenhill · edenhill · commit 05be02c8f61a · 2018-02-12T12:25:43.000+01:00
diff --git a/confluent_kafka/src/Consumer.c b/confluent_kafka/src/Consumer.c
@@ -699,6 +699,55 @@ static PyObject *Consumer_resume (Handle *self, PyObject *args,
 	Py_RETURN_NONE;
 }
 
+
+static PyObject *Consumer_seek (Handle *self, PyObject *args, PyObject *kwargs) {
+
+        TopicPartition *tp;
+        rd_kafka_resp_err_t err;
+        static char *kws[] = { "partition", NULL };
+        rd_kafka_topic_t *rkt;
+
+        if (!self->rk) {
+                PyErr_SetString(PyExc_RuntimeError, "Consumer closed");
+                return NULL;
+        }
+
+        if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O", kws,
+                                         (PyObject **)&tp))
+                return NULL;
+
+
+        if (PyObject_Type((PyObject *)tp) != (PyObject *)&TopicPartitionType) {
+                PyErr_Format(PyExc_TypeError,
+                             "expected %s", TopicPartitionType.tp_name);
+                return NULL;
+        }
+
+        rkt = rd_kafka_topic_new(self->rk, tp->topic, NULL);
+        if (!rkt) {
+                cfl_PyErr_Format(rd_kafka_last_error(),
+                                 "Failed to get topic object for "
+                                 "topic \"%s\": %s",
+                                 tp->topic,
+                                 rd_kafka_err2str(rd_kafka_last_error()));
+                return NULL;
+        }
+
+        err = rd_kafka_seek(rkt, tp->partition, tp->offset, -1);
+
+        rd_kafka_topic_destroy(rkt);
+
+        if (err) {
+                cfl_PyErr_Format(err,
+                                 "Failed to seek to offset %"PRId64": %s",
+                                 tp->offset, rd_kafka_err2str(err));
+                return NULL;
+        }
+
+        Py_RETURN_NONE;
+}
+
+
 static PyObject *Consumer_get_watermark_offsets (Handle *self, PyObject *args,
                                                  PyObject *kwargs) {
 
@@ -1136,6 +1185,24 @@ static PyMethodDef Consumer_methods[] = {
 	  "  :raises: KafkaException\n"
 	  "\n"
 	},
+        { "seek", (PyCFunction)Consumer_seek,
+          METH_VARARGS|METH_KEYWORDS,
+          ".. py:function:: seek(partition)\n"
+          "\n"
+          "  Set consume position for partition to offset.\n"
+          "  The offset may be an absolute (>=0) or a\n"
+          "  logical offset (:py:const:`OFFSET_BEGINNING` et.al).\n"
+          "\n"
+          "  seek() may only be used to update the consume offset of an\n"
+          "  actively consumed partition (i.e., after :py:const:`assign()`),\n"
+          "  to set the starting offset of partition not being consumed instead\n"
+          "  pass the offset in an `assign()` call.\n"
+          "\n"
+          "  :param TopicPartition partition: Topic+partition+offset to seek to.\n"
+          "\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/examples/integration_test.py b/examples/integration_test.py
@@ -402,6 +402,32 @@ def print_commit_result(err, partitions):
         print('# Committed offsets for: %s' % partitions)
 
 
+def verify_consumer_seek(c, seek_to_msg):
+    """ Seek to message and verify the next consumed message matches.
+        Must only be performed on an actively consuming consumer. """
+
+    tp = confluent_kafka.TopicPartition(seek_to_msg.topic(),
+                                        seek_to_msg.partition(),
+                                        seek_to_msg.offset())
+    print('seek: Seeking to %s' % tp)
+    c.seek(tp)
+
+    while True:
+        msg = c.poll()
+        assert msg is not None
+        if msg.error():
+            print('seek: Ignoring non-message: %s' % msg)
+            continue
+
+        if msg.topic() != seek_to_msg.topic() or msg.partition() != seek_to_msg.partition():
+            continue
+
+        print('seek: message at offset %d' % msg.offset())
+        assert msg.offset() == seek_to_msg.offset(), \
+            'expected message at offset %d, not %d' % (seek_to_msg.offset(), msg.offset())
+        break
+
+
 def verify_consumer():
     """ Verify basic Consumer functionality """
 
@@ -433,6 +459,8 @@ def print_wmark(consumer, parts):
     max_msgcnt = 100
     msgcnt = 0
 
+    first_msg = None
+
     while True:
         # Consume until EOF or error
 
@@ -455,6 +483,9 @@ def print_wmark(consumer, parts):
               (msg.topic(), msg.partition(), msg.offset(),
                msg.key(), msg.value(), tstype, timestamp))
 
+        if first_msg is None:
+            first_msg = msg
+
         if (msgcnt == 11):
             parts = c.assignment()
             print('Pausing partitions briefly')
@@ -499,6 +530,8 @@ def print_wmark(consumer, parts):
     offsets = c.offsets_for_times(topic_partions_to_search, timeout=1.0)
     print("offsets_for_times results: %s" % offsets)
 
+    verify_consumer_seek(c, first_msg)
+
     # Close consumer
     c.close()
 
diff --git a/tests/test_Consumer.py b/tests/test_Consumer.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
-from confluent_kafka import (Consumer, TopicPartition, KafkaError, KafkaException, TIMESTAMP_NOT_AVAILABLE,
+from confluent_kafka import (Consumer, TopicPartition, KafkaError,
+                             KafkaException, TIMESTAMP_NOT_AVAILABLE,
                              OFFSET_INVALID, libversion)
 import pytest
 
@@ -55,6 +56,10 @@ def dummy_assign_revoke(consumer, partitions):
     partitions = list(map(lambda part: TopicPartition("test", part), range(0, 100, 3)))
     kc.assign(partitions)
 
+    with pytest.raises(KafkaException) as ex:
+        kc.seek(TopicPartition("test", 0, 123))
+    assert 'Erroneous state' in str(ex.value)
+
     # Verify assignment
     assignment = kc.assignment()
     assert partitions == assignment
@@ -267,6 +272,10 @@ def test_any_method_after_close_throws_exception():
         c.position([TopicPartition("test", 0)])
     assert 'Consumer closed' == str(ex.value)
 
+    with pytest.raises(RuntimeError) as ex:
+        c.seek([TopicPartition("test", 0, 0)])
+    assert 'Consumer closed' == str(ex.value)
+
     with pytest.raises(RuntimeError) as ex:
         lo, hi = c.get_watermark_offsets(TopicPartition("test", 0))
     assert 'Consumer closed' == str(ex.value)
