From 3ce4844d4b30f63ed5e7e7c816b1b5fd69f78145 Mon Sep 17 00:00:00 2001
From: Azat Khuzhin <a3at.mail@gmail.com>
Date: Sat, 2 Apr 2022 18:12:51 +0300
Subject: [PATCH] Properly cancel the query after client format error

In case of format error (i.e. select 2 format Template settings
format_template_row='/dev/null') the client will reset the connection,
since it will not expect Data blocks.

To fix this, catch this client error and properly cancel the query. This
should also fix query hang checks (the one that executed after each test).

v2: use getCurrentExceptionMessage()/getCurrentExceptionCode()
Signed-off-by: Azat Khuzhin <a.khuzhin@semrush.com>
---
 src/Client/ClientBase.cpp | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/src/Client/ClientBase.cpp b/src/Client/ClientBase.cpp
index 9e61356566c..4faf3a8028e 100644
--- a/src/Client/ClientBase.cpp
+++ b/src/Client/ClientBase.cpp
@@ -240,6 +240,14 @@ void interruptSignalHandler(int signum)
 }
 
 
+/// To cancel the query on local format error.
+class LocalFormatError : public DB::Exception
+{
+public:
+    using Exception::Exception;
+};
+
+
 ClientBase::~ClientBase() = default;
 ClientBase::ClientBase() = default;
 
@@ -442,6 +450,7 @@ void ClientBase::onProfileInfo(const ProfileInfo & profile_info)
 
 
 void ClientBase::initBlockOutputStream(const Block & block, ASTPtr parsed_query)
+try
 {
     if (!output_format)
     {
@@ -530,6 +539,10 @@ void ClientBase::initBlockOutputStream(const Block & block, ASTPtr parsed_query)
         output_format->setAutoFlush();
     }
 }
+catch (...)
+{
+    throw LocalFormatError(getCurrentExceptionMessage(print_stack_trace), getCurrentExceptionCode());
+}
 
 
 void ClientBase::initLogsOutputStream()
@@ -721,6 +734,9 @@ void ClientBase::receiveResult(ASTPtr parsed_query)
         = std::max(min_poll_interval, std::min<size_t>(receive_timeout.totalMicroseconds(), default_poll_interval));
 
     bool break_on_timeout = connection->getConnectionType() != IServerConnection::Type::LOCAL;
+
+    std::exception_ptr local_format_error;
+
     while (true)
     {
         Stopwatch receive_watch(CLOCK_MONOTONIC_COARSE);
@@ -769,10 +785,21 @@ void ClientBase::receiveResult(ASTPtr parsed_query)
                 break;
         }
 
-        if (!receiveAndProcessPacket(parsed_query, cancelled))
-            break;
+        try
+        {
+            if (!receiveAndProcessPacket(parsed_query, cancelled))
+                break;
+        }
+        catch (const LocalFormatError &)
+        {
+            local_format_error = std::current_exception();
+            connection->sendCancel();
+        }
     }
 
+    if (local_format_error)
+        std::rethrow_exception(local_format_error);
+
     if (cancelled && is_interactive)
         std::cout << "Query was cancelled." << std::endl;
 }