topic metadata

Author: vgvoleg

examples/topic/basic_example.py

@@ -7,9 +7,9 @@
 
 async def connect(endpoint: str, database: str) -> ydb.aio.Driver:
     config = ydb.DriverConfig(endpoint=endpoint, database=database)
-    config.credentials = ydb.credentials_from_env_variables()
+    # config.credentials = ydb.credentials_from_env_variables()
     driver = ydb.aio.Driver(config)
-    await driver.wait(15)
+    await driver.wait(5,fail_fast=True)
     return driver
 
 
@@ -25,7 +25,11 @@ async def create_topic(driver: ydb.aio.Driver, topic: str, consumer: str):
 async def write_messages(driver: ydb.aio.Driver, topic: str):
     async with driver.topic_client.writer(topic) as writer:
         for i in range(10):
-            await writer.write(f"mess-{i}")
+            mess = ydb.TopicWriterMessage(
+                data = f"mess-{i}",
+                metadata_items= {"index": f"{i}"}
+            )
+            await writer.write(mess)
             await asyncio.sleep(1)
 
 
@@ -38,6 +42,7 @@ async def read_messages(driver: ydb.aio.Driver, topic: str, consumer: str):
                 print(mess.seqno)
                 print(mess.created_at)
                 print(mess.data.decode())
+                print(mess.metadata_items)
                 reader.commit(mess)
             except asyncio.TimeoutError:
                 return

ydb/_grpc/grpcwrapper/ydb_topic.py

@@ -208,6 +208,7 @@ class MessageData(IToProto):
             data: bytes
             uncompressed_size: int
             partitioning: "StreamWriteMessage.PartitioningType"
+            metadata_items: Dict[str, bytes]
 
             def to_proto(
                 self,
@@ -218,6 +219,11 @@ def to_proto(
                 proto.data = self.data
                 proto.uncompressed_size = self.uncompressed_size
 
+                for key, value in self.metadata_items.items():
+                    # TODO: CHECK
+                    item = ydb_topic_pb2.MetadataItem(key=key, value=value)
+                    proto.metadata_items.append(item)
+
                 if self.partitioning is None:
                     pass
                 elif isinstance(self.partitioning, StreamWriteMessage.PartitioningPartitionID):
@@ -489,16 +495,19 @@ class MessageData(IFromProto):
             data: bytes
             uncompresed_size: int
             message_group_id: str
+            metadata_items: Dict[str, bytes]
 
             @staticmethod
             def from_proto(
                 msg: ydb_topic_pb2.StreamReadMessage.ReadResponse.MessageData,
             ) -> "StreamReadMessage.ReadResponse.MessageData":
+                metadata_items = {meta.key: meta.value for meta in msg.metadata_items}
                 return StreamReadMessage.ReadResponse.MessageData(
                     offset=msg.offset,
                     seq_no=msg.seq_no,
                     created_at=msg.created_at.ToDatetime(),
                     data=msg.data,
+                    metadata_items=metadata_items,
                     uncompresed_size=msg.uncompressed_size,
                     message_group_id=msg.message_group_id,
                 )

ydb/_topic_reader/datatypes.py

@@ -40,6 +40,7 @@ class PublicMessage(ICommittable, ISessionAlive):
     written_at: datetime.datetime
     producer_id: str
     data: Union[bytes, Any]  # set as original decompressed bytes or deserialized object if deserializer set in reader
+    metadata_items: Dict[str, bytes]
     _partition_session: PartitionSession
     _commit_start_offset: int
     _commit_end_offset: int

ydb/_topic_reader/topic_reader_asyncio.py

@@ -627,6 +627,7 @@ def _read_response_to_batches(self, message: StreamReadMessage.ReadResponse) ->
                         written_at=server_batch.written_at,
                         producer_id=server_batch.producer_id,
                         data=message_data.data,
+                        metadata_items=message_data.metadata_items,
                         _partition_session=partition_session,
                         _commit_start_offset=partition_session._next_message_start_commit_offset,
                         _commit_end_offset=message_data.offset + 1,

ydb/_topic_reader/topic_reader_asyncio_test.py

@@ -74,6 +74,7 @@ def stub_message(id: int):
         written_at=datetime.datetime(2023, 3, 18, 14, 15),
         producer_id="",
         data=bytes(),
+        metadata_items={},
         _partition_session=stub_partition_session(),
         _commit_start_offset=0,
         _commit_end_offset=1,
@@ -445,6 +446,7 @@ async def test_commit_ranges_for_received_messages(
                             written_at=datetime.datetime(2023, 3, 14, 15, 42),
                             producer_id="asd",
                             data=rb"123",
+                            metadata_items={},
                             _partition_session=None,
                             _commit_start_offset=5,
                             _commit_end_offset=15,
@@ -468,6 +470,7 @@ async def test_commit_ranges_for_received_messages(
                             written_at=datetime.datetime(2023, 3, 14, 15, 42),
                             producer_id="asd",
                             data=gzip.compress(rb"123"),
+                            metadata_items={},
                             _partition_session=None,
                             _commit_start_offset=5,
                             _commit_end_offset=15,
@@ -490,6 +493,7 @@ async def test_commit_ranges_for_received_messages(
                             offset=1,
                             written_at=datetime.datetime(2023, 3, 14, 15, 42),
                             producer_id="asd",
+                            metadata_items={},
                             data=rb"123",
                             _partition_session=None,
                             _commit_start_offset=5,
@@ -504,6 +508,7 @@ async def test_commit_ranges_for_received_messages(
                             written_at=datetime.datetime(2023, 3, 14, 15, 42),
                             producer_id="asd",
                             data=rb"456",
+                            metadata_items={},
                             _partition_session=None,
                             _commit_start_offset=5,
                             _commit_end_offset=15,
@@ -527,6 +532,7 @@ async def test_commit_ranges_for_received_messages(
                             written_at=datetime.datetime(2023, 3, 14, 15, 42),
                             producer_id="asd",
                             data=gzip.compress(rb"123"),
+                            metadata_items={},
                             _partition_session=None,
                             _commit_start_offset=5,
                             _commit_end_offset=15,
@@ -540,6 +546,7 @@ async def test_commit_ranges_for_received_messages(
                             written_at=datetime.datetime(2023, 3, 14, 15, 42),
                             producer_id="asd",
                             data=gzip.compress(rb"456"),
+                            metadata_items={},
                             _partition_session=None,
                             _commit_start_offset=5,
                             _commit_end_offset=15,

ydb/_topic_writer/topic_writer.py

@@ -15,7 +15,7 @@
 from .._grpc.grpcwrapper.ydb_topic_public_types import PublicCodec
 from .. import connection
 
-Message = typing.Union["PublicMessage", "PublicMessage.SimpleMessageSourceType"]
+Message = typing.Union["PublicMessage", "PublicMessage.SimpleSourceType"]
 
 
 @dataclass
@@ -91,20 +91,23 @@ class PublicWriterInitInfo:
 class PublicMessage:
     seqno: Optional[int]
     created_at: Optional[datetime.datetime]
-    data: "PublicMessage.SimpleMessageSourceType"
+    data: "PublicMessage.SimpleSourceType"
+    metadata_items: Optional[Dict[str, "PublicMessage.SimpleSourceType"]]
 
-    SimpleMessageSourceType = Union[str, bytes]  # Will be extend
+    SimpleSourceType = Union[str, bytes]  # Will be extend
 
     def __init__(
         self,
-        data: SimpleMessageSourceType,
+        data: SimpleSourceType,
         *,
+        metadata_items: Optional[Dict[str, "PublicMessage.SimpleSourceType"]] = None,
         seqno: Optional[int] = None,
         created_at: Optional[datetime.datetime] = None,
     ):
         self.seqno = seqno
         self.created_at = created_at
         self.data = data
+        self.metadata_items = metadata_items
 
     @staticmethod
     def _create_message(data: Message) -> "PublicMessage":
@@ -121,26 +124,32 @@ def __init__(self, mess: PublicMessage):
             seq_no=mess.seqno,
             created_at=mess.created_at,
             data=mess.data,
+            metadata_items=mess.metadata_items,
             uncompressed_size=len(mess.data),
             partitioning=None,
         )
         self.codec = PublicCodec.RAW
 
-    def get_bytes(self) -> bytes:
-        if self.data is None:
+    def _get_bytes(self, obj: Optional[PublicMessage.SimpleSourceType]) -> bytes:
+        if obj is None:
             return bytes()
-        if isinstance(self.data, bytes):
-            return self.data
-        if isinstance(self.data, str):
-            return self.data.encode("utf-8")
+        if isinstance(obj, bytes):
+            return obj
+        if isinstance(obj, str):
+            return obj.encode("utf-8")
         raise ValueError("Bad data type")
 
+    def get_data_bytes(self) -> bytes:
+        return self._get_bytes(self.data)
+
     def to_message_data(self) -> StreamWriteMessage.WriteRequest.MessageData:
-        data = self.get_bytes()
+        data = self.get_data_bytes()
+        metadata_items = {key: self._get_bytes(value) for key, value in self.metadata_items.items()}
         return StreamWriteMessage.WriteRequest.MessageData(
             seq_no=self.seq_no,
             created_at=self.created_at,
             data=data,
+            metadata_items=metadata_items,
             uncompressed_size=len(data),
             partitioning=None,  # unsupported by server now
         )
@@ -221,6 +230,7 @@ def messages_to_proto_requests(
                     seq_no=_max_int,
                     created_at=datetime.datetime(3000, 1, 1, 1, 1, 1, 1),
                     data=bytes(1),
+                    metadata_items={},
                     uncompressed_size=_max_int,
                     partitioning=StreamWriteMessage.PartitioningMessageGroupID(
                         message_group_id="a" * 100,

ydb/_topic_writer/topic_writer_asyncio.py

@@ -427,7 +427,7 @@ async def _encode_data_inplace(self, codec: PublicCodec, messages: List[Internal
 
         for message in messages:
             encoded_data_futures = eventloop.run_in_executor(
-                self._encode_executor, encoder_function, message.get_bytes()
+                self._encode_executor, encoder_function, message.get_data_bytes()
             )
             encode_waiters.append(encoded_data_futures)
 
@@ -493,7 +493,7 @@ def get_compressed_size(codec) -> int:
             f = self._codec_functions[codec]
 
             for m in test_messages:
-                encoded = f(m.get_bytes())
+                encoded = f(m.get_data_bytes())
                 s += len(encoded)
 
             return s

ydb/_topic_writer/topic_writer_asyncio_test.py

@@ -544,7 +544,7 @@ def add_messages(_self, messages: typing.List[InternalMessage]):
             mess = mess[0]
 
             assert mess.codec == expected_codecs[i]
-            assert mess.get_bytes() == expected_datas[i]
+            assert mess.get_data_bytes() == expected_datas[i]
 
         await reconnector.close(flush=False)
 
@@ -575,7 +575,7 @@ async def test_encode_data_inplace(
 
         for index, mess in enumerate(messages):
             assert mess.codec == codec
-            assert mess.get_bytes() == expected_datas[index]
+            assert mess.get_data_bytes() == expected_datas[index]
 
         await reconnector.close(flush=True)