Merge pull request ClickHouse#86875 from ClickHouse/backport/25.8/86606

Backport ClickHouse#86606 to 25.8: Better fix for ShellCommandSource

Author: clickhouse-gh[bot]

src/Processors/Formats/IOutputFormat.h

@@ -107,6 +107,11 @@ class IOutputFormat : public IProcessor
         progress_write_frequency_us = value;
     }
 
+    /// Derived classes can use some wrappers around out WriteBuffer
+    /// and can override this method to return wrapper
+    /// that should be used in its derived classes.
+    virtual WriteBuffer * getWriteBufferPtr() { return &out; }
+
 protected:
     friend class ParallelFormattingOutputFormat;
 
@@ -182,11 +187,6 @@ class IOutputFormat : public IProcessor
     /// outputs them in finalize() method.
     virtual bool areTotalsAndExtremesUsedInFinalize() const { return false; }
 
-    /// Derived classes can use some wrappers around out WriteBuffer
-    /// and can override this method to return wrapper
-    /// that should be used in its derived classes.
-    virtual WriteBuffer * getWriteBufferPtr() { return &out; }
-
     WriteBuffer & out;
 
     Chunk current_chunk;

src/Processors/Formats/Impl/JSONEachRowRowOutputFormat.h

@@ -1,6 +1,5 @@
 #pragma once
 
-#include <IO/PeekableWriteBuffer.h>
 #include <Processors/Formats/OutputFormatWithUTF8ValidationAdaptor.h>
 #include <Processors/Formats/RowOutputFormatWithExceptionHandlerAdaptor.h>
 #include <Formats/FormatSettings.h>

src/Processors/Formats/Impl/JSONRowOutputFormat.h

@@ -1,7 +1,6 @@
 #pragma once
 
 #include <IO/Progress.h>
-#include <IO/PeekableWriteBuffer.h>
 #include <Common/Stopwatch.h>
 #include <Processors/Formats/OutputFormatWithUTF8ValidationAdaptor.h>
 #include <Processors/Formats/RowOutputFormatWithExceptionHandlerAdaptor.h>

src/Processors/Sources/ShellCommandSource.cpp

@@ -566,7 +566,7 @@ namespace
     class SendingChunkHeaderTransform final : public ISimpleTransform
     {
     public:
-        SendingChunkHeaderTransform(SharedHeader header, std::shared_ptr<TimeoutWriteBufferFromFileDescriptor> buffer_)
+        SendingChunkHeaderTransform(SharedHeader header, WriteBuffer & buffer_)
             : ISimpleTransform(header, header, false)
             , buffer(buffer_)
         {
@@ -578,12 +578,12 @@ namespace
 
         void transform(Chunk & chunk) override
         {
-            writeText(chunk.getNumRows(), *buffer);
-            writeChar('\n', *buffer);
+            writeText(chunk.getNumRows(), buffer);
+            writeChar('\n', buffer);
         }
 
     private:
-        std::shared_ptr<TimeoutWriteBufferFromFileDescriptor> buffer;
+        WriteBuffer & buffer;
     };
 
 }
@@ -677,17 +677,19 @@ Pipe ShellCommandSourceCoordinator::createPipe(
 
         input_pipes[i].resize(1);
 
+        auto out = context->getOutputFormat(configuration.format, *timeout_write_buffer, materializeBlock(input_pipes[i].getHeader()));
+        out->setAutoFlush();
+
         if (configuration.send_chunk_header)
         {
-            auto transform = std::make_shared<SendingChunkHeaderTransform>(input_pipes[i].getSharedHeader(), timeout_write_buffer);
+            /// We cannot use timeout_write_buffer directly since the output format may wrap the buffer, so we need to use a wrapper
+            auto transform = std::make_shared<SendingChunkHeaderTransform>(input_pipes[i].getSharedHeader(), *out->getWriteBufferPtr());
             input_pipes[i].addTransform(std::move(transform));
         }
 
         auto num_streams = input_pipes[i].maxParallelStreams();
         auto pipeline = std::make_shared<QueryPipeline>(std::move(input_pipes[i]));
         pipeline->setNumThreads(num_streams);
-        auto out = context->getOutputFormat(configuration.format, *timeout_write_buffer, materializeBlock(pipeline->getHeader()));
-        out->setAutoFlush();
         pipeline->complete(std::move(out));
 
         ShellCommandSource::SendDataTask task = [pipeline, timeout_write_buffer, write_buffer, is_executable_pool]()

tests/integration/test_executable_user_defined_function/functions/test_function_config.xml

@@ -48,23 +48,25 @@
         <name>test_function_send_chunk_header_python</name>
         <return_type>String</return_type>
         <argument>
+            <name>number</name>
             <type>UInt64</type>
         </argument>
-        <format>TabSeparated</format>
+        <format>JSONEachRow</format>
         <send_chunk_header>1</send_chunk_header>
-        <command>input_chunk_header.py</command>
+        <command>input_chunk_header_json.py</command>
     </function>
 
     <function>
         <type>executable_pool</type>
         <name>test_function_send_chunk_header_pool_python</name>
         <return_type>String</return_type>
         <argument>
+            <name>number</name>
             <type>UInt64</type>
         </argument>
-        <format>TabSeparated</format>
+        <format>JSONEachRow</format>
         <send_chunk_header>1</send_chunk_header>
-        <command>input_chunk_header.py</command>
+        <command>input_chunk_header_json.py</command>
     </function>
 
     <function>

tests/integration/test_executable_user_defined_function/test.py

@@ -81,19 +81,23 @@ def test_executable_function_python(started_cluster):
 
 def test_executable_function_send_chunk_header_python(started_cluster):
     skip_test_msan(node)
-    assert (
-        node.query("SELECT test_function_send_chunk_header_python(toUInt64(1))")
-        == "Key 1\n"
-    )
-    assert node.query("SELECT test_function_send_chunk_header_python(1)") == "Key 1\n"
 
-    assert (
-        node.query("SELECT test_function_send_chunk_header_pool_python(toUInt64(1))")
-        == "Key 1\n"
-    )
-    assert (
-        node.query("SELECT test_function_send_chunk_header_pool_python(1)") == "Key 1\n"
-    )
+    for function_name in [
+        "test_function_send_chunk_header_python",
+        "test_function_send_chunk_header_pool_python",
+    ]:
+        assert node.query(f"SELECT {function_name}(toUInt64(1))") == "Key 1\n"
+        assert node.query(f"SELECT {function_name}(1)") == "Key 1\n"
+
+        assert node.query(f"SELECT {function_name}(toUInt64(1))") == "Key 1\n"
+        assert node.query(f"SELECT {function_name}(1)") == "Key 1\n"
+
+        # Test specifically HTTP protocol
+        # This ensures that http_write_exception_in_output_format works as expected
+        assert node.http_query(
+            f"SELECT {function_name}(number) FROM numbers(10)",
+            params={"max_block_size": 3, "http_write_exception_in_output_format": True},
+        ) == "".join(f"Key {i}\n" for i in range(10))
 
 
 def test_executable_function_sum_python(started_cluster):

tests/integration/test_executable_user_defined_function/user_scripts/input_chunk_header.py

@@ -0,0 +1,14 @@
+#!/usr/bin/python3
+
+import sys
+import json
+
+if __name__ == "__main__":
+    for chunk_header in sys.stdin:
+        chunk_length = int(chunk_header)
+
+        for _ in range(chunk_length):
+            row = json.loads(sys.stdin.readline().strip())
+            print(f'{{"result": "Key {row["number"]}"}}')
+
+        sys.stdout.flush()