From 81ab6b118566865e84c2c54f3f84c006732ac540 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mat=C4=9Bj=20Jehli=C4=8Dka?= <matej.jehlicka@cdn77.com>
Date: Thu, 5 Nov 2026 11:32:39 +0100
Subject: [PATCH] [client-v2] Fix compressed buffer size for incompressible
 request data

---
 .../internal/CompressedBlockOutputStream.java |  8 +--
 .../CompressedBlockOutputStreamTest.java      | 72 +++++++++++++++++++
 2 files changed, 76 insertions(+), 4 deletions(-)
 create mode 100644 client-v2/src/test/java/com/clickhouse/client/api/internal/CompressedBlockOutputStreamTest.java

diff --git a/client-v2/src/main/java/com/clickhouse/client/api/internal/CompressedBlockOutputStream.java b/client-v2/src/main/java/com/clickhouse/client/api/internal/CompressedBlockOutputStream.java
index 750bdf80c..a808432a7 100644
--- a/client-v2/src/main/java/com/clickhouse/client/api/internal/CompressedBlockOutputStream.java
+++ b/client-v2/src/main/java/com/clickhouse/client/api/internal/CompressedBlockOutputStream.java
@@ -21,12 +21,12 @@ public abstract class CompressedBlockOutputStream extends OutputStream {
 
     private final ByteBuffer compressedBuffer;
 
-    public CompressedBlockOutputStream(OutputStream out, int bufferSize) {
+    public CompressedBlockOutputStream(OutputStream out, int bufferSize, int maxCompressedBufferSize) {
         super();
         LOG.debug("Using compressor with buffer size {}", bufferSize);
         this.inBuffer = ByteBuffer.allocate(bufferSize);
         this.out = out;
-        this.compressedBuffer = ByteBuffer.allocate(bufferSize + CompressedBlockInputStream.HEADER_LENGTH);
+        this.compressedBuffer = ByteBuffer.allocate(maxCompressedBufferSize + CompressedBlockInputStream.HEADER_LENGTH);
     }
 
     @Override
@@ -116,7 +116,7 @@ public static class LZ4OutputStream extends CompressedBlockOutputStream {
         private final LZ4Compressor compressor;
 
         public LZ4OutputStream(OutputStream out, LZ4Compressor compressor, int bufferSize) {
-            super(out, bufferSize);
+            super(out, bufferSize, compressor.maxCompressedLength(bufferSize));
             this.compressor = compressor;
         }
 
@@ -137,7 +137,7 @@ public static class ZSTDOutputStream extends CompressedBlockOutputStream {
         private static final int COMPRESSION_LEVEL = 3;
 
         public ZSTDOutputStream(OutputStream out, int bufferSize) {
-            super(out, bufferSize);
+            super(out, bufferSize, (int) Zstd.compressBound(bufferSize));
         }
 
         @Override
diff --git a/client-v2/src/test/java/com/clickhouse/client/api/internal/CompressedBlockOutputStreamTest.java b/client-v2/src/test/java/com/clickhouse/client/api/internal/CompressedBlockOutputStreamTest.java
new file mode 100644
index 000000000..b69720828
--- /dev/null
+++ b/client-v2/src/test/java/com/clickhouse/client/api/internal/CompressedBlockOutputStreamTest.java
@@ -0,0 +1,72 @@
+package com.clickhouse.client.api.internal;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.charset.StandardCharsets;
+import java.util.Random;
+import java.util.function.Function;
+
+import net.jpountz.lz4.LZ4Factory;
+import org.testng.Assert;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+public class CompressedBlockOutputStreamTest {
+
+    private static final LZ4Factory LZ4_FACTORY = LZ4Factory.fastestInstance();
+
+    private static final int BUFFER_SIZE = CompressedBlockOutputStream.UNCOMPRESSED_BUFF_SIZE;
+
+    @DataProvider(name = "streamsAndPayloads")
+    public static Object[][] streamsAndPayloads() {
+        Function<OutputStream, CompressedBlockOutputStream> lz4 =
+                out -> new CompressedBlockOutputStream.LZ4OutputStream(out, LZ4_FACTORY.fastCompressor(), BUFFER_SIZE);
+        Function<OutputStream, CompressedBlockOutputStream> zstd =
+                out -> new CompressedBlockOutputStream.ZSTDOutputStream(out, BUFFER_SIZE);
+
+        return new Object[][] {
+                { "LZ4", "incompressible", lz4, incompressibleData(BUFFER_SIZE) },
+                { "ZSTD", "incompressible", zstd, incompressibleData(BUFFER_SIZE) },
+                { "LZ4", "compressible", lz4, compressibleData(BUFFER_SIZE) },
+                { "ZSTD", "compressible", zstd, compressibleData(BUFFER_SIZE) },
+        };
+    }
+
+    @Test(dataProvider = "streamsAndPayloads")
+    public void testFullBlockRoundTrip(String method, String dataKind,
+                                       Function<OutputStream, CompressedBlockOutputStream> streamFactory,
+                                       byte[] payload) throws IOException {
+        ByteArrayOutputStream compressed = new ByteArrayOutputStream();
+        try (CompressedBlockOutputStream out = streamFactory.apply(compressed)) {
+            out.write(payload);
+        }
+
+        try (CompressedBlockInputStream in = new CompressedBlockInputStream(
+                new ByteArrayInputStream(compressed.toByteArray()), LZ4_FACTORY.fastDecompressor(), BUFFER_SIZE)) {
+            ByteArrayOutputStream decompressed = new ByteArrayOutputStream();
+            byte[] buf = new byte[8192];
+            int n;
+            while ((n = in.read(buf)) != -1) {
+                decompressed.write(buf, 0, n);
+            }
+            Assert.assertEquals(decompressed.toByteArray(), payload, method + " " + dataKind + " round trip");
+        }
+    }
+
+    private static byte[] incompressibleData(int length) {
+        byte[] data = new byte[length];
+        new Random(42).nextBytes(data);
+        return data;
+    }
+
+    private static byte[] compressibleData(int length) {
+        byte[] pattern = "ClickHouse compressed block test line. ".getBytes(StandardCharsets.US_ASCII);
+        byte[] data = new byte[length];
+        for (int i = 0; i < length; i++) {
+            data[i] = pattern[i % pattern.length];
+        }
+        return data;
+    }
+}