Merge pull request #35867 from azat/client-format-error-fix

Properly cancel the query after client format error
This commit is contained in:
Alexey Milovidov 2022-04-07 17:42:16 +02:00 committed by GitHub
commit ffb40c4ff4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -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;
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) void ClientBase::initBlockOutputStream(const Block & block, ASTPtr parsed_query)
try
{ {
if (!output_format) if (!output_format)
{ {
@ -530,6 +539,10 @@ void ClientBase::initBlockOutputStream(const Block & block, ASTPtr parsed_query)
output_format->setAutoFlush(); output_format->setAutoFlush();
} }
} }
catch (...)
{
throw LocalFormatError(getCurrentExceptionMessage(print_stack_trace), getCurrentExceptionCode());
}
void ClientBase::initLogsOutputStream() 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)); = std::max(min_poll_interval, std::min<size_t>(receive_timeout.totalMicroseconds(), default_poll_interval));
bool break_on_timeout = connection->getConnectionType() != IServerConnection::Type::LOCAL; bool break_on_timeout = connection->getConnectionType() != IServerConnection::Type::LOCAL;
std::exception_ptr local_format_error;
while (true) while (true)
{ {
Stopwatch receive_watch(CLOCK_MONOTONIC_COARSE); Stopwatch receive_watch(CLOCK_MONOTONIC_COARSE);
@ -769,9 +785,20 @@ void ClientBase::receiveResult(ASTPtr parsed_query)
break; break;
} }
try
{
if (!receiveAndProcessPacket(parsed_query, cancelled)) if (!receiveAndProcessPacket(parsed_query, cancelled))
break; 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) if (cancelled && is_interactive)
std::cout << "Query was cancelled." << std::endl; std::cout << "Query was cancelled." << std::endl;