From 4326809a7e4b4c9ef5132045854809c85389ebc7 Mon Sep 17 00:00:00 2001
From: robot-clickhouse <robot-clickhouse@users.noreply.github.com>
Date: Tue, 5 Nov 2024 18:10:43 +0000
Subject: [PATCH] Backport #68975 to 24.8: `Content-Encoding` was not sent on
 progress

---
 .../HTTP/WriteBufferFromHTTPServerResponse.cpp   | 16 +++++++++++++++-
 .../03172_http_content_encoding.reference        |  2 ++
 .../0_stateless/03172_http_content_encoding.sh   | 16 ++++++++++++++++
 3 files changed, 33 insertions(+), 1 deletion(-)
 create mode 100644 tests/queries/0_stateless/03172_http_content_encoding.reference
 create mode 100755 tests/queries/0_stateless/03172_http_content_encoding.sh

diff --git a/src/Server/HTTP/WriteBufferFromHTTPServerResponse.cpp b/src/Server/HTTP/WriteBufferFromHTTPServerResponse.cpp
index 2fcb66ae606..946eaf8aea4 100644
--- a/src/Server/HTTP/WriteBufferFromHTTPServerResponse.cpp
+++ b/src/Server/HTTP/WriteBufferFromHTTPServerResponse.cpp
@@ -83,7 +83,11 @@ void WriteBufferFromHTTPServerResponse::finishSendHeaders()
         return;
 
     if (!headers_started_sending)
+    {
+        if (compression_method != CompressionMethod::None)
+            response.set("Content-Encoding", toContentEncodingName(compression_method));
         startSendHeaders();
+    }
 
     writeHeaderSummary();
     writeExceptionCode();
@@ -105,7 +109,13 @@ void WriteBufferFromHTTPServerResponse::nextImpl()
         initialized = true;
 
         if (compression_method != CompressionMethod::None)
-            response.set("Content-Encoding", toContentEncodingName(compression_method));
+        {
+            /// If we've already sent headers, just send the `Content-Encoding` down the socket directly
+            if (headers_started_sending)
+                socketSendStr("Content-Encoding: " + toContentEncodingName(compression_method) + "\r\n");
+            else
+                response.set("Content-Encoding", toContentEncodingName(compression_method));
+        }
 
         startSendHeaders();
         finishSendHeaders();
@@ -177,8 +187,12 @@ void WriteBufferFromHTTPServerResponse::finalizeImpl()
         /// If no body data just send header
         startSendHeaders();
 
+        /// `finalizeImpl` must be idempotent, so set `initialized` here to not send stuff twice
         if (!initialized && offset() && compression_method != CompressionMethod::None)
+        {
+            initialized = true;
             socketSendStr("Content-Encoding: " + toContentEncodingName(compression_method) + "\r\n");
+        }
 
         finishSendHeaders();
     }
diff --git a/tests/queries/0_stateless/03172_http_content_encoding.reference b/tests/queries/0_stateless/03172_http_content_encoding.reference
new file mode 100644
index 00000000000..1671d148388
--- /dev/null
+++ b/tests/queries/0_stateless/03172_http_content_encoding.reference
@@ -0,0 +1,2 @@
+< Content-Encoding: zstd
+< Content-Encoding: zstd
diff --git a/tests/queries/0_stateless/03172_http_content_encoding.sh b/tests/queries/0_stateless/03172_http_content_encoding.sh
new file mode 100755
index 00000000000..bc55e6d44cc
--- /dev/null
+++ b/tests/queries/0_stateless/03172_http_content_encoding.sh
@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+# shellcheck source=../shell_config.sh
+. "$CURDIR"/../shell_config.sh
+
+URL="${CLICKHOUSE_PORT_HTTP_PROTO}://${CLICKHOUSE_HOST}:${CLICKHOUSE_PORT_HTTP}/"
+
+# with progress
+${CLICKHOUSE_CURL} -vsS "${URL}?send_progress_in_http_headers=1&enable_http_compression=1&wait_end_of_query=0" -o /dev/null \
+  -H 'Accept-Encoding: zstd' --compressed --data-binary @- <<< "select distinct sleep(.1),name from generateRandom('name String',1,1000,2) limit 100009 format TSV" 2>&1 \
+  | perl -lnE 'print if /Content-Encoding/';
+# no progress
+${CLICKHOUSE_CURL} -vsS "${URL}?send_progress_in_http_headers=0&enable_http_compression=1&wait_end_of_query=0" -o /dev/null \
+  -H 'Accept-Encoding: zstd' --compressed --data-binary @- <<< "select distinct sleep(.1),name from generateRandom('name String',1,1000,2) limit 100009 format TSV" 2>&1 \
+  | perl -lnE 'print if /Content-Encoding/';