Merge pull request #198 Codec support for writer

ydb-platform · Mar 7, 2023 · b31efcb · b31efcb
2 parents 95ab809 + ccc6eec
commit b31efcb
Show file tree

Hide file tree

Showing 8 changed files with 497 additions and 49 deletions.
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -163,7 +163,9 @@ async def topic_path(driver, topic_consumer, database) -> str:
 @pytest.fixture()
 @pytest.mark.asyncio()
 async def topic_with_messages(driver, topic_path):
-    writer = driver.topic_client.writer(topic_path, producer_id="fixture-producer-id")
+    writer = driver.topic_client.writer(
+        topic_path, producer_id="fixture-producer-id", codec=ydb.TopicCodec.RAW
+    )
     await writer.write_with_ack(
         [
             ydb.TopicWriterMessage(data="123".encode()),

diff --git a/tests/topics/test_topic_writer.py b/tests/topics/test_topic_writer.py
@@ -84,6 +84,20 @@ async def test_write_multi_message_with_ack(
         assert batch.messages[0].seqno == 2
         assert batch.messages[0].data == "456".encode()
 
+    @pytest.mark.parametrize(
+        "codec",
+        [
+            ydb.TopicCodec.RAW,
+            ydb.TopicCodec.GZIP,
+            None,
+        ],
+    )
+    async def test_write_encoded(self, driver: ydb.Driver, topic_path: str, codec):
+        async with driver.topic_client.writer(topic_path, codec=codec) as writer:
+            writer.write("a" * 1000)
+            writer.write("b" * 1000)
+            writer.write("c" * 1000)
+
 
 class TestTopicWriterSync:
     def test_send_message(self, driver_sync: ydb.Driver, topic_path):
@@ -163,3 +177,17 @@ def test_write_multi_message_with_ack(
         assert batch.messages[0].offset == 1
         assert batch.messages[0].seqno == 2
         assert batch.messages[0].data == "456".encode()
+
+    @pytest.mark.parametrize(
+        "codec",
+        [
+            ydb.TopicCodec.RAW,
+            ydb.TopicCodec.GZIP,
+            None,
+        ],
+    )
+    def test_write_encoded(self, driver_sync: ydb.Driver, topic_path: str, codec):
+        with driver_sync.topic_client.writer(topic_path, codec=codec) as writer:
+            writer.write("a" * 1000)
+            writer.write("b" * 1000)
+            writer.write("c" * 1000)
diff --git a/ydb/_grpc/grpcwrapper/ydb_topic_public_types.py b/ydb/_grpc/grpcwrapper/ydb_topic_public_types.py
@@ -31,11 +31,18 @@ class CreateTopicRequestParams:
 
 
 class PublicCodec(int):
+    """
+    Codec value may contain any int number.
+
+    Values below is only well-known predefined values,
+    but protocol support custom codecs.
+    """
+
     UNSPECIFIED = 0
     RAW = 1
     GZIP = 2
-    LZOP = 3
-    ZSTD = 4
+    LZOP = 3  # Has not supported codec in standard library
+    ZSTD = 4  # Has not supported codec in standard library
 
 
 class PublicMeteringMode(IntEnum):

diff --git a/ydb/_topic_writer/topic_writer.py b/ydb/_topic_writer/topic_writer.py
@@ -1,3 +1,4 @@
+import concurrent.futures
 import datetime
 import enum
 import uuid
@@ -10,9 +11,9 @@
 import ydb.aio
 from .._grpc.grpcwrapper.ydb_topic import Codec, StreamWriteMessage
 from .._grpc.grpcwrapper.common_utils import IToProto
+from .._grpc.grpcwrapper.ydb_topic_public_types import PublicCodec
 
-
-MessageType = typing.Union["PublicMessage", "PublicMessage.SimpleMessageSourceType"]
+Message = typing.Union["PublicMessage", "PublicMessage.SimpleMessageSourceType"]
 
 
 @dataclass
@@ -29,8 +30,14 @@ class PublicWriterSettings:
     partition_id: Optional[int] = None
     auto_seqno: bool = True
     auto_created_at: bool = True
+    codec: Optional[PublicCodec] = None  # default mean auto-select
+    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
@@ -85,8 +92,9 @@ class SendMode(Enum):
 
 @dataclass
 class PublicWriterInitInfo:
-    __slots__ = "last_seqno"
+    __slots__ = ("last_seqno", "supported_codecs")
     last_seqno: Optional[int]
+    supported_codecs: List[PublicCodec]
 
 
 class PublicMessage:
@@ -108,24 +116,24 @@ def __init__(
         self.data = data
 
     @staticmethod
-    def _create_message(
-        data: Union["PublicMessage", "PublicMessage.SimpleMessageSourceType"]
-    ) -> "PublicMessage":
+    def _create_message(data: Message) -> "PublicMessage":
         if isinstance(data, PublicMessage):
             return data
         return PublicMessage(data=data)
 
 
 class InternalMessage(StreamWriteMessage.WriteRequest.MessageData, IToProto):
+    codec: PublicCodec
+
     def __init__(self, mess: PublicMessage):
-        StreamWriteMessage.WriteRequest.MessageData.__init__(
-            self,
+        super().__init__(
             seq_no=mess.seqno,
             created_at=mess.created_at,
             data=mess.data,
             uncompressed_size=len(mess.data),
             partitioning=None,
         )
+        self.codec = PublicCodec.RAW
 
     def get_bytes(self) -> bytes:
         if self.data is None: