Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
195 changes: 144 additions & 51 deletions clickhouse/client.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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(); });
}
Expand All @@ -377,6 +381,7 @@ void Client::Impl::BeginExecuteQuery(const Query& query, bool finalize) {
SendQuery(query_, finalize);
}
catch (...) {
InvalidateConnection();
ResetState();
throw;
}
Expand All @@ -399,16 +404,24 @@ std::optional<Block> Client::Impl::NextBlock() {
return {std::move(block)};
}
case VariantIndex<ServerError, decltype(packet)>():
case VariantIndex<std::monostate, decltype(packet)>():
case VariantIndex<EndOfStream, decltype(packet)>():
ResetState();
return std::nullopt;
case VariantIndex<std::monostate, decltype(packet)>():
InvalidateConnection();
ResetState();
return std::nullopt;
default:
continue;
}
}
}
catch (const ServerException&) {
ResetState();
throw;
}
catch (...) {
InvalidateConnection();
ResetState();
throw;
}
Expand All @@ -433,6 +446,7 @@ void Client::Impl::SelectWithExternalData(Query query, const ExternalTables& ext
FinalizeQuery();
}
catch (...) {
InvalidateConnection();
ResetState();
throw;
}
Expand Down Expand Up @@ -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();

Expand All @@ -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) {
Expand All @@ -535,86 +560,126 @@ Block Client::Impl::BeginInsert(Query query) {
throw ValidationError("Query callbacks are not supported in BeginInsert");
}

EnsureConnection();

EnsureNull en(static_cast<QueryEvents*>(&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) {
block = std::move(b);
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) {
if (state_ != State::Inserting) {
throw ValidationError("illegal to send insert data without first calling BeginInsert");
}

SendData(block);
try {
SendData(block);
} catch (...) {
InvalidateConnection();
ResetState();
throw;
}
}

void Client::Impl::EndInsert() {
if (state_ != State::Inserting) {
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() {
if (state_ != State::Idle) {
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;
}
}

Expand Down Expand Up @@ -839,19 +904,35 @@ bool Client::Impl::ProcessPacket(uint64_t* server_packet) {
auto packet = ReceivePacket(server_packet);
switch (packet.index()) {
case VariantIndex<ServerError, decltype(packet)>():
case VariantIndex<std::monostate, decltype(packet)>():
case VariantIndex<EndOfStream, decltype(packet)>():
return false;
case VariantIndex<std::monostate, decltype(packet)>():
InvalidateConnection();
return false;
default:
return true;
}
}

void Client::Impl::ResetState()
{
state_ = State::Idle;
query_ = {};
events_ = nullptr;
query_ = {};
state_ = State::Idle;
}

void Client::Impl::EnsureConnection()
{
if (!socket_) {
ResetConnection();
}
}
Comment on lines +924 to +929

void Client::Impl::InvalidateConnection() noexcept
{
input_.reset();
output_.reset();
socket_.reset();
}

bool Client::Impl::ReadBlock(InputStream& input, Block* block) {
Expand Down Expand Up @@ -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;
}

Comment on lines +1051 to +1054
if (events_) {
events_->OnServerException(*e);
}
Expand All @@ -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() {
Expand All @@ -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()) {
;
}
Expand Down
1 change: 1 addition & 0 deletions ut/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
1 change: 1 addition & 0 deletions ut/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading
Loading