diff --git a/asynch/proto/streams/buffered.py b/asynch/proto/streams/buffered.py index e4357733..0c1b1eb8 100644 --- a/asynch/proto/streams/buffered.py +++ b/asynch/proto/streams/buffered.py @@ -219,6 +219,7 @@ def __init__( async def flush(self): await self.compressor.write(self.buffer) self.position = 0 + self.buffer = bytearray() class CompressedBufferedReader(BufferedReader): diff --git a/tests/test_proto/test_proto_connection.py b/tests/test_proto/test_proto_connection.py index c620af86..2581a45a 100644 --- a/tests/test_proto/test_proto_connection.py +++ b/tests/test_proto/test_proto_connection.py @@ -10,14 +10,15 @@ from asynch.proto.cs import ServerInfo -@pytest.fixture() -async def proto_conn(config) -> AsyncIterator[ProtoConnection]: +@pytest.fixture(params=[False, "lz4", "zstd"]) +async def proto_conn(request, config) -> AsyncIterator[ProtoConnection]: _conn = ProtoConnection( user=config.user, password=config.password, host=config.host, port=config.port, database=config.database, + compression=request.param, ) await _conn.connect() yield _conn @@ -103,6 +104,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")