Intermediate changes

1 files changed, 48 insertions, 26 deletions

diff --git a/contrib/clickhouse/src/Server/HTTPHandler.cpp b/contrib/clickhouse/src/Server/HTTPHandler.cpp
index 41ed78bc69..a58f3f64e6 100644
--- a/contrib/clickhouse/src/Server/HTTPHandler.cpp
+++ b/contrib/clickhouse/src/Server/HTTPHandler.cpp
@@ -28,6 +28,8 @@
 #include <Common/setThreadName.h>
 #include <Common/typeid_cast.h>
 #include <Parsers/ASTSetQuery.h>
+#include <Processors/Formats/IOutputFormat.h>
+#include <Formats/FormatFactory.h>
 
 #include <base/getFQDNOrHostName.h>
 #include <base/scope_guard.h>
@@ -833,23 +835,40 @@ void HTTPHandler::processQuery(
     customizeContext(request, context, *in_post_maybe_compressed);
     in = has_external_data ? std::move(in_param) : std::make_unique<ConcatReadBuffer>(*in_param, *in_post_maybe_compressed);
 
-    executeQuery(*in, *used_output.out_maybe_delayed_and_compressed, /* allow_into_outfile = */ false, context,
-        [&response, this] (const QueryResultDetails & details)
-        {
-            response.add("X-ClickHouse-Query-Id", details.query_id);
+    auto set_query_result = [&response, this] (const QueryResultDetails & details)
+    {
+        response.add("X-ClickHouse-Query-Id", details.query_id);
 
-            if (content_type_override)
-                response.setContentType(*content_type_override);
-            else if (details.content_type)
-                response.setContentType(*details.content_type);
+        if (content_type_override)
+            response.setContentType(*content_type_override);
+        else if (details.content_type)
+            response.setContentType(*details.content_type);
 
-            if (details.format)
-                response.add("X-ClickHouse-Format", *details.format);
+        if (details.format)
+            response.add("X-ClickHouse-Format", *details.format);
+
+        if (details.timezone)
+            response.add("X-ClickHouse-Timezone", *details.timezone);
+    };
 
-            if (details.timezone)
-                response.add("X-ClickHouse-Timezone", *details.timezone);
+    auto handle_exception_in_output_format = [&](IOutputFormat & output_format)
+    {
+        if (settings.http_write_exception_in_output_format && output_format.supportsWritingException())
+        {
+            output_format.setException(getCurrentExceptionMessage(false));
+            output_format.finalize();
+            used_output.exception_is_written = true;
         }
-    );
+    };
+
+    executeQuery(
+        *in,
+        *used_output.out_maybe_delayed_and_compressed,
+        /* allow_into_outfile = */ false,
+        context,
+        set_query_result,
+        {},
+        handle_exception_in_output_format);
 
     if (used_output.hasDelayed())
     {
@@ -893,7 +912,7 @@ try
         response.setStatusAndReason(exceptionCodeToHTTPStatus(exception_code));
     }
 
-    if (!response.sent() && !used_output.out_maybe_compressed)
+    if (!response.sent() && !used_output.out_maybe_compressed && !used_output.exception_is_written)
     {
         /// If nothing was sent yet and we don't even know if we must compress the response.
         *response.send() << s << std::endl;
@@ -909,21 +928,24 @@ try
             used_output.out_maybe_delayed_and_compressed.reset();
         }
 
-        /// Send the error message into already used (and possibly compressed) stream.
-        /// Note that the error message will possibly be sent after some data.
-        /// Also HTTP code 200 could have already been sent.
+        if (!used_output.exception_is_written)
+        {
+            /// Send the error message into already used (and possibly compressed) stream.
+            /// Note that the error message will possibly be sent after some data.
+            /// Also HTTP code 200 could have already been sent.
 
-        /// If buffer has data, and that data wasn't sent yet, then no need to send that data
-        bool data_sent = used_output.out->count() != used_output.out->offset();
+            /// If buffer has data, and that data wasn't sent yet, then no need to send that data
+            bool data_sent = used_output.out->count() != used_output.out->offset();
 
-        if (!data_sent)
-        {
-            used_output.out_maybe_compressed->position() = used_output.out_maybe_compressed->buffer().begin();
-            used_output.out->position() = used_output.out->buffer().begin();
-        }
+            if (!data_sent)
+            {
+                used_output.out_maybe_compressed->position() = used_output.out_maybe_compressed->buffer().begin();
+                used_output.out->position() = used_output.out->buffer().begin();
+            }
 
-        writeString(s, *used_output.out_maybe_compressed);
-        writeChar('\n', *used_output.out_maybe_compressed);
+            writeString(s, *used_output.out_maybe_compressed);
+            writeChar('\n', *used_output.out_maybe_compressed);
+        }
 
         used_output.out_maybe_compressed->next();
     }