support seek and fix tests

Author: BewareMyPower

pulsar/asyncio.py

@@ -235,6 +235,36 @@ async def unsubscribe(self) -> None:
         self._consumer.unsubscribe_async(functools.partial(_set_future, future, value=None))
         await future
 
+    async def seek(self, messageid: Union[pulsar.MessageId, int]) -> None:
+        """
+        Reset the subscription associated with this consumer to a specific
+        message id or publish timestamp asynchronously.
+
+        The message id can either be a specific message or represent the first
+        or last messages in the topic.
+
+        Parameters
+        ----------
+        messageid : MessageId or int
+            The message id for seek, OR an integer event time (timestamp) to
+            seek to
+
+        Raises
+        ------
+        PulsarException
+        """
+        future = asyncio.get_running_loop().create_future()
+        if isinstance(messageid, pulsar.MessageId):
+            msg_id = messageid._msg_id
+        elif isinstance(messageid, int):
+            msg_id = messageid
+        else:
+            raise ValueError(f"invalid messageid type {type(messageid)}")
+        self._consumer.seek_async(
+            msg_id, functools.partial(_set_future, future, value=None)
+        )
+        await future
+
     async def close(self) -> None:
         """
         Close the consumer asynchronously.
@@ -286,7 +316,7 @@ async def create_producer(self, topic: str,
                               max_pending_messages: int = 1000,
                               max_pending_messages_across_partitions: int = 50000,
                               block_if_queue_full: bool = False,
-                              batching_enabled: bool = False,
+                              batching_enabled: bool = True,
                               batching_max_messages: int = 1000,
                               batching_max_allowed_size_in_bytes: int = 128*1024,
                               batching_max_publish_delay_ms: int = 10,

src/consumer.cc

@@ -142,6 +142,16 @@ void Consumer_unsubscribeAsync(Consumer& consumer, ResultCallback callback) {
     consumer.unsubscribeAsync(callback);
 }
 
+void Consumer_seekAsync(Consumer& consumer, const MessageId& msgId, ResultCallback callback) {
+    py::gil_scoped_release release;
+    consumer.seekAsync(msgId, callback);
+}
+
+void Consumer_seekAsync_timestamp(Consumer& consumer, uint64_t timestamp, ResultCallback callback) {
+    py::gil_scoped_release release;
+    consumer.seekAsync(timestamp, callback);
+}
+
 void export_consumer(py::module_& m) {
     py::class_<Consumer>(m, "Consumer")
         .def(py::init<>())
@@ -173,5 +183,7 @@ void export_consumer(py::module_& m) {
         .def("acknowledge_cumulative_async", &Consumer_acknowledgeCumulativeAsync)
         .def("acknowledge_cumulative_async", &Consumer_acknowledgeCumulativeAsync_message_id)
         .def("close_async", &Consumer_closeAsync)
-        .def("unsubscribe_async", &Consumer_unsubscribeAsync);
+        .def("unsubscribe_async", &Consumer_unsubscribeAsync)
+        .def("seek_async", &Consumer_seekAsync)
+        .def("seek_async", &Consumer_seekAsync_timestamp);
 }

tests/asyncio_test.py

@@ -19,6 +19,7 @@
 #
 
 import asyncio
+from typing import List
 import pulsar
 import time
 from pulsar.asyncio import (
@@ -99,9 +100,12 @@ async def test_close_producer(self):
         except PulsarException as e:
             self.assertEqual(e.error(), pulsar.Result.AlreadyClosed)
 
-    async def _prepare_messages(self, producer: Producer):
+    async def _prepare_messages(self, producer: Producer) -> List[pulsar.MessageId]:
+        msg_ids = []
         for i in range(5):
-            await producer.send(f'msg-{i}'.encode())
+            msg_id = await producer.send(f'msg-{i}'.encode())
+            msg_ids.append(msg_id)
+        return msg_ids
 
     async def test_consumer_cumulative_acknowledge(self):
         topic = f'asyncio-test-consumer-cumulative-ack-{time.time()}'
@@ -185,6 +189,55 @@ async def test_unsubscribe(self):
         await consumer.unsubscribe()
         consumer = await self._client.subscribe(topic, sub)
 
+    async def test_seek_message_id(self):
+        topic = f'asyncio-test-seek-message-id-{time.time()}'
+        sub = 'sub'
+        consumer = await self._client.subscribe(
+            topic, sub, initial_position=pulsar.InitialPosition.Earliest
+        )
+
+        producer = await self._client.create_producer(topic)
+        msg_ids = await self._prepare_messages(producer)
+
+        for i in range(5):
+            msg = await consumer.receive()
+            self.assertEqual(msg.data(), f'msg-{i}'.encode())
+
+        await consumer.seek(msg_ids[2])
+
+        msg = await consumer.receive()
+        self.assertEqual(msg.data(), b'msg-3')
+
+    async def test_seek_timestamp(self):
+        topic = f'asyncio-test-seek-timestamp-{time.time()}'
+        sub = 'sub'
+        consumer = await self._client.subscribe(
+            topic, sub, initial_position=pulsar.InitialPosition.Earliest
+        )
+
+        producer = await self._client.create_producer(topic)
+
+        # Send first 3 messages
+        for i in range(3):
+            await producer.send(f'msg-{i}'.encode())
+
+        seek_time = int(time.time() * 1000)
+
+        # Send 2 more messages
+        for i in range(3, 5):
+            await producer.send(f'msg-{i}'.encode())
+
+        # Consume all messages first
+        for i in range(5):
+            msg = await consumer.receive()
+            self.assertEqual(msg.data(), f'msg-{i}'.encode())
+
+        # Seek to the timestamp (should start from msg-3)
+        await consumer.seek(seek_time)
+
+        msg = await consumer.receive()
+        self.assertEqual(msg.data(), b'msg-3')
+
 
 if __name__ == '__main__':
     main()