diff --git a/clickhouse/client.cpp b/clickhouse/client.cpp index 0a8950c7..58766abe 100644 --- a/clickhouse/client.cpp +++ b/clickhouse/client.cpp @@ -237,6 +237,8 @@ class Client::Impl { DecodedPacket ReceivePacket(uint64_t* server_packet = nullptr); bool ProcessPacket(uint64_t* server_packet = nullptr); void ResetState(); + void EnsureConnection(); + void InvalidateConnection() noexcept; void SendQuery(const Query& query, bool finalize = true); void FinalizeQuery(); @@ -365,6 +367,8 @@ void Client::Impl::BeginExecuteQuery(const Query& query, bool finalize) { throw ValidationError("cannot execute query while executing another operation"); } + EnsureConnection(); + if (options_.ping_before_query) { RetryGuard([this]() { Ping(); }); } @@ -377,6 +381,7 @@ void Client::Impl::BeginExecuteQuery(const Query& query, bool finalize) { SendQuery(query_, finalize); } catch (...) { + InvalidateConnection(); ResetState(); throw; } @@ -399,16 +404,24 @@ std::optional Client::Impl::NextBlock() { return {std::move(block)}; } case VariantIndex(): - case VariantIndex(): case VariantIndex(): ResetState(); return std::nullopt; + case VariantIndex(): + InvalidateConnection(); + ResetState(); + return std::nullopt; default: continue; } } } + catch (const ServerException&) { + ResetState(); + throw; + } catch (...) { + InvalidateConnection(); ResetState(); throw; } @@ -433,6 +446,7 @@ void Client::Impl::SelectWithExternalData(Query query, const ExternalTables& ext FinalizeQuery(); } catch (...) { + InvalidateConnection(); ResetState(); throw; } @@ -493,12 +507,12 @@ void Client::Impl::Insert(const std::string& table_name, const std::string& quer throw ValidationError("cannot execute query while executing another operation"); } + EnsureConnection(); + if (options_.ping_before_query) { RetryGuard([this]() { Ping(); }); } - state_ = State::Inserting; - std::stringstream fields_section; const auto num_columns = block.GetColumnCount(); @@ -511,19 +525,30 @@ void Client::Impl::Insert(const std::string& table_name, const std::string& quer } Query query("INSERT INTO " + table_name + " ( " + fields_section.str() + " ) VALUES", query_id); - SendQuery(query); + state_ = State::Inserting; - // Wait for a data packet and return - uint64_t server_packet = 0; - while (ProcessPacket(&server_packet)) { - if (server_packet == ServerCodes::Data) { - SendData(block); - EndInsert(); - return; + try { + SendQuery(query); + + // Wait for a data packet and return + uint64_t server_packet = 0; + while (ProcessPacket(&server_packet)) { + if (server_packet == ServerCodes::Data) { + SendData(block); + EndInsert(); + return; + } } - } - throw ProtocolError("fail to receive data packet"); + throw ProtocolError("fail to receive data packet"); + } catch (const ServerException&) { + ResetState(); + throw; + } catch (...) { + InvalidateConnection(); + ResetState(); + throw; + } } Block Client::Impl::BeginInsert(Query query) { @@ -535,14 +560,14 @@ Block Client::Impl::BeginInsert(Query query) { throw ValidationError("Query callbacks are not supported in BeginInsert"); } + EnsureConnection(); + EnsureNull en(static_cast(&query), &events_); if (options_.ping_before_query) { RetryGuard([this]() { Ping(); }); } - state_ = State::Inserting; - // Create a callback to extract the block with the proper query columns. Block block; query.OnData([&block](const Block& b) { @@ -550,17 +575,28 @@ Block Client::Impl::BeginInsert(Query query) { return true; }); - SendQuery(query); + state_ = State::Inserting; + + try { + SendQuery(query); - // Wait for a data packet and return - uint64_t server_packet = 0; - while (ProcessPacket(&server_packet)) { - if (server_packet == ServerCodes::Data) { - return block; + // Wait for a data packet and return + uint64_t server_packet = 0; + while (ProcessPacket(&server_packet)) { + if (server_packet == ServerCodes::Data) { + return block; + } } - } - throw ProtocolError("fail to receive data packet"); + throw ProtocolError("fail to receive data packet"); + } catch (const ServerException&) { + ResetState(); + throw; + } catch (...) { + InvalidateConnection(); + ResetState(); + throw; + } } void Client::Impl::SendInsertBlock(const Block& block) { @@ -568,7 +604,13 @@ void Client::Impl::SendInsertBlock(const Block& block) { throw ValidationError("illegal to send insert data without first calling BeginInsert"); } - SendData(block); + try { + SendData(block); + } catch (...) { + InvalidateConnection(); + ResetState(); + throw; + } } void Client::Impl::EndInsert() { @@ -576,21 +618,30 @@ void Client::Impl::EndInsert() { return; } - // Send empty block as marker of end of data. - SendData(Block()); + try { + // Send empty block as marker of end of data. + SendData(Block()); - // Wait for EOS. - uint64_t eos_packet{0}; - while (ProcessPacket(&eos_packet)) { - ; - } + // Wait for EOS. + uint64_t eos_packet{0}; + while (ProcessPacket(&eos_packet)) { + ; + } - if (eos_packet != ServerCodes::EndOfStream && eos_packet != ServerCodes::Exception - && eos_packet != ServerCodes::Log && options_.rethrow_exceptions) { - throw ProtocolError(std::string{"unexpected packet from server while receiving end of query, expected (expected Exception, EndOfStream or Log, got: "} - + (eos_packet ? std::to_string(eos_packet) : "nothing") + ")"); + if (eos_packet != ServerCodes::EndOfStream && eos_packet != ServerCodes::Exception + && eos_packet != ServerCodes::Log && options_.rethrow_exceptions) { + throw ProtocolError(std::string{"unexpected packet from server while receiving end of query, expected (expected Exception, EndOfStream or Log, got: "} + + (eos_packet ? std::to_string(eos_packet) : "nothing") + ")"); + } + state_ = State::Idle; + } catch (const ServerException&) { + ResetState(); + throw; + } catch (...) { + InvalidateConnection(); + ResetState(); + throw; } - state_ = State::Idle; } void Client::Impl::Ping() { @@ -598,23 +649,37 @@ void Client::Impl::Ping() { throw ValidationError("cannot execute query while executing another operation"); } - WireFormat::WriteUInt64(*output_, ClientCodes::Ping); - output_->Flush(); + EnsureConnection(); + + try { + WireFormat::WriteUInt64(*output_, ClientCodes::Ping); + output_->Flush(); - uint64_t server_packet; - const bool ret = ProcessPacket(&server_packet); + uint64_t server_packet; + const bool ret = ProcessPacket(&server_packet); - if (!ret || server_packet != ServerCodes::Pong) { - throw ProtocolError("fail to ping server"); + if (!ret || server_packet != ServerCodes::Pong) { + throw ProtocolError("fail to ping server"); + } + } catch (const ServerException&) { + throw; + } catch (...) { + InvalidateConnection(); + throw; } } void Client::Impl::ResetConnection() { InitializeStreams(socket_factory_->connect(options_, current_endpoint_.value())); - state_ = State::Idle; + ResetState(); - if (!Handshake()) { - throw ProtocolError("fail to connect to " + options_.host); + try { + if (!Handshake()) { + throw ProtocolError("fail to connect to " + options_.host); + } + } catch (...) { + InvalidateConnection(); + throw; } } @@ -839,9 +904,11 @@ bool Client::Impl::ProcessPacket(uint64_t* server_packet) { auto packet = ReceivePacket(server_packet); switch (packet.index()) { case VariantIndex(): - case VariantIndex(): case VariantIndex(): return false; + case VariantIndex(): + InvalidateConnection(); + return false; default: return true; } @@ -849,9 +916,23 @@ bool Client::Impl::ProcessPacket(uint64_t* server_packet) { void Client::Impl::ResetState() { - state_ = State::Idle; - query_ = {}; events_ = nullptr; + query_ = {}; + state_ = State::Idle; +} + +void Client::Impl::EnsureConnection() +{ + if (!socket_) { + ResetConnection(); + } +} + +void Client::Impl::InvalidateConnection() noexcept +{ + input_.reset(); + output_.reset(); + socket_.reset(); } bool Client::Impl::ReadBlock(InputStream& input, Block* block) { @@ -967,6 +1048,10 @@ bool Client::Impl::ReceiveException(bool rethrow, ServerError * error) { && WireFormat::ReadString(*input_, &e->stack_trace) && WireFormat::ReadFixed(*input_, &has_nested); + if (!exception_received) { + return false; + } + if (events_) { events_->OnServerException(*e); } @@ -975,10 +1060,10 @@ bool Client::Impl::ReceiveException(bool rethrow, ServerError * error) { throw ServerError(e); } - if (exception_received && error != nullptr) { + if (error != nullptr) { *error = ServerError(e); } - return exception_received; + return true; } void Client::Impl::SendCancel() { @@ -990,7 +1075,15 @@ void Client::Impl::Cancel() { if (state_ != State::Selecting) { throw ValidationError("cannot cancel while not executing a query"); } - SendCancel(); + + try { + SendCancel(); + } catch (...) { + InvalidateConnection(); + ResetState(); + throw; + } + while (NextBlock().has_value()) { ; } diff --git a/ut/BUILD.bazel b/ut/BUILD.bazel index 6845fbb5..7cea2826 100644 --- a/ut/BUILD.bazel +++ b/ut/BUILD.bazel @@ -101,6 +101,7 @@ cc_test( "bignum_samples.cpp", "bignum_round_trip.cpp", "Column_ut.cpp", + "client_socket_state_test.cpp", "client_ut.cpp", "connection_failed_client_test.cpp", "connection_failed_client_test.h", diff --git a/ut/CMakeLists.txt b/ut/CMakeLists.txt index c827e5e3..51a8f46d 100644 --- a/ut/CMakeLists.txt +++ b/ut/CMakeLists.txt @@ -5,6 +5,7 @@ SET ( clickhouse-cpp-ut-src bignum_ut.cpp bignum_round_trip.cpp block_ut.cpp + client_socket_state_test.cpp client_ut.cpp columns_ut.cpp column_as_ut.cpp diff --git a/ut/client_socket_state_test.cpp b/ut/client_socket_state_test.cpp new file mode 100644 index 00000000..e4720867 --- /dev/null +++ b/ut/client_socket_state_test.cpp @@ -0,0 +1,324 @@ +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "utils.h" + +namespace { + +using namespace clickhouse; + +constexpr uint64_t kMarker = 12648430; + +ClientOptions MakeClientOptions() { + return ClientOptions() + .SetHost(getEnvOrDefault("CLICKHOUSE_HOST", "localhost")) + .SetPort(getEnvOrDefault("CLICKHOUSE_PORT", "9000")) + .SetUser(getEnvOrDefault("CLICKHOUSE_USER", "default")) + .SetPassword(getEnvOrDefault("CLICKHOUSE_PASSWORD", "")) + .SetDefaultDatabase(getEnvOrDefault("CLICKHOUSE_DB", "default")) + .SetConnectionRecvTimeout(std::chrono::seconds(5)); +} + +bool IsUnsupportedCompressionError(const ServerException& error) { + // Verified with ClickHouse 25.12. Older releases may reject or ignore compression NONE + switch (error.GetCode()) { + case ErrorCodes::INVALID_SETTING_VALUE: + case ErrorCodes::UNKNOWN_COMPRESSION_METHOD: + case ErrorCodes::UNKNOWN_SETTING: + case ErrorCodes::SUPPORT_IS_DISABLED: + return true; + default: + return false; + } +} + +bool IsUnsupportedJsonError(const ServerException& error) { + // The native JSON type was introduced in ClickHouse 24.8. Earlier releases do + // not have this type. + switch (error.GetCode()) { + case ErrorCodes::NOT_IMPLEMENTED: + case ErrorCodes::UNKNOWN_SETTING: + case ErrorCodes::UNKNOWN_TYPE: + case ErrorCodes::SUPPORT_IS_DISABLED: + return true; + default: + return false; + } +} + +bool SelectReturnsMarker(Client& client) { + bool valid = true; + size_t rows = 0; + + try { + client.Select("SELECT CAST(12648430 AS UInt64) AS marker", [&](const Block& block) { + if (block.GetRowCount() == 0) { + return; + } + if (block.GetColumnCount() != 1) { + valid = false; + return; + } + + const auto column = block[0]->As(); + if (!column) { + valid = false; + return; + } + + for (size_t i = 0; i < column->Size(); ++i) { + valid = valid && column->At(i) == kMarker; + ++rows; + } + }); + } catch (...) { + return false; + } + + return valid && rows == 1; +} + +void ExpectSocketIsHealthy(Client& client) { + // Stale EndOfStream may end the probe without throwing, so verify its exact result. + EXPECT_TRUE(SelectReturnsMarker(client)); + + // ASSERT_NO_THROW(client.ResetConnection()); + // EXPECT_TRUE(SelectReturnsMarker(client)); +} + +class CallbackError final : public std::exception {}; + +class ClientSocketStateTest : public testing::Test { +protected: + void SetUp() override { client_ = std::make_unique(MakeClientOptions()); } + + std::unique_ptr client_; +}; + +TEST_F(ClientSocketStateTest, UnsupportedResultTypeBreaksSocketState) { + EXPECT_THROW(client_->Execute("SELECT sumState(number) FROM numbers(10)"), UnimplementedError); + EXPECT_FALSE(client_->IsSelecting()); + + ExpectSocketIsHealthy(*client_); +} + +TEST_F(ClientSocketStateTest, TotalsBreakSocketState) { + EXPECT_THROW(client_->Execute("SELECT number % 2 AS key, count() " + "FROM numbers(10) GROUP BY key WITH TOTALS"), + UnimplementedError); + EXPECT_FALSE(client_->IsSelecting()); + + ExpectSocketIsHealthy(*client_); +} + +TEST_F(ClientSocketStateTest, ExtremesBreakSocketState) { + Query query("SELECT number FROM numbers(10)"); + query.SetSetting("extremes", {"1"}); + + EXPECT_THROW(client_->Execute(query), UnimplementedError); + EXPECT_FALSE(client_->IsSelecting()); + + ExpectSocketIsHealthy(*client_); +} + +TEST_F(ClientSocketStateTest, ThrowingOnDataBreaksSocketState) { + bool callback_called = false; + Query query( + "SELECT number FROM system.numbers " + "LIMIT 100000 SETTINGS max_block_size = 100"); + query.OnData([&](const Block& block) { + if (block.GetRowCount() > 0) { + callback_called = true; + throw CallbackError{}; + } + }); + + EXPECT_THROW(client_->Execute(query), CallbackError); + EXPECT_TRUE(callback_called); + EXPECT_FALSE(client_->IsSelecting()); + + ExpectSocketIsHealthy(*client_); +} + +TEST_F(ClientSocketStateTest, ThrowingOnDataCancelableBreaksSocketState) { + bool callback_called = false; + Query query( + "SELECT number FROM system.numbers " + "LIMIT 100000 SETTINGS max_block_size = 100"); + query.OnDataCancelable([&](const Block& block) { + if (block.GetRowCount() > 0) { + callback_called = true; + throw CallbackError{}; + } + return true; + }); + + EXPECT_THROW(client_->Execute(query), CallbackError); + EXPECT_TRUE(callback_called); + EXPECT_FALSE(client_->IsSelecting()); + + ExpectSocketIsHealthy(*client_); +} + +TEST_F(ClientSocketStateTest, ThrowingOnProfileBreaksSocketState) { + bool callback_called = false; + Query query("SELECT * FROM system.numbers LIMIT 10"); + query.OnProfile([&](const Profile&) { + callback_called = true; + throw CallbackError{}; + }); + + EXPECT_THROW(client_->Execute(query), CallbackError); + EXPECT_TRUE(callback_called); + EXPECT_FALSE(client_->IsSelecting()); + + ExpectSocketIsHealthy(*client_); +} + +TEST_F(ClientSocketStateTest, ThrowingOnProgressBreaksSocketState) { + client_->Execute("CREATE TEMPORARY TABLE socket_state_progress (value String) ENGINE = Memory"); + + bool callback_called = false; + Query query("INSERT INTO socket_state_progress VALUES ('Foo'), ('Bar')"); + query.OnProgress([&](const Progress&) { + callback_called = true; + throw CallbackError{}; + }); + + EXPECT_THROW(client_->Execute(query), CallbackError); + EXPECT_TRUE(callback_called); + EXPECT_FALSE(client_->IsSelecting()); + + ExpectSocketIsHealthy(*client_); +} + +TEST_F(ClientSocketStateTest, ThrowingOnServerLogBreaksSocketState) { + client_->Execute("CREATE TEMPORARY TABLE socket_state_server_log (value String) ENGINE = Memory"); + + bool callback_called = false; + Query query("INSERT INTO socket_state_server_log VALUES ('Foo'), ('Bar')"); + query.SetSetting("send_logs_level", {"trace"}); + query.OnServerLog([&](const Block&) -> bool { + callback_called = true; + throw CallbackError{}; + }); + + EXPECT_THROW(client_->Execute(query), CallbackError); + EXPECT_TRUE(callback_called); + EXPECT_FALSE(client_->IsSelecting()); + + ExpectSocketIsHealthy(*client_); +} + +TEST_F(ClientSocketStateTest, ThrowingOnProfileEventsBreaksSocketState) { + constexpr uint64_t kMinRevisionWithIncrementalProfileEvents = 54451; + if (client_->GetServerInfo().revision < kMinRevisionWithIncrementalProfileEvents) { + GTEST_SKIP() << "Server does not support incremental profile events"; + } + + client_->Execute("CREATE TEMPORARY TABLE socket_state_profile_events (value String) ENGINE = Memory"); + client_->Execute("INSERT INTO socket_state_profile_events VALUES ('Foo'), ('Bar')"); + + bool callback_called = false; + Query query("SELECT * FROM socket_state_profile_events"); + query.OnProfileEvents([&](const Block&) -> bool { + callback_called = true; + throw CallbackError{}; + }); + + EXPECT_THROW(client_->Execute(query), CallbackError); + EXPECT_TRUE(callback_called); + EXPECT_FALSE(client_->IsSelecting()); + + ExpectSocketIsHealthy(*client_); +} + +TEST_F(ClientSocketStateTest, ReceiveTimeoutBreaksSocketState) { + ClientOptions options = MakeClientOptions(); + options.SetConnectionRecvTimeout(std::chrono::milliseconds(500)); + Client client(options); + + Query query("SELECT sleep(1)"); + query.SetSetting("interactive_delay", {"10000000"}); + + EXPECT_THROW(client.Execute(query), std::system_error); + EXPECT_FALSE(client.IsSelecting()); + + std::this_thread::sleep_for(std::chrono::seconds(2)); + ExpectSocketIsHealthy(client); +} + +TEST_F(ClientSocketStateTest, NoneNetworkCompressionBreaksSocketState) { + ClientOptions options = MakeClientOptions(); + options.SetCompressionMethod(CompressionMethod::LZ4); + Client client(options); + + Query query("SELECT number FROM numbers(10)"); + query.SetSetting("network_compression_method", {"NONE"}); + + bool compression_error = false; + try { + client.Execute(query); + } catch (const CompressionError&) { + compression_error = true; + } catch (const ServerException& error) { + if (IsUnsupportedCompressionError(error)) { + GTEST_SKIP() << "Server does not support NONE network compression: " << error.what(); + } + throw; + } + + if (!compression_error) { + GTEST_SKIP() << "Server did not emit NONE-compressed frames"; + } + EXPECT_FALSE(client.IsSelecting()); + + ExpectSocketIsHealthy(client); +} + +TEST_F(ClientSocketStateTest, NativeJsonSerializationBreaksSocketState) { + try { + Query create("CREATE TEMPORARY TABLE socket_state_json (value JSON) ENGINE = Memory"); + create.SetSetting("allow_experimental_json_type", {"1"}); + client_->Execute(create); + + Query insert("INSERT INTO socket_state_json VALUES ('{\"key\": 1}')"); + insert.SetSetting("allow_experimental_json_type", {"1"}); + client_->Execute(insert); + } catch (const ServerException& error) { + if (IsUnsupportedJsonError(error)) { + GTEST_SKIP() << "Server does not support the JSON test setup: " << error.what(); + } + throw; + } + + Query query("SELECT value FROM socket_state_json"); + query.SetSetting("allow_experimental_json_type", {"1"}); + + bool protocol_error = false; + try { + client_->Execute(query); + } catch (const ProtocolError&) { + protocol_error = true; + } + + if (!protocol_error) { + GTEST_SKIP() << "Server did not emit native JSON serialization"; + } + EXPECT_FALSE(client_->IsSelecting()); + + ExpectSocketIsHealthy(*client_); +} + +} // namespace