ClickHouse/dbms/src/Client/Connection.cpp

#include <iomanip>

#include <Poco/Net/NetException.h>

#include <common/ClickHouseRevision.h>

#include <DB/Core/Defines.h>
#include <DB/Common/Exception.h>

#include <DB/IO/CompressedReadBuffer.h>
#include <DB/IO/CompressedWriteBuffer.h>
#include <DB/IO/ReadBufferFromPocoSocket.h>
#include <DB/IO/WriteBufferFromPocoSocket.h>
#include <DB/IO/ReadHelpers.h>
#include <DB/IO/WriteHelpers.h>
#include <DB/IO/copyData.h>

#include <DB/DataStreams/NativeBlockInputStream.h>
#include <DB/DataStreams/NativeBlockOutputStream.h>

#include <DB/Client/Connection.h>

#include <DB/Common/NetException.h>
#include <DB/Common/CurrentMetrics.h>


namespace CurrentMetrics
{
	extern const Metric SendExternalTables;
}

namespace DB
{

namespace ErrorCodes
{
	extern const int NETWORK_ERROR;
	extern const int SOCKET_TIMEOUT;
	extern const int SERVER_REVISION_IS_TOO_OLD;
	extern const int UNEXPECTED_PACKET_FROM_SERVER;
	extern const int UNKNOWN_PACKET_FROM_SERVER;
}


void Connection::connect()
{
	try
	{
		if (connected)
			disconnect();

		LOG_TRACE(log_wrapper.get(), "Connecting. Database: " << (default_database.empty() ? "(not specified)" : default_database) << ". User: " << user);

		socket.connect(resolved_address, connect_timeout);
		socket.setReceiveTimeout(receive_timeout);
		socket.setSendTimeout(send_timeout);
		socket.setNoDelay(true);

		in = std::make_shared<ReadBufferFromPocoSocket>(socket);
		out = std::make_shared<WriteBufferFromPocoSocket>(socket);

		connected = true;

		sendHello();
		receiveHello();

		LOG_TRACE(log_wrapper.get(), "Connected to " << server_name
			<< " server version " << server_version_major
			<< "." << server_version_minor
			<< "." << server_revision
			<< ".");
	}
	catch (Poco::Net::NetException & e)
	{
		disconnect();

		/// Add server address to exception. Also Exception will remember stack trace. It's a pity that more precise exception type is lost.
		throw NetException(e.displayText(), "(" + getDescription() + ")", ErrorCodes::NETWORK_ERROR);
	}
	catch (Poco::TimeoutException & e)
	{
		disconnect();

		/// Add server address to exception. Also Exception will remember stack trace. It's a pity that more precise exception type is lost.
		throw NetException(e.displayText(), "(" + getDescription() + ")", ErrorCodes::SOCKET_TIMEOUT);
	}
}


void Connection::disconnect()
{
	//LOG_TRACE(log_wrapper.get(), "Disconnecting");

	socket.close();
	in = nullptr;
	out = nullptr;
	connected = false;
}


void Connection::sendHello()
{
	//LOG_TRACE(log_wrapper.get(), "Sending hello");

	writeVarUInt(Protocol::Client::Hello, *out);
	writeStringBinary((DBMS_NAME " ") + client_name, *out);
	writeVarUInt(DBMS_VERSION_MAJOR, *out);
	writeVarUInt(DBMS_VERSION_MINOR, *out);
	writeVarUInt(ClickHouseRevision::get(), *out);
	writeStringBinary(default_database, *out);
	writeStringBinary(user, *out);
	writeStringBinary(password, *out);

	out->next();
}


void Connection::receiveHello()
{
	//LOG_TRACE(log_wrapper.get(), "Receiving hello");

	/// Receive hello packet.
	UInt64 packet_type = 0;

	readVarUInt(packet_type, *in);
	if (packet_type == Protocol::Server::Hello)
	{
		readStringBinary(server_name, *in);
		readVarUInt(server_version_major, *in);
		readVarUInt(server_version_minor, *in);
		readVarUInt(server_revision, *in);
	}
	else if (packet_type == Protocol::Server::Exception)
		receiveException()->rethrow();
	else
	{
		/// Close connection, to not stay in unsynchronised state.
		disconnect();

		throw NetException("Unexpected packet from server " + getDescription() + " (expected Hello or Exception, got "
			+ String(Protocol::Server::toString(packet_type)) + ")", ErrorCodes::UNEXPECTED_PACKET_FROM_SERVER);
	}
}

void Connection::setDefaultDatabase(const String & database)
{
	default_database = database;
}

const String & Connection::getDefaultDatabase() const
{
	return default_database;
}

const String & Connection::getDescription() const
{
	return description;
}

const String & Connection::getHost() const
{
	return host;
}

UInt16 Connection::getPort() const
{
	return port;
}

void Connection::getServerVersion(String & name, UInt64 & version_major, UInt64 & version_minor, UInt64 & revision)
{
	if (!connected)
		connect();

	name = server_name;
	version_major = server_version_major;
	version_minor = server_version_minor;
	revision = server_revision;
}


void Connection::forceConnected()
{
	if (!connected)
	{
		connect();
	}
	else if (!ping())
	{
		LOG_TRACE(log_wrapper.get(), "Connection was closed, will reconnect.");
		connect();
	}
}

struct PingTimeoutSetter
{
	PingTimeoutSetter(Poco::Net::StreamSocket & socket_, const Poco::Timespan & ping_timeout_)
	: socket(socket_), ping_timeout(ping_timeout_)
	{
		old_send_timeout = socket.getSendTimeout();
		old_receive_timeout = socket.getReceiveTimeout();

		if (old_send_timeout > ping_timeout)
			socket.setSendTimeout(ping_timeout);
		if (old_receive_timeout > ping_timeout)
			socket.setReceiveTimeout(ping_timeout);
	}

	~PingTimeoutSetter()
	{
		socket.setSendTimeout(old_send_timeout);
		socket.setReceiveTimeout(old_receive_timeout);
	}

	Poco::Net::StreamSocket & socket;
	Poco::Timespan ping_timeout;
	Poco::Timespan old_send_timeout;
	Poco::Timespan old_receive_timeout;
};

bool Connection::ping()
{
	// LOG_TRACE(log_wrapper.get(), "Ping");

	PingTimeoutSetter timeout_setter(socket, ping_timeout);
	try
	{
		UInt64 pong = 0;
		writeVarUInt(Protocol::Client::Ping, *out);
		out->next();

		if (in->eof())
			return false;

		readVarUInt(pong, *in);

		/// Could receive late packets with progress. TODO: Maybe possible to fix.
		while (pong == Protocol::Server::Progress)
		{
			receiveProgress();

			if (in->eof())
				return false;

			readVarUInt(pong, *in);
		}

		if (pong != Protocol::Server::Pong)
		{
			throw Exception("Unexpected packet from server " + getDescription() + " (expected Pong, got "
				+ String(Protocol::Server::toString(pong)) + ")",
				ErrorCodes::UNEXPECTED_PACKET_FROM_SERVER);
		}
	}
	catch (const Poco::Exception & e)
	{
		LOG_TRACE(log_wrapper.get(), e.displayText());
		return false;
	}

	return true;
}


void Connection::sendQuery(const String & query, const String & query_id_, UInt64 stage, const Settings * settings, bool with_pending_data)
{
	network_compression_method = settings ? settings->network_compression_method.value : CompressionMethod::LZ4;

	forceConnected();

	query_id = query_id_;

	//LOG_TRACE(log_wrapper.get(), "Sending query");

	writeVarUInt(Protocol::Client::Query, *out);
	writeStringBinary(query_id, *out);

	/// Per query settings.
	if (settings)
		settings->serialize(*out);
	else
		writeStringBinary("", *out);

	writeVarUInt(stage, *out);
	writeVarUInt(compression, *out);

	writeStringBinary(query, *out);

	maybe_compressed_in.reset();
	maybe_compressed_out.reset();
	block_in.reset();
	block_out.reset();

	/// If server version is new enough, send empty block which meand end of data.
	if (server_revision >= DBMS_MIN_REVISION_WITH_TEMPORARY_TABLES && !with_pending_data)
	{
		sendData(Block());
		out->next();
	}
}


void Connection::sendCancel()
{
	//LOG_TRACE(log_wrapper.get(), "Sending cancel");

	writeVarUInt(Protocol::Client::Cancel, *out);
	out->next();
}


void Connection::sendData(const Block & block, const String & name)
{
	//LOG_TRACE(log_wrapper.get(), "Sending data");

	if (!block_out)
	{
		if (compression == Protocol::Compression::Enable)
			maybe_compressed_out = std::make_shared<CompressedWriteBuffer>(*out, network_compression_method);
		else
			maybe_compressed_out = out;

		block_out = std::make_shared<NativeBlockOutputStream>(*maybe_compressed_out, server_revision);
	}

	writeVarUInt(Protocol::Client::Data, *out);

	if (server_revision >= DBMS_MIN_REVISION_WITH_TEMPORARY_TABLES)
		writeStringBinary(name, *out);

	size_t prev_bytes = out->count();

	block.checkNestedArraysOffsets();
	block_out->write(block);
	maybe_compressed_out->next();
	out->next();

	if (throttler)
		throttler->add(out->count() - prev_bytes);
}


void Connection::sendPreparedData(ReadBuffer & input, size_t size, const String & name)
{
	/// NOTE 'Throttler' is not used in this method (could use, but it's not important right now).

	writeVarUInt(Protocol::Client::Data, *out);

	if (server_revision >= DBMS_MIN_REVISION_WITH_TEMPORARY_TABLES)
		writeStringBinary(name, *out);

	if (0 == size)
		copyData(input, *out);
	else
		copyData(input, *out, size);
	out->next();
}


void Connection::sendExternalTablesData(ExternalTablesData & data)
{
	/// If working with older server, don't send any info.
	if (server_revision < DBMS_MIN_REVISION_WITH_TEMPORARY_TABLES)
	{
		out->next();
		return;
	}

	if (data.empty())
	{
		/// Send empty block, which means end of data transfer.
		sendData(Block());
		return;
	}

	Stopwatch watch;
	size_t out_bytes = out ? out->count() : 0;
	size_t maybe_compressed_out_bytes = maybe_compressed_out ? maybe_compressed_out->count() : 0;
	size_t rows = 0;

	CurrentMetrics::Increment metric_increment{CurrentMetrics::SendExternalTables};

	for (auto & elem : data)
	{
		elem.first->readPrefix();
		while (Block block = elem.first->read())
		{
			rows += block.rowsInFirstColumn();
			sendData(block, elem.second);
		}
		elem.first->readSuffix();
	}

	/// Send empty block, which means end of data transfer.
	sendData(Block());

	out_bytes = out->count() - out_bytes;
	maybe_compressed_out_bytes = maybe_compressed_out->count() - maybe_compressed_out_bytes;
	double elapsed = watch.elapsedSeconds();

	std::stringstream msg;
	msg << std::fixed << std::setprecision(3);
	msg << "Sent data for " << data.size() << " external tables, total " << rows << " rows in " << elapsed << " sec., "
		<< static_cast<size_t>(rows / watch.elapsedSeconds()) << " rows/sec., "
		<< maybe_compressed_out_bytes / 1048576.0 << " MiB (" << maybe_compressed_out_bytes / 1048576.0 / watch.elapsedSeconds() << " MiB/sec.)";

	if (compression == Protocol::Compression::Enable)
		msg << ", compressed " << static_cast<double>(maybe_compressed_out_bytes) / out_bytes << " times to "
			<< out_bytes / 1048576.0 << " MiB (" << out_bytes / 1048576.0 / watch.elapsedSeconds() << " MiB/sec.)";
	else
		msg << ", no compression.";

	LOG_DEBUG(log_wrapper.get(), msg.rdbuf());
}


bool Connection::poll(size_t timeout_microseconds)
{
	return static_cast<ReadBufferFromPocoSocket &>(*in).poll(timeout_microseconds);
}


bool Connection::hasReadBufferPendingData() const
{
	return static_cast<const ReadBufferFromPocoSocket &>(*in).hasPendingData();
}


Connection::Packet Connection::receivePacket()
{
	//LOG_TRACE(log_wrapper.get(), "Receiving packet");

	try
	{
		Packet res;
		readVarUInt(res.type, *in);

		switch (res.type)
		{
			case Protocol::Server::Data:
				res.block = receiveData();
				return res;

			case Protocol::Server::Exception:
				res.exception = receiveException();
				return res;

			case Protocol::Server::Progress:
				res.progress = receiveProgress();
				return res;

			case Protocol::Server::ProfileInfo:
				res.profile_info = receiveProfileInfo();
				return res;

			case Protocol::Server::Totals:
				/// Block with total values is passed in same form as ordinary block. The only difference is packed id.
				res.block = receiveData();
				return res;

			case Protocol::Server::Extremes:
				/// Same as above.
				res.block = receiveData();
				return res;

			case Protocol::Server::EndOfStream:
				return res;

			default:
				/// In unknown state, disconnect - to not leave unsynchronised connection.
				disconnect();
				throw Exception("Unknown packet "
					+ toString(res.type)
					+ " from server " + getDescription(), ErrorCodes::UNKNOWN_PACKET_FROM_SERVER);
		}
	}
	catch (Exception & e)
	{
		/// Add server address to exception message, if need.
		if (e.code() != ErrorCodes::UNKNOWN_PACKET_FROM_SERVER)
			e.addMessage("while receiving packet from " + getDescription());

		throw;
	}
}


Block Connection::receiveData()
{
	//LOG_TRACE(log_wrapper.get(), "Receiving data");

	initBlockInput();

	String external_table_name;

	if (server_revision >= DBMS_MIN_REVISION_WITH_TEMPORARY_TABLES)
		readStringBinary(external_table_name, *in);

	size_t prev_bytes = in->count();

	/// Read one block from network.
	Block res = block_in->read();

	if (throttler)
		throttler->add(in->count() - prev_bytes);

	return res;
}


void Connection::initBlockInput()
{
	if (!block_in)
	{
		if (compression == Protocol::Compression::Enable)
			maybe_compressed_in = std::make_shared<CompressedReadBuffer>(*in);
		else
			maybe_compressed_in = in;

		block_in = std::make_shared<NativeBlockInputStream>(*maybe_compressed_in, server_revision);
	}
}


void Connection::setDescription()
{
	description = host + ":" + toString(resolved_address.port());
	auto ip_address =  resolved_address.host().toString();

	if (host != ip_address)
		description += ", " + ip_address;
}


std::unique_ptr<Exception> Connection::receiveException()
{
	//LOG_TRACE(log_wrapper.get(), "Receiving exception");

	Exception e;
	readException(e, *in, "Received from " + getDescription());
	return std::unique_ptr<Exception>{ e.clone() };
}


Progress Connection::receiveProgress()
{
	//LOG_TRACE(log_wrapper.get(), "Receiving progress");

	Progress progress;
	progress.read(*in, server_revision);
	return progress;
}


BlockStreamProfileInfo Connection::receiveProfileInfo()
{
	BlockStreamProfileInfo profile_info;
	profile_info.read(*in);
	return profile_info;
}

void Connection::fillBlockExtraInfo(BlockExtraInfo & info) const
{
	info.is_valid = true;
	info.host = host;
	info.resolved_address = resolved_address.toString();
	info.port = port;
	info.user = user;
}

}