remove difficult variants of send message

fix examples
add wait container timeout - for tests on m1

Author: rekby

examples/topic/writer_async_example.py

@@ -37,15 +37,15 @@ async def connect_without_context_manager(db: ydb.aio.Driver):
 async def send_messages(writer: ydb.TopicWriterAsyncIO):
     # simple str/bytes without additional metadata
     await writer.write("mess")  # send text
-    await writer.write(bytes([1, 2, 3]))  # send bytes
-    await writer.write("mess-1", "mess-2")  # send two messages
+    await writer.write(bytes([1, 2, 3]))  # send single message with bytes 1,2,3
+    await writer.write(["mess-1", "mess-2"])  # send two messages
 
     # full forms
     await writer.write(ydb.TopicWriterMessage("mess"))  # send text
     await writer.write(ydb.TopicWriterMessage(bytes([1, 2, 3])))  # send bytes
-    await writer.write(
+    await writer.write([
         ydb.TopicWriterMessage("mess-1"), ydb.TopicWriterMessage("mess-2")
-    )  # send few messages by one call
+    ])  # send few messages by one call
 
     # with meta
     await writer.write(
@@ -71,12 +71,12 @@ async def send_messages_with_manual_seqno(writer: ydb.TopicWriter):
 
 async def send_messages_with_wait_ack(writer: ydb.TopicWriterAsyncIO):
     # future wait
-    await writer.write_with_result(
+    await writer.write_with_result([
         ydb.TopicWriterMessage("mess", seqno=1), ydb.TopicWriterMessage("mess", seqno=2)
-    )
+    ])
 
     # send with flush
-    await writer.write("1", "2", "3")
+    await writer.write(["1", "2", "3"])
     await writer.flush()
 
 

examples/topic/writer_example.py

@@ -44,21 +44,21 @@ def connect_without_context_manager(db: ydb.Driver):
     try:
         pass  # some code
     finally:
-        await writer.close()
+        writer.close()
 
 
 def send_messages(writer: ydb.TopicWriter):
     # simple str/bytes without additional metadata
     writer.write("mess")  # send text
-    writer.write(bytes([1, 2, 3]))  # send bytes
-    writer.write("mess-1", "mess-2")  # send two messages
+    writer.write(bytes([1, 2, 3]))  # send single message with bytes 1,2,3
+    writer.write(["mess-1", "mess-2"])  # send two messages
 
     # full forms
     writer.write(ydb.TopicWriterMessage("mess"))  # send text
     writer.write(ydb.TopicWriterMessage(bytes([1, 2, 3])))  # send bytes
-    writer.write(
+    writer.write([
         ydb.TopicWriterMessage("mess-1"), ydb.TopicWriterMessage("mess-2")
-    )  # send few messages by one call
+    ])  # send few messages by one call
 
     # with meta
     writer.write(ydb.TopicWriterMessage("asd", seqno=123, created_at_ns=time.time_ns()))
@@ -87,13 +87,13 @@ def send_messages_with_wait_ack(writer: ydb.TopicWriter):
     ).result()
 
     # implicit, by sync call
-    writer.write_with_ack(
+    writer.write_with_ack([
         ydb.TopicWriterMessage("mess", seqno=1), ydb.TopicWriterMessage("mess", seqno=2)
-    )
+    ])
     # write_with_ack
 
     # send with flush
-    writer.write("1", "2", "3")
+    writer.write(["1", "2", "3"])
     writer.flush()
 
 

tests/conftest.py

@@ -12,7 +12,7 @@ def docker_compose_file(pytestconfig):
 
 
 def wait_container_ready(driver):
-    driver.wait(timeout=10)
+    driver.wait(timeout=30)
 
     with ydb.SessionPool(driver) as pool:
 
@@ -133,12 +133,16 @@ async def topic_path(driver, topic_consumer, database) -> str:
 async def topic_with_messages(driver, topic_path):
     writer = driver.topic_client.writer(topic_path, producer_id="fixture-producer-id")
     await writer.write_with_ack(
-        ydb.TopicWriterMessage(data="123".encode()),
-        ydb.TopicWriterMessage(data="456".encode()),
+        [
+            ydb.TopicWriterMessage(data="123".encode()),
+            ydb.TopicWriterMessage(data="456".encode()),
+        ]
     )
     await writer.write_with_ack(
-        ydb.TopicWriterMessage(data="789".encode()),
-        ydb.TopicWriterMessage(data="0".encode()),
+        [
+            ydb.TopicWriterMessage(data="789".encode()),
+            ydb.TopicWriterMessage(data="0".encode()),
+        ]
     )
     await writer.close()
 

tests/topics/test_topic_writer.py

@@ -60,6 +60,30 @@ async def test_auto_flush_on_close(self, driver: ydb.aio.Driver, topic_path):
             init_info = await writer.wait_init()
             assert init_info.last_seqno == last_seqno
 
+    async def test_write_multi_message_with_ack(
+        self, driver: ydb.aio.Driver, topic_path, topic_reader: ydb.TopicReaderAsyncIO
+    ):
+        async with driver.topic_client.writer(topic_path) as writer:
+            await writer.write_with_ack(
+                [
+                    ydb.TopicWriterMessage(data="123".encode()),
+                    ydb.TopicWriterMessage(data="456".encode()),
+                ]
+            )
+
+        batch = await topic_reader.receive_batch()
+
+        assert batch.messages[0].offset == 0
+        assert batch.messages[0].seqno == 1
+        assert batch.messages[0].data == "123".encode()
+
+        # remove second recieve batch when implement batching
+        # https://github.com/ydb-platform/ydb-python-sdk/issues/142
+        batch = await topic_reader.receive_batch()
+        assert batch.messages[0].offset == 1
+        assert batch.messages[0].seqno == 2
+        assert batch.messages[0].data == "456".encode()
+
 
 class TestTopicWriterSync:
     def test_send_message(self, driver_sync: ydb.Driver, topic_path):
@@ -115,3 +139,27 @@ def test_random_producer_id(
         batch2 = topic_reader_sync.receive_batch()
 
         assert batch1.messages[0].producer_id != batch2.messages[0].producer_id
+
+    def test_write_multi_message_with_ack(
+        self, driver_sync: ydb.Driver, topic_path, topic_reader_sync: ydb.TopicReader
+    ):
+        with driver_sync.topic_client.writer(topic_path) as writer:
+            writer.write_with_ack(
+                [
+                    ydb.TopicWriterMessage(data="123".encode()),
+                    ydb.TopicWriterMessage(data="456".encode()),
+                ]
+            )
+
+        batch = topic_reader_sync.receive_batch()
+
+        assert batch.messages[0].offset == 0
+        assert batch.messages[0].seqno == 1
+        assert batch.messages[0].data == "123".encode()
+
+        # remove second recieve batch when implement batching
+        # https://github.com/ydb-platform/ydb-python-sdk/issues/142
+        batch = topic_reader_sync.receive_batch()
+        assert batch.messages[0].offset == 1
+        assert batch.messages[0].seqno == 2
+        assert batch.messages[0].data == "456".encode()

ydb/_topic_reader/topic_reader_asyncio.py

@@ -115,7 +115,7 @@ async def receive_batch(
         *,
         max_messages: typing.Union[int, None] = None,
         max_bytes: typing.Union[int, None] = None,
-    ) -> typing.Union[topic_reader.PublicBatch, None]:
+    ) -> typing.Union[datatypes.PublicBatch, None]:
         """
         Get one messages batch from reader.
         All messages in a batch from same partition.
@@ -243,7 +243,8 @@ def commit(
         return self._stream_reader.commit(batch)
 
     async def close(self):
-        await self._stream_reader.close()
+        if self._stream_reader:
+            await self._stream_reader.close()
         for task in self._background_tasks:
             task.cancel()
 

ydb/_topic_writer/topic_writer.py

@@ -3,7 +3,7 @@
 import uuid
 from dataclasses import dataclass
 from enum import Enum
-from typing import List, Union, TextIO, BinaryIO, Optional, Any, Dict
+from typing import List, Union, Optional, Any, Dict
 
 import typing
 
@@ -92,9 +92,9 @@ class PublicWriterInitInfo:
 class PublicMessage:
     seqno: Optional[int]
     created_at: Optional[datetime.datetime]
-    data: Union[str, bytes, TextIO, BinaryIO]
+    data: "PublicMessage.SimpleMessageSourceType"
 
-    SimpleMessageSourceType = Union[str, bytes, TextIO, BinaryIO]
+    SimpleMessageSourceType = Union[str, bytes]  # Will be extend
 
     def __init__(
         self,
@@ -107,6 +107,14 @@ def __init__(
         self.created_at = created_at
         self.data = data
 
+    @staticmethod
+    def _create_message(
+        data: Union["PublicMessage", "PublicMessage.SimpleMessageSourceType"]
+    ) -> "PublicMessage":
+        if isinstance(data, PublicMessage):
+            return data
+        return PublicMessage(data=data)
+
 
 class InternalMessage(StreamWriteMessage.WriteRequest.MessageData, IToProto):
     def __init__(self, mess: PublicMessage):

ydb/_topic_writer/topic_writer_asyncio.py

@@ -1,7 +1,7 @@
 import asyncio
 import datetime
 from collections import deque
-from typing import Deque, AsyncIterator, Union, List, Optional
+from typing import Deque, AsyncIterator, Union, List
 
 import ydb
 from .topic_writer import (
@@ -76,7 +76,6 @@ async def close(self, *, flush: bool = True):
     async def write_with_ack(
         self,
         messages: Union[MessageType, List[MessageType]],
-        *args: Optional[MessageType],
     ) -> Union[PublicWriteResultTypes, List[PublicWriteResultTypes]]:
         """
         IT IS SLOWLY WAY. IT IS BAD CHOISE IN MOST CASES.
@@ -86,7 +85,7 @@ async def write_with_ack(
 
         For wait with timeout use asyncio.wait_for.
         """
-        futures = await self.write_with_ack_future(messages, *args)
+        futures = await self.write_with_ack_future(messages)
         if not isinstance(futures, list):
             futures = [futures]
 
@@ -98,7 +97,6 @@ async def write_with_ack(
     async def write_with_ack_future(
         self,
         messages: Union[MessageType, List[MessageType]],
-        *args: Optional[MessageType],
     ) -> Union[asyncio.Future, List[asyncio.Future]]:
         """
         send one or number of messages to server.
@@ -108,20 +106,22 @@ async def write_with_ack_future(
 
         For wait with timeout use asyncio.wait_for.
         """
+        input_single_message = not isinstance(messages, list)
         if isinstance(messages, PublicMessage):
-            futures = await self._reconnector.write_with_ack_future([messages])
-            return futures[0]
+            messages = [PublicMessage._create_message(messages)]
         if isinstance(messages, list):
-            for m in messages:
-                if not isinstance(m, PublicMessage):
-                    raise NotImplementedError()
-            return await self._reconnector.write_with_ack_future(messages)
-        raise NotImplementedError()
+            for index, m in enumerate(messages):
+                messages[index] = PublicMessage._create_message(m)
+
+        futures = await self._reconnector.write_with_ack_future(messages)
+        if input_single_message:
+            return futures[0]
+        else:
+            return futures
 
     async def write(
         self,
         messages: Union[MessageType, List[MessageType]],
-        *args: Optional[MessageType],
     ):
         """
         send one or number of messages to server.