Merge pull request #124 First iteration of topic reader

Author: rekby

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Initial implementation of topic reader
+
 ## 3.0.1b3 ##
 * Fix error of check retriable error for idempotent operations (error exist since 2.12.1)
 

tests/conftest.py

@@ -100,11 +100,17 @@ async def driver(endpoint, database, event_loop):
 
 
 @pytest.fixture()
-def topic_path(endpoint) -> str:
+def topic_consumer():
+    return "fixture-consumer"
+
+
+@pytest.fixture()
+def topic_path(endpoint, topic_consumer) -> str:
     subprocess.run(
         """docker-compose exec -T ydb /ydb -e grpc://%s -d /local topic drop /local/test-topic"""
         % endpoint,
         shell=True,
+        capture_output=True,
     )
     res = subprocess.run(
         """docker-compose exec -T ydb /ydb -e grpc://%s -d /local topic create /local/test-topic"""
@@ -114,4 +120,26 @@ def topic_path(endpoint) -> str:
     )
     assert res.returncode == 0, res.stderr + res.stdout
 
+    res = subprocess.run(
+        """docker-compose exec -T ydb /ydb -e grpc://%s -d /local topic consumer add --consumer %s /local/test-topic"""
+        % (endpoint, topic_consumer),
+        shell=True,
+        capture_output=True,
+    )
+    assert res.returncode == 0, res.stderr + res.stdout
+
     return "/local/test-topic"
+
+
+@pytest.fixture()
+@pytest.mark.asyncio()
+async def topic_with_messages(driver, topic_path):
+    pass
+    writer = driver.topic_client.topic_writer(
+        topic_path, producer_and_message_group_id="fixture-producer-id"
+    )
+    await writer.write_with_ack(
+        ydb.TopicWriterMessage(data="123".encode()),
+        ydb.TopicWriterMessage(data="456".encode()),
+    )
+    await writer.close()

tests/topics/test_topic_reader.py

@@ -0,0 +1,11 @@
+import pytest
+
+
+@pytest.mark.asyncio
+class TestTopicWriterAsyncIO:
+    async def test_read_message(
+        self, driver, topic_path, topic_with_messages, topic_consumer
+    ):
+        reader = driver.topic_client.topic_reader(topic_consumer, topic_path)
+
+        assert await reader.receive_batch() is not None

ydb/_topic_reader/datatypes.py

@@ -0,0 +1,93 @@
+import abc
+import enum
+from dataclasses import dataclass
+import datetime
+from typing import Mapping, Union, Any, List, Dict
+
+
+class ICommittable(abc.ABC):
+    @property
+    @abc.abstractmethod
+    def start_offset(self) -> int:
+        pass
+
+    @property
+    @abc.abstractmethod
+    def end_offset(self) -> int:
+        pass
+
+
+class ISessionAlive(abc.ABC):
+    @property
+    @abc.abstractmethod
+    def is_alive(self) -> bool:
+        pass
+
+
+@dataclass
+class PublicMessage(ICommittable, ISessionAlive):
+    seqno: int
+    created_at: datetime.datetime
+    message_group_id: str
+    session_metadata: Dict[str, str]
+    offset: int
+    written_at: datetime.datetime
+    producer_id: str
+    data: Union[
+        bytes, Any
+    ]  # set as original decompressed bytes or deserialized object if deserializer set in reader
+    _partition_session: "PartitionSession"
+
+    @property
+    def start_offset(self) -> int:
+        raise NotImplementedError()
+
+    @property
+    def end_offset(self) -> int:
+        raise NotImplementedError()
+
+    # ISessionAlive implementation
+    @property
+    def is_alive(self) -> bool:
+        raise NotImplementedError()
+
+
+@dataclass
+class PartitionSession:
+    id: int
+    state: "PartitionSession.State"
+    topic_path: str
+    partition_id: int
+
+    def stop(self):
+        self.state = PartitionSession.State.Stopped
+
+    class State(enum.Enum):
+        Active = 1
+        GracefulShutdown = 2
+        Stopped = 3
+
+
+@dataclass
+class PublicBatch(ICommittable, ISessionAlive):
+    session_metadata: Mapping[str, str]
+    messages: List[PublicMessage]
+    _partition_session: PartitionSession
+    _bytes_size: int
+
+    @property
+    def start_offset(self) -> int:
+        raise NotImplementedError()
+
+    @property
+    def end_offset(self) -> int:
+        raise NotImplementedError()
+
+    # ISessionAlive implementation
+    @property
+    def is_alive(self) -> bool:
+        state = self._partition_session.state
+        return (
+            state == PartitionSession.State.Active
+            or state == PartitionSession.State.GracefulShutdown
+        )