ClickHouse/dbms/include/DB/DataStreams/RemoteBlockInputStream.h

281 lines
9.5 KiB
C
Raw Normal View History

2012-05-17 19:15:53 +00:00
#pragma once
#include <Yandex/logger_useful.h>
2012-05-17 19:15:53 +00:00
#include <DB/DataStreams/IProfilingBlockInputStream.h>
#include <DB/Common/VirtualColumnUtils.h>
2012-05-17 19:15:53 +00:00
#include <DB/Client/ConnectionPool.h>
2012-05-17 19:15:53 +00:00
namespace DB
{
/** Позволяет выполнить запрос (SELECT) на удалённом сервере и получить результат.
*/
class RemoteBlockInputStream : public IProfilingBlockInputStream
{
public:
RemoteBlockInputStream(Connection & connection_, const String & query_, const Settings * settings_,
QueryProcessingStage::Enum stage_ = QueryProcessingStage::Complete)
: connection(connection_), query(query_), stage(stage_),
2012-10-18 19:00:49 +00:00
sent_query(false), finished(false), was_cancelled(false), got_exception_from_server(false),
log(&Logger::get("RemoteBlockInputStream (" + connection.getServerAddress() + ")"))
2012-05-17 19:15:53 +00:00
{
if (settings_)
{
send_settings = true;
settings = *settings_;
}
else
send_settings = false;
2012-05-17 19:15:53 +00:00
}
2012-10-18 19:24:46 +00:00
/// Захватывает владение соединением из пула.
RemoteBlockInputStream(ConnectionPool::Entry pool_entry_, const String & query_, const Settings * settings_,
QueryProcessingStage::Enum stage_ = QueryProcessingStage::Complete)
: pool_entry(pool_entry_), connection(*pool_entry), query(query_),
stage(stage_), sent_query(false), finished(false), was_cancelled(false),
got_exception_from_server(false), log(&Logger::get("RemoteBlockInputStream (" + connection.getServerAddress() + ")"))
{
if (settings_)
{
send_settings = true;
settings = *settings_;
}
else
send_settings = false;
}
RemoteBlockInputStream(ConnectionPool::Entry pool_entry_, const String & query_, const Settings * settings_,
const String & _host_column_, const String & _port_column_, QueryProcessingStage::Enum stage_ = QueryProcessingStage::Complete)
: pool_entry(pool_entry_), connection(*pool_entry), query(query_), _host_column(_host_column_),
_port_column(_port_column_), stage(stage_), sent_query(false), finished(false), was_cancelled(false),
got_exception_from_server(false), log(&Logger::get("RemoteBlockInputStream (" + connection.getServerAddress() + ")"))
{
if (settings_)
{
send_settings = true;
settings = *settings_;
}
else
send_settings = false;
}
2012-10-18 19:24:46 +00:00
2012-05-17 19:15:53 +00:00
String getName() const { return "RemoteBlockInputStream"; }
String getID() const
{
2013-05-04 04:05:15 +00:00
std::stringstream res;
res << this;
return res.str();
}
2012-10-18 19:24:46 +00:00
2012-05-17 19:15:53 +00:00
/** Отменяем умолчальное уведомление о прогрессе,
* так как колбэк прогресса вызывается самостоятельно.
*/
void progress(size_t rows, size_t bytes) {}
2012-05-17 19:15:53 +00:00
2012-10-18 19:24:46 +00:00
void cancel()
{
2012-11-10 05:13:46 +00:00
if (!__sync_bool_compare_and_swap(&is_cancelled, false, true))
return;
2012-10-20 07:23:13 +00:00
if (sent_query && !was_cancelled && !finished && !got_exception_from_server)
2012-10-18 19:24:46 +00:00
{
LOG_TRACE(log, "Cancelling query");
/// Если запрошено прервать запрос - попросим удалённый сервер тоже прервать запрос.
connection.sendCancel();
was_cancelled = true;
2012-10-18 19:24:46 +00:00
}
}
~RemoteBlockInputStream()
{
/** Если прервались в середине цикла общения с сервером, то закрываем соединение,
* чтобы оно не осталось висеть в рассихронизированном состоянии.
*/
if (sent_query && !finished)
connection.disconnect();
}
2012-10-20 02:10:47 +00:00
protected:
void populateBlock(Block & res)
{
2014-02-16 21:02:20 +00:00
if (!_host_column.empty() && !res.has(_host_column))
{
ColumnPtr column_ptr = ColumnConst<String>(res.rows(), connection.getHost(), new DataTypeString).convertToFullColumn();
ColumnWithNameAndType column(column_ptr, new DataTypeString, _host_column);
res.insert(column);
}
2014-02-16 21:02:20 +00:00
if (!_port_column.empty() && !res.has(_port_column))
{
ColumnPtr column_ptr = ColumnConst<UInt16>(res.rows(), connection.getPort(), new DataTypeUInt16).convertToFullColumn();
ColumnWithNameAndType column(column_ptr, new DataTypeUInt16, _port_column);
res.insert(column);
}
}
2012-10-20 02:10:47 +00:00
Block readImpl()
{
if (!sent_query)
{
connection.sendQuery(query, "", stage, send_settings ? &settings : NULL);
2012-10-20 02:10:47 +00:00
sent_query = true;
}
while (true)
{
Connection::Packet packet = connection.receivePacket();
switch (packet.type)
{
case Protocol::Server::Data:
/// Если блок не пуст и не является заголовочным блоком
if (packet.block && packet.block.rows() > 0)
{
populateBlock(packet.block);
2012-10-20 02:10:47 +00:00
return packet.block;
}
2012-10-20 02:10:47 +00:00
break; /// Если блок пустой - получим другие пакеты до EndOfStream.
case Protocol::Server::Exception:
got_exception_from_server = true;
packet.exception->rethrow();
break;
case Protocol::Server::EndOfStream:
finished = true;
return Block();
case Protocol::Server::Progress:
/** Используем прогресс с удалённого сервера.
* В том числе, запишем его в ProcessList,
* и будем использовать его для проверки
* ограничений (например, минимальная скорость выполнения запроса)
* и квот (например, на количество строчек для чтения).
*/
progressImpl(packet.progress.rows, packet.progress.bytes);
2012-10-20 02:10:47 +00:00
if (!was_cancelled && !finished && isCancelled())
cancel();
break;
2013-06-01 18:57:04 +00:00
case Protocol::Server::ProfileInfo:
info = packet.profile_info;
2013-06-01 18:57:04 +00:00
break;
case Protocol::Server::Totals:
totals = packet.block;
break;
case Protocol::Server::Extremes:
extremes = packet.block;
break;
2012-10-20 02:10:47 +00:00
default:
throw Exception("Unknown packet from server", ErrorCodes::UNKNOWN_PACKET_FROM_SERVER);
}
}
}
2013-09-13 20:33:09 +00:00
void readSuffixImpl()
{
/** Если одно из:
* - ничего не начинали делать;
* - получили все пакеты до EndOfStream;
* - получили с сервера эксепшен;
* - то больше читать ничего не нужно.
*/
if (!sent_query || finished || got_exception_from_server)
return;
/** Если ещё прочитали не все данные, но они больше не нужны.
* Это может быть из-за того, что данных достаточно (например, при использовании LIMIT).
*/
/// Отправим просьбу прервать выполнение запроса, если ещё не отправляли.
if (!was_cancelled)
{
LOG_TRACE(log, "Cancelling query because enough data has been read");
was_cancelled = true;
connection.sendCancel();
}
/// Получим оставшиеся пакеты, чтобы не было рассинхронизации в соединении с сервером.
while (true)
{
Connection::Packet packet = connection.receivePacket();
switch (packet.type)
{
case Protocol::Server::Data:
case Protocol::Server::Progress:
case Protocol::Server::ProfileInfo:
case Protocol::Server::Totals:
case Protocol::Server::Extremes:
break;
case Protocol::Server::EndOfStream:
return;
case Protocol::Server::Exception:
got_exception_from_server = true;
packet.exception->rethrow();
break;
default:
throw Exception("Unknown packet from server", ErrorCodes::UNKNOWN_PACKET_FROM_SERVER);
}
}
finished = true;
2013-09-13 20:33:09 +00:00
}
2012-05-17 19:15:53 +00:00
private:
/// Используется, если нужно владеть соединением из пула
ConnectionPool::Entry pool_entry;
2012-05-17 19:15:53 +00:00
Connection & connection;
2012-05-23 19:51:30 +00:00
const String query;
bool send_settings;
Settings settings;
/// Имя столбца, куда записать имя хоста (Например "_host"). Пустая строка, если записывать не надо.
String _host_column;
/// Имя столбца, куда записать номер порта (Например "_port"). Пустая строка, если записывать не надо.
String _port_column;
2012-05-17 19:15:53 +00:00
QueryProcessingStage::Enum stage;
2012-10-18 19:00:49 +00:00
/// Отправили запрос (это делается перед получением первого блока).
2012-05-17 19:15:53 +00:00
bool sent_query;
2012-10-18 19:00:49 +00:00
/** Получили все данные от сервера, до пакета EndOfStream.
* Если при уничтожении объекта, ещё не все данные считаны,
* то для того, чтобы не было рассинхронизации, на сервер отправляется просьба прервать выполнение запроса,
* и после этого считываются все пакеты до EndOfStream.
*/
bool finished;
2012-10-18 19:00:49 +00:00
/** На сервер была отправлена просьба прервать выполенение запроса, так как данные больше не нужны.
* Это может быть из-за того, что данных достаточно (например, при использовании LIMIT),
* или если на стороне клиента произошло исключение.
*/
bool was_cancelled;
2012-10-18 19:00:49 +00:00
/// С сервера было получено исключение. В этом случае получать больше пакетов или просить прервать запрос не нужно.
bool got_exception_from_server;
Logger * log;
2012-05-17 19:15:53 +00:00
};
}