upgrade protocol 54456

Author: nils-borrmann-tacto

asynch/proto/connection.py

@@ -828,30 +828,49 @@ async def process_insert_query(
         await self.send_query(query_without_data, query_id=query_id)
         await self.send_external_tables(external_tables, types_check=types_check)
 
-        sample_block = await self.receive_sample_block()
+        sample_block = await self._receive_sample_block()
         if sample_block:
             rv = await self.send_data(
                 sample_block, data, types_check=types_check, columnar=columnar
             )
-            packet = await self._receive_packet()
-            if packet.exception:
-                raise packet.exception
+            await self._receive_end_of_stream()
+
             return rv
 
-    async def receive_sample_block(self):
+    async def _receive_sample_block(self):
         while True:
             packet = await self._receive_packet()
 
             if packet.type == ServerPacket.DATA:
                 return packet.block
-
             elif packet.type == ServerPacket.EXCEPTION:
                 raise packet.exception
             elif packet.type == ServerPacket.LOG:
-                self.log_block(packet.block)
+                pass
             elif packet.type == ServerPacket.TABLE_COLUMNS:
                 pass
+            else:
+                message = self.unexpected_packet_message(
+                    "Data, Exception or TableColumns", packet.type
+                )
+                raise UnexpectedPacketFromServerError(message)
 
+    async def _receive_end_of_stream(self):
+        while True:
+            packet = await self._receive_packet()
+
+            if packet.type == ServerPacket.END_OF_STREAM:
+                return
+            elif packet.type == ServerPacket.EXCEPTION:
+                raise packet.exception
+            elif packet.type == ServerPacket.LOG:
+                pass
+            elif packet.type == ServerPacket.PROFILE_INFO:
+                pass
+            elif packet.type == ServerPacket.PROFILE_EVENTS:
+                pass
+            elif packet.type == ServerPacket.PROGRESS:
+                pass
             else:
                 message = self.unexpected_packet_message(
                     "Data, Exception or TableColumns", packet.type

asynch/proto/constants.py

@@ -32,6 +32,7 @@
 DBMS_MIN_PROTOCOL_VERSION_WITH_INCREMENTAL_PROFILE_EVENTS = 54451
 DBMS_MIN_REVISION_WITH_PARALLEL_REPLICAS = 54453
 DBMS_MIN_REVISION_WITH_CUSTOM_SERIALIZATION = 54454
+DBMS_MIN_PROTOCOL_VERSION_WITH_PROFILE_EVENTS_IN_INSERT = 54456
 
 # Timeouts
 DBMS_DEFAULT_CONNECT_TIMEOUT_SEC = 10
@@ -47,7 +48,7 @@
 CLIENT_VERSION_MAJOR = 20
 CLIENT_VERSION_MINOR = 10
 CLIENT_VERSION_PATCH = 2
-CLIENT_REVISION = 54454
+CLIENT_REVISION = 54456
 
 BUFFER_SIZE = 1048576
 

tests/test_proto/test_proto_connection.py

@@ -103,6 +103,17 @@ async def test_execute_with_missing_arg(proto_conn: ProtoConnection):
         await proto_conn.execute(query, args={"foo": 1})
 
 
+@pytest.mark.asyncio
+async def test_large_insert(proto_conn: ProtoConnection):
+    data = [(1,)] * 10_000
+    async with create_table(proto_conn, "a Int64"):
+        await proto_conn.execute(
+            "INSERT INTO test.test (a) VALUES", data, settings={"insert_block_size": 1000}
+        )
+        rv = await proto_conn.execute("SELECT * FROM test.test")
+        assert rv == data
+
+
 @asynccontextmanager
 async def create_table(connection, spec):
     await connection.execute("DROP TABLE IF EXISTS test.test")