add support reader codecs (#214)

* add support reader codecs
* fix writer codec bug

Author: rekby

tests/topics/test_topic_reader.py

@@ -1,5 +1,7 @@
 import pytest
 
+import ydb
+
 
 @pytest.mark.asyncio
 class TestTopicReaderAsyncIO:
@@ -25,6 +27,36 @@ async def test_read_and_commit_message(
 
         await reader.close()
 
+    async def test_read_compressed_messages(self, driver, topic_path, topic_consumer):
+        async with driver.topic_client.writer(
+            topic_path, codec=ydb.TopicCodec.GZIP
+        ) as writer:
+            await writer.write("123")
+
+        async with driver.topic_client.reader(topic_consumer, topic_path) as reader:
+            batch = await reader.receive_batch()
+            assert batch.messages[0].data.decode() == "123"
+
+    async def test_read_custom_encoded(self, driver, topic_path, topic_consumer):
+        codec = 10001
+
+        def encode(b: bytes):
+            return bytes(reversed(b))
+
+        def decode(b: bytes):
+            return bytes(reversed(b))
+
+        async with driver.topic_client.writer(
+            topic_path, codec=codec, encoders={codec: encode}
+        ) as writer:
+            await writer.write("123")
+
+        async with driver.topic_client.reader(
+            topic_consumer, topic_path, decoders={codec: decode}
+        ) as reader:
+            batch = await reader.receive_batch()
+            assert batch.messages[0].data.decode() == "123"
+
 
 class TestTopicReaderSync:
     def test_read_message(
@@ -45,3 +77,33 @@ def test_read_and_commit_message(
         reader = driver_sync.topic_client.reader(topic_consumer, topic_path)
         batch2 = reader.receive_batch()
         assert batch.messages[0] != batch2.messages[0]
+
+    def test_read_compressed_messages(self, driver_sync, topic_path, topic_consumer):
+        with driver_sync.topic_client.writer(
+            topic_path, codec=ydb.TopicCodec.GZIP
+        ) as writer:
+            writer.write("123")
+
+        with driver_sync.topic_client.reader(topic_consumer, topic_path) as reader:
+            batch = reader.receive_batch()
+            assert batch.messages[0].data.decode() == "123"
+
+    def test_read_custom_encoded(self, driver_sync, topic_path, topic_consumer):
+        codec = 10001
+
+        def encode(b: bytes):
+            return bytes(reversed(b))
+
+        def decode(b: bytes):
+            return bytes(reversed(b))
+
+        with driver_sync.topic_client.writer(
+            topic_path, codec=codec, encoders={codec: encode}
+        ) as writer:
+            writer.write("123")
+
+        with driver_sync.topic_client.reader(
+            topic_consumer, topic_path, decoders={codec: decode}
+        ) as reader:
+            batch = reader.receive_batch()
+            assert batch.messages[0].data.decode() == "123"

ydb/_topic_reader/datatypes.py

@@ -9,7 +9,7 @@
 import datetime
 from typing import Mapping, Union, Any, List, Dict, Deque, Optional
 
-from ydb._grpc.grpcwrapper.ydb_topic import OffsetsRange
+from ydb._grpc.grpcwrapper.ydb_topic import OffsetsRange, Codec
 from ydb._topic_reader import topic_reader_asyncio
 
 
@@ -168,6 +168,7 @@ class PublicBatch(ICommittable, ISessionAlive):
     messages: List[PublicMessage]
     _partition_session: PartitionSession
     _bytes_size: int
+    _codec: Codec
 
     def _commit_get_partition_session(self) -> PartitionSession:
         return self.messages[0]._commit_get_partition_session()

ydb/_topic_reader/topic_reader.py

@@ -1,10 +1,13 @@
+import concurrent.futures
 import enum
 import datetime
 from dataclasses import dataclass
 from typing import (
     Union,
     Optional,
     List,
+    Mapping,
+    Callable,
 )
 
 from ..table import RetrySettings
@@ -27,6 +30,13 @@ class PublicReaderSettings:
     consumer: str
     topic: str
     buffer_size_bytes: int = 50 * 1024 * 1024
+
+    decoders: Union[Mapping[int, Callable[[bytes], bytes]], None] = None
+    """decoders: map[codec_code] func(encoded_bytes)->decoded_bytes"""
+
+    # decoder_executor, must be set for handle non raw messages
+    decoder_executor: Optional[concurrent.futures.Executor] = None
+
     # on_commit: Callable[["Events.OnCommit"], None] = None
     # on_get_partition_start_offset: Callable[
     #     ["Events.OnPartitionGetStartOffsetRequest"],
@@ -35,7 +45,6 @@ class PublicReaderSettings:
     # on_partition_session_start: Callable[["StubEvent"], None] = None
     # on_partition_session_stop: Callable[["StubEvent"], None] = None
     # on_partition_session_close: Callable[["StubEvent"], None] = None  # todo?
-    # decoder: Union[Mapping[int, Callable[[bytes], bytes]], None] = None
     # deserializer: Union[Callable[[bytes], Any], None] = None
     # one_attempt_connection_timeout: Union[float, None] = 1
     # connection_timeout: Union[float, None] = None

ydb/_topic_reader/topic_reader_asyncio.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import asyncio
+import concurrent.futures
+import gzip
 import typing
 from asyncio import Task
 from collections import deque
@@ -17,14 +19,18 @@
     SupportedDriverType,
     GrpcWrapperAsyncIO,
 )
-from .._grpc.grpcwrapper.ydb_topic import StreamReadMessage
+from .._grpc.grpcwrapper.ydb_topic import StreamReadMessage, Codec
 from .._errors import check_retriable_error
 
 
 class TopicReaderError(YdbError):
     pass
 
 
+class TopicReaderUnexpectedCodec(YdbError):
+    pass
+
+
 class TopicReaderCommitToExpiredPartition(TopicReaderError):
     """
     Commit message when partition read session are dropped.
@@ -57,10 +63,10 @@ def __init__(self, driver: Driver, settings: topic_reader.PublicReaderSettings):
         self._reconnector = ReaderReconnector(driver, settings)
 
     async def __aenter__(self):
-        raise NotImplementedError()
+        return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        raise NotImplementedError()
+        await self.close()
 
     def __del__(self):
         if not self._closed:
@@ -259,6 +265,7 @@ def _set_first_error(self, err: issues.Error):
 class ReaderStream:
     _static_id_counter = AtomicCounter()
 
+    _loop: asyncio.AbstractEventLoop
     _id: int
     _reader_reconnector_id: int
     _session_id: str
@@ -267,6 +274,15 @@ class ReaderStream:
     _background_tasks: Set[asyncio.Task]
     _partition_sessions: Dict[int, datatypes.PartitionSession]
     _buffer_size_bytes: int  # use for init request, then for debug purposes only
+    _decode_executor: concurrent.futures.Executor
+    _decoders: Dict[
+        int, typing.Callable[[bytes], bytes]
+    ]  # dict[codec_code] func(encoded_bytes)->decoded_bytes
+
+    if typing.TYPE_CHECKING:
+        _batches_to_decode: asyncio.Queue[datatypes.PublicBatch]
+    else:
+        _batches_to_decode: asyncio.Queue
 
     _state_changed: asyncio.Event
     _closed: bool
@@ -276,6 +292,7 @@ class ReaderStream:
     def __init__(
         self, reader_reconnector_id: int, settings: topic_reader.PublicReaderSettings
     ):
+        self._loop = asyncio.get_running_loop()
         self._id = ReaderStream._static_id_counter.inc_and_get()
         self._reader_reconnector_id = reader_reconnector_id
         self._session_id = "not initialized"
@@ -284,10 +301,16 @@ def __init__(
         self._background_tasks = set()
         self._partition_sessions = dict()
         self._buffer_size_bytes = settings.buffer_size_bytes
+        self._decode_executor = settings.decoder_executor
+
+        self._decoders = {Codec.CODEC_GZIP: gzip.decompress}
+        if settings.decoders:
+            self._decoders.update(settings.decoders)
 
         self._state_changed = asyncio.Event()
         self._closed = False
         self._first_error = asyncio.get_running_loop().create_future()
+        self._batches_to_decode = asyncio.Queue()
         self._message_batches = deque()
 
     @staticmethod
@@ -324,8 +347,10 @@ async def _start(
                 "Unexpected message after InitRequest: %s", init_response
             )
 
-        read_messages_task = asyncio.create_task(self._read_messages_loop(stream))
-        self._background_tasks.add(read_messages_task)
+        self._background_tasks.add(
+            asyncio.create_task(self._read_messages_loop(stream))
+        )
+        self._background_tasks.add(asyncio.create_task(self._decode_batches_loop()))
 
     async def wait_error(self):
         raise await self._first_error
@@ -486,10 +511,12 @@ def _on_partition_session_stop(
             )
 
     def _on_read_response(self, message: StreamReadMessage.ReadResponse):
-        batches = self._read_response_to_batches(message)
-        self._message_batches.extend(batches)
         self._buffer_consume_bytes(message.bytes_size)
 
+        batches = self._read_response_to_batches(message)
+        for batch in batches:
+            self._batches_to_decode.put_nowait(batch)
+
     def _on_commit_response(self, message: StreamReadMessage.CommitOffsetResponse):
         for partition_offset in message.partitions_committed_offsets:
             session = self._partition_sessions.get(
@@ -561,12 +588,44 @@ def _read_response_to_batches(
                         messages=messages,
                         _partition_session=partition_session,
                         _bytes_size=bytes_per_batch,
+                        _codec=Codec(server_batch.codec),
                     )
                     batches.append(batch)
 
         batches[-1]._bytes_size += additional_bytes_to_last_batch
         return batches
 
+    async def _decode_batches_loop(self):
+        while True:
+            batch = await self._batches_to_decode.get()
+            await self._decode_batch_inplace(batch)
+            self._message_batches.append(batch)
+            self._state_changed.set()
+
+    async def _decode_batch_inplace(self, batch):
+        if batch._codec == Codec.CODEC_RAW:
+            return
+
+        try:
+            decode_func = self._decoders[batch._codec]
+        except KeyError:
+            raise TopicReaderUnexpectedCodec(
+                "Receive message with unexpected codec: %s" % batch._codec
+            )
+
+        decode_data_futures = []
+        for message in batch.messages:
+            future = self._loop.run_in_executor(
+                self._decode_executor, decode_func, message.data
+            )
+            decode_data_futures.append(future)
+
+        decoded_data = await asyncio.gather(*decode_data_futures)
+        for index, message in enumerate(batch.messages):
+            message.data = decoded_data[index]
+
+        batch._codec = Codec.CODEC_RAW
+
     def _set_first_error(self, err: YdbError):
         try:
             self._first_error.set_result(err)