feat(asyncio): add Reader API

pulsar/asyncio.py

@@ -449,6 +449,107 @@ def consumer_name(self) -> str:
  """
  return self._consumer.consumer_name()
 
+class Reader:
+ """
+ The Pulsar topic reader, used to read messages from a topic.
+ """
+
+ def __init__(self, reader: _pulsar.Reader, schema: pulsar.schema.Schema) -> None:
+ """
+ Create the reader.
+ Users should not call this constructor directly. Instead, create the
+ reader via ``Client.create_reader``.
+
+ Parameters
+ ----------
+ reader: _pulsar.Reader
+ The underlying Reader object from the C extension.
+ schema: pulsar.schema.Schema
+ The schema of the data that will be received by this reader.
+ """
+ self._reader = reader
+ self._schema = schema
+
+ async def read_next(self) -> pulsar.Message:
+ """
+ Read a single message asynchronously.
+
+ Returns
+ -------
+ pulsar.Message
+ The message received.
+
+ Raises
+ ------
+ PulsarException
+ """
+ future = asyncio.get_running_loop().create_future()
+ self._reader.read_next_async(functools.partial(_set_future, future))
+ msg = await future
+ m = pulsar.Message()
+ m._message = msg
+ m._schema = self._schema
+ return m
+
+ async def has_message_available(self) -> bool:
+ """
+ Check if there is any message available to read from the current
+ position.
+ """
+ future = asyncio.get_running_loop().create_future()
+ self._reader.has_message_available_async(functools.partial(_set_future, future))
+ return await future
+
+ async def seek(self, messageid: Union[pulsar.MessageId, int]) -> None:
+ """
+ Reset this reader to a specific message id or publish timestamp
+ asynchronously.
+
+ 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._reader.seek_async(msg_id, functools.partial(_set_future, future, value=None))
+ await future
+
+ async def close(self) -> None:
+ """
+ Close the reader asynchronously.
+
+ Raises
+ ------
+ PulsarException
+ """
+ future = asyncio.get_running_loop().create_future()
+ self._reader.close_async(functools.partial(_set_future, future, value=None))
+ await future
+
+ def topic(self) -> str:
+ """
+ Return the topic this reader is reading from.
+ """
+ return self._reader.topic()
+
+ def is_connected(self) -> bool:
+ """
+ Check if the reader is connected or not.
+ """
+ return self._reader.is_connected()
+
+
 class Client:
  """
  The asynchronous version of `pulsar.Client`.
@@ -777,6 +878,93 @@ async def subscribe(self, topic: Union[str, List[str]],
  schema.attach_client(self._client)
  return Consumer(await future, schema)
 
+ # pylint: disable=too-many-arguments,too-many-locals,too-many-positional-arguments
+ async def create_reader(self, topic: str,
+ start_message_id: Union[pulsar.MessageId, _pulsar.MessageId],
+ schema: pulsar.schema.Schema | None = None,
+ receiver_queue_size: int = 1000,
+ reader_name: str | None = None,
+ subscription_role_prefix: str | None = None,
+ is_read_compacted: bool = False,
+ crypto_key_reader: pulsar.CryptoKeyReader | None = None,
+ start_message_id_inclusive: bool = False,
+ crypto_failure_action: ConsumerCryptoFailureAction =
+ ConsumerCryptoFailureAction.FAIL,
+ ) -> Reader:
+ """
+ Create a reader on a particular topic.
+
+ Parameters
+ ----------
+ topic: str
+ The name of the topic.
+ start_message_id: MessageId or _pulsar.MessageId
+ The initial reader positioning is done by specifying a message id.
+ The options are:
+
+ * ``MessageId.earliest``: Start reading from the earliest message
+ available in the topic.
+ * ``MessageId.latest``: Start reading from the end topic, only
+ getting messages published after the reader was created.
+ * ``MessageId``: When passing a particular message id, the reader
+ will position itself on that specific position.
+ schema: pulsar.schema.Schema | None, default=None
+ Define the schema of the data that will be received by this reader.
+ receiver_queue_size: int, default=1000
+ Sets the size of the reader receive queue.
+ reader_name: str | None, default=None
+ Sets the reader name.
+ subscription_role_prefix: str | None, default=None
+ Sets the subscription role prefix.
+ is_read_compacted: bool, default=False
+ Selects whether to read the compacted version of the topic.
+ crypto_key_reader: pulsar.CryptoKeyReader | None, default=None
+ Symmetric encryption class implementation.
+ start_message_id_inclusive: bool, default=False
+ Set the reader to include the startMessageId or given position of
+ any reset operation like Reader.seek.
+ crypto_failure_action: ConsumerCryptoFailureAction, \
+ default=ConsumerCryptoFailureAction.FAIL
+ Set the behavior when the decryption fails.
+
+ Returns
+ -------
+ Reader
+ The reader created
+
+ Raises
+ ------
+ PulsarException
+ """
+ if schema is None:
+ schema = pulsar.schema.BytesSchema()
+
+ if isinstance(start_message_id, pulsar.MessageId):
+ start_message_id = start_message_id._msg_id
+
+ _check_type(_pulsar.MessageId, start_message_id, 'start_message_id')
+
+ conf = _pulsar.ReaderConfiguration()
+ conf.receiver_queue_size(receiver_queue_size)
+ if reader_name is not None:
+ conf.reader_name(reader_name)
+ if subscription_role_prefix is not None:
+ conf.subscription_role_prefix(subscription_role_prefix)
+ conf.schema(schema.schema_info())
+ conf.read_compacted(is_read_compacted)
+ if crypto_key_reader is not None:
+ conf.crypto_key_reader(crypto_key_reader.cryptoKeyReader)
+ conf.start_message_id_inclusive(start_message_id_inclusive)
+ conf.crypto_failure_action(crypto_failure_action)
+
+ future = asyncio.get_running_loop().create_future()
+ self._client.create_reader_async_v2(
+ topic, start_message_id, conf, functools.partial(_set_future_v2, future)
+ )
+ reader = await future
+ schema.attach_client(self._client)
+ return Reader(reader, schema)
+
  def shutdown(self) -> None:
  """
  Shutdown the client and all the associated producers and consumers

src/client.cc

@@ -118,6 +118,18 @@ Reader Client_createReader(Client& client, const std::string& topic, const Messa
  [&](ReaderCallback callback) { client.createReaderAsync(topic, startMessageId, conf, callback); });
 }
 
+void Client_createReaderAsync(Client& client, const std::string& topic, const MessageId& startMessageId,
+ ReaderConfiguration conf, ReaderCallback callback) {
+ py::gil_scoped_release release;
+ client.createReaderAsync(topic, startMessageId, conf, callback);
+}
+
+void Client_createReaderAsyncV2(Client& client, const std::string& topic, const MessageId& startMessageId,
+ ReaderConfiguration conf, ReaderV2Callback callback) {
+ py::gil_scoped_release release;
+ client.createReaderAsyncV2(topic, startMessageId, conf, std::move(callback));
+}
+
 std::vector<std::string> Client_getTopicPartitions(Client& client, const std::string& topic) {
  return waitForAsyncValue<std::vector<std::string>>(
  [&](GetPartitionsCallback callback) { client.getPartitionsForTopicAsync(topic, callback); });
@@ -204,6 +216,8 @@ void export_client(py::module_& m) {
  .def("subscribe_topics", &Client_subscribe_topics)
  .def("subscribe_pattern", &Client_subscribe_pattern)
  .def("create_reader", &Client_createReader)
+ .def("create_reader_async", &Client_createReaderAsync)
+ .def("create_reader_async_v2", &Client_createReaderAsyncV2)
  .def("create_table_view",
  [](Client& client, const std::string& topic, const TableViewConfiguration& config) {
  return waitForAsyncValue<TableView>([&](TableViewCallback callback) {

src/reader.cc

@@ -17,6 +17,7 @@
  * under the License.
  */
 #include "utils.h"
+#include <pybind11/functional.h>
 #include <pybind11/pybind11.h>
 
 namespace py = pybind11;
@@ -54,16 +55,46 @@ void Reader_seek_timestamp(Reader& reader, uint64_t timestamp) {
 
 bool Reader_is_connected(Reader& reader) { return reader.isConnected(); }
 
+void Reader_readNextAsync(Reader& reader, ReadNextCallback callback) {
+ py::gil_scoped_release release;
+ reader.readNextAsync(callback);
+}
+
+void Reader_closeAsync(Reader& reader, ResultCallback callback) {
+ py::gil_scoped_release release;
+ reader.closeAsync(callback);
+}
+
+void Reader_seekAsync(Reader& reader, const MessageId& msgId, ResultCallback callback) {
+ py::gil_scoped_release release;
+ reader.seekAsync(msgId, callback);
+}
+
+void Reader_seekAsync_timestamp(Reader& reader, uint64_t timestamp, ResultCallback callback) {
+ py::gil_scoped_release release;
+ reader.seekAsync(timestamp, callback);
+}
+
+void Reader_hasMessageAvailableAsync(Reader& reader, HasMessageAvailableCallback callback) {
+ py::gil_scoped_release release;
+ reader.hasMessageAvailableAsync(callback);
+}
+
 void export_reader(py::module_& m) {
  using namespace py;
 
  class_<Reader>(m, "Reader")
  .def("topic", &Reader::getTopic, return_value_policy::copy)
  .def("read_next", &Reader_readNext)
  .def("read_next", &Reader_readNextTimeout)
+ .def("read_next_async", &Reader_readNextAsync)
  .def("has_message_available", &Reader_hasMessageAvailable)
+ .def("has_message_available_async", &Reader_hasMessageAvailableAsync)
  .def("close", &Reader_close)
+ .def("close_async", &Reader_closeAsync)
  .def("seek", &Reader_seek)
  .def("seek", &Reader_seek_timestamp)
+ .def("seek_async", &Reader_seekAsync)
+ .def("seek_async", &Reader_seekAsync_timestamp)
  .def("is_connected", &Reader_is_connected);
 }

tests/asyncio_test.py

@@ -39,6 +39,7 @@
  Consumer,
  Producer,
  PulsarException,
+ Reader,
  _set_future,
 )
 from pulsar.schema import ( # pylint: disable=import-error
@@ -465,6 +466,86 @@ async def test_seek_timestamp(self):
  msg = await consumer.receive()
  self.assertEqual(msg.data(), b'msg-3')
 
+ async def test_reader_simple(self):
+ topic = f'asyncio-test-reader-simple-{time.time()}'
+ reader = await self._client.create_reader(topic, pulsar.MessageId.earliest)
+ self.assertTrue(reader.is_connected())
+ self.assertEqual(reader.topic(), f'persistent://public/default/{topic}')
+
+ producer = await self._client.create_producer(topic)
+ await producer.send(b'hello')
+ msg = await reader.read_next()
+ self.assertEqual(msg.data(), b'hello')
+ with self.assertRaises(asyncio.TimeoutError):
+ await asyncio.wait_for(reader.read_next(), 1)
+ await reader.close()
+ self.assertFalse(reader.is_connected())
+
+ async def test_reader_on_last_message(self):
+ topic = f'asyncio-test-reader-on-last-message-{time.time()}'
+ producer = await self._client.create_producer(topic)
+ for i in range(10):
+ await producer.send(f'hello-{i}'.encode())
+ reader = await self._client.create_reader(topic, pulsar.MessageId.latest)
+ for i in range(10, 20):
+ await producer.send(f'hello-{i}'.encode())
+ for i in range(10, 20):
+ msg = await reader.read_next()
+ self.assertEqual(msg.data(), f'hello-{i}'.encode())
+ await reader.close()
+
+ async def test_reader_on_specific_message(self):
+ topic = f'asyncio-test-reader-on-specific-msg-{time.time()}'
+ producer = await self._client.create_producer(topic)
+ msg_ids = []
+ for i in range(10):
+ msg_id = await producer.send(f'hello-{i}'.encode())
+ msg_ids.append(msg_id)
+ reader1 = await self._client.create_reader(topic, pulsar.MessageId.earliest)
+ for i in range(5):
+ msg = await reader1.read_next()
+ self.assertEqual(msg.data(), f'hello-{i}'.encode())
+ last_msg_id = msg_ids[4]
+ reader2 = await self._client.create_reader(topic, last_msg_id)
+ for i in range(5, 10):
+ msg = await reader2.read_next()
+ self.assertEqual(msg.data(), f'hello-{i}'.encode())
+ await reader1.close()
+ await reader2.close()
+
+ async def test_reader_has_message_available(self):
+ topic = f'asyncio-test-reader-has-message-available-{time.time()}'
+ producer = await self._client.create_producer(topic)
+ reader = await self._client.create_reader(topic, pulsar.MessageId.latest)
+ self.assertFalse(await reader.has_message_available())
+ for i in range(10):
+ await producer.send(f'hello-{i}'.encode())
+ for _ in range(10):
+ self.assertTrue(await reader.has_message_available())
+ await reader.read_next()
+ self.assertFalse(await reader.has_message_available())
+ await reader.close()
+
+ async def test_reader_seek(self):
+ topic = f'asyncio-test-reader-seek-{time.time()}'
+ producer = await self._client.create_producer(topic)
+ msg_ids = []
+ for i in range(10):
+ msg_id = await producer.send(f'msg-{i}'.encode())
+ msg_ids.append(msg_id)
+ reader = await self._client.create_reader(topic, pulsar.MessageId.latest,
+ start_message_id_inclusive=False)
+ await reader.seek(msg_ids[2])
+ msg = await reader.read_next()
+ self.assertEqual(msg.data(), b'msg-3')
+ await reader.close()
+ reader_inclusive = await self._client.create_reader(topic, pulsar.MessageId.latest,
+ start_message_id_inclusive=True)
+ await reader_inclusive.seek(msg_ids[2])
+ msg = await reader_inclusive.read_next()
+ self.assertEqual(msg.data(), b'msg-2')
+ await reader_inclusive.close()
+
  async def test_schema(self):
  class ExampleRecord(Record): # pylint: disable=too-few-public-methods
  """Example record schema for testing."""
@@ -507,6 +588,12 @@ def raise_exception():
  self.assertEqual(e.exception.error(), pulsar.Result.AuthenticationError)
  # TODO: we should fix the error message not included in pattern subscription case
 
+ with self.assertRaises(PulsarException) as e:
+ await client.create_reader("private/auth/asyncio-test-token-auth-reader",
+ pulsar.MessageId.earliest)
+ self.assertEqual(e.exception.error(), pulsar.Result.AuthenticationError)
+ self.assertIn("token supplier failed", str(e.exception))
+
  await client.close()