add custom encoders

Author: rekby

ydb/_topic_writer/topic_writer.py

@@ -34,8 +34,10 @@ class PublicWriterSettings:
     encoder_executor: Optional[
         concurrent.futures.Executor
     ] = None  # default shared client executor pool
+    encoders: Optional[
+        typing.Mapping[PublicCodec, typing.Callable[[bytes], bytes]]
+    ] = None
     # get_last_seqno: bool = False
-    # encoders: Union[Mapping[int, Callable[[bytes], bytes]], None] = None
     # serializer: Union[Callable[[Any], bytes], None] = None
     # send_buffer_count: Optional[int] = 10000
     # send_buffer_bytes: Optional[int] = 100 * 1024 * 1024

ydb/_topic_writer/topic_writer_asyncio.py

@@ -198,6 +198,11 @@ def __init__(self, driver: SupportedDriverType, settings: WriterSettings):
             PublicCodec.RAW: lambda data: data,
             PublicCodec.GZIP: gzip.compress,
         }
+
+        if settings.encoders:
+            for codec, encoder in settings.encoders.items():
+                self._codec_functions[codec] = encoder
+
         self._encode_executor = settings.encoder_executor
 
         self._codec_selector_batch_num = 0

ydb/_topic_writer/topic_writer_asyncio_test.py

@@ -27,7 +27,7 @@
     TopicWriterError,
 )
 from .._grpc.grpcwrapper.ydb_topic_public_types import PublicCodec
-from .._topic_common.test_helpers import StreamMock
+from .._topic_common.test_helpers import StreamMock, wait_for_fast
 
 from .topic_writer_asyncio import (
     WriterAsyncIOStream,
@@ -572,6 +572,35 @@ async def test_encode_data_inplace(
             assert mess.codec == codec
             assert mess.get_bytes() == expected_datas[index]
 
+    async def test_custom_encoder(
+        self, default_driver, default_settings, get_stream_writer
+    ):
+        codec = 10001
+
+        settings = copy.copy(default_settings)
+        settings.encoders = {codec: lambda x: bytes(reversed(x))}
+        settings.codec = codec
+        reconnector = WriterAsyncIOReconnector(default_driver, settings)
+
+        now = datetime.datetime.now()
+        seqno = self.init_last_seqno + 1
+
+        await reconnector.write_with_ack_future(
+            [PublicMessage(data=b"123", seqno=seqno, created_at=now)]
+        )
+
+        stream_writer = get_stream_writer()
+        sent_messages = await wait_for_fast(stream_writer.from_client.get())
+
+        expected_mess = InternalMessage(
+            PublicMessage(data=b"321", seqno=seqno, created_at=now)
+        )
+        expected_mess.codec = codec
+
+        assert sent_messages == [expected_mess]
+
+        await reconnector.close(flush=False)
+
 
 @pytest.mark.asyncio
 class TestWriterAsyncIO:

ydb/topic.py

@@ -3,7 +3,7 @@
 import concurrent.futures
 import datetime
 from dataclasses import dataclass
-from typing import List, Union, Mapping, Optional, Dict
+from typing import List, Union, Mapping, Optional, Dict, Callable
 
 from . import aio, Credentials, _apis, issues
 
@@ -167,6 +167,9 @@ def writer(
         auto_seqno: bool = True,
         auto_created_at: bool = True,
         codec: Optional[TopicCodec] = None,  # default mean auto-select
+        encoders: Optional[
+            Mapping[_ydb_topic_public_types.PublicCodec, Callable[[bytes], bytes]]
+        ] = None,
         encoder_executor: Optional[
             concurrent.futures.Executor
         ] = None,  # default shared client executor pool
@@ -325,6 +328,9 @@ def writer(
         auto_seqno: bool = True,
         auto_created_at: bool = True,
         codec: Optional[TopicCodec] = None,  # default mean auto-select
+        encoders: Optional[
+            Mapping[_ydb_topic_public_types.PublicCodec, Callable[[bytes], bytes]]
+        ] = None,
         encoder_executor: Optional[
             concurrent.futures.Executor
         ] = None,  # default shared client executor pool