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

273 lines
9.3 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
{
private:
void init(const Settings * settings_)
{
if (settings_)
{
send_settings = true;
settings = *settings_;
}
else
send_settings = false;
}
public:
/// Принимает готовое соединение.
RemoteBlockInputStream(Connection & connection_, const String & query_, const Settings * settings_,
const Tables & external_tables_ = Tables(), QueryProcessingStage::Enum stage_ = QueryProcessingStage::Complete)
: connection(&connection_), query(query_), external_tables(external_tables_), stage(stage_)
{
init(settings_);
}
/// Принимает готовое соединение. Захватывает владение соединением из пула.
RemoteBlockInputStream(ConnectionPool::Entry & pool_entry_, const String & query_, const Settings * settings_,
const Tables & external_tables_ = Tables(), QueryProcessingStage::Enum stage_ = QueryProcessingStage::Complete)
: pool_entry(pool_entry_), connection(&*pool_entry_), query(query_), external_tables(external_tables_), stage(stage_)
{
init(settings_);
}
/// Принимает пул, из которого нужно будет достать соединение.
RemoteBlockInputStream(IConnectionPool * pool_, const String & query_, const Settings * settings_,
const Tables & external_tables_ = Tables(), QueryProcessingStage::Enum stage_ = QueryProcessingStage::Complete)
: pool(pool_), query(query_), external_tables(external_tables_), stage(stage_)
{
init(settings_);
}
2012-10-18 19:24:46 +00:00
String getName() const override { return "RemoteBlockInputStream"; }
2012-05-17 19:15:53 +00:00
String getID() const override
{
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(const Progress & value) override {}
2012-05-17 19:15:53 +00:00
2012-10-18 19:24:46 +00:00
void cancel() override
2012-10-18 19:24:46 +00:00
{
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, "(" + connection->getServerAddress() + ") Cancelling query");
2012-10-18 19:24:46 +00:00
/// Если запрошено прервать запрос - попросим удалённый сервер тоже прервать запрос.
2014-04-06 23:36:39 +00:00
connection->sendCancel();
was_cancelled = true;
2012-10-18 19:24:46 +00:00
}
}
~RemoteBlockInputStream() override
{
/** Если прервались в середине цикла общения с сервером, то закрываем соединение,
* чтобы оно не осталось висеть в рассихронизированном состоянии.
*/
if (sent_query && !finished)
2014-04-06 23:36:39 +00:00
connection->disconnect();
}
2012-10-20 02:10:47 +00:00
protected:
/// Отправить на удаленные сервера все временные таблицы
void sendExternalTables()
{
ExternalTablesData res;
for (const auto & table : external_tables)
{
StoragePtr cur = table.second;
QueryProcessingStage::Enum stage = QueryProcessingStage::Complete;
DB::BlockInputStreams input = cur->read(cur->getColumnNamesList(), ASTPtr(), settings, stage, DEFAULT_BLOCK_SIZE, 1);
if (input.size() == 0)
res.push_back(std::make_pair(new OneBlockInputStream(cur->getSampleBlock()), table.first));
else
res.push_back(std::make_pair(input[0], table.first));
}
2014-04-06 23:36:39 +00:00
connection->sendExternalTablesData(res);
}
Block readImpl() override
2012-10-20 02:10:47 +00:00
{
if (!sent_query)
{
/// Если надо - достаём соединение из пула.
if (pool)
{
pool_entry = pool->get(send_settings ? &settings : nullptr);
connection = &*pool_entry;
}
connection->sendQuery(query, "", stage, send_settings ? &settings : nullptr, true);
sendExternalTables();
2012-10-20 02:10:47 +00:00
sent_query = true;
}
while (true)
{
2014-04-06 23:36:39 +00:00
Connection::Packet packet = connection->receivePacket();
2012-10-20 02:10:47 +00:00
switch (packet.type)
{
case Protocol::Server::Data:
/// Если блок не пуст и не является заголовочным блоком
if (packet.block && packet.block.rows() > 0)
2012-10-20 02:10:47 +00:00
return packet.block;
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);
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);
}
}
}
void readSuffixImpl() override
2013-09-13 20:33:09 +00:00
{
/** Если одно из:
* - ничего не начинали делать;
* - получили все пакеты до EndOfStream;
* - получили с сервера эксепшен;
* - то больше читать ничего не нужно.
*/
if (!sent_query || finished || got_exception_from_server)
return;
/** Если ещё прочитали не все данные, но они больше не нужны.
* Это может быть из-за того, что данных достаточно (например, при использовании LIMIT).
*/
/// Отправим просьбу прервать выполнение запроса, если ещё не отправляли.
if (!was_cancelled)
{
LOG_TRACE(log, "(" + connection->getServerAddress() + ") Cancelling query because enough data has been read");
2013-09-13 20:33:09 +00:00
was_cancelled = true;
2014-04-06 23:36:39 +00:00
connection->sendCancel();
2013-09-13 20:33:09 +00:00
}
/// Получим оставшиеся пакеты, чтобы не было рассинхронизации в соединении с сервером.
while (true)
{
2014-04-06 23:36:39 +00:00
Connection::Packet packet = connection->receivePacket();
2013-09-13 20:33:09 +00:00
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:
IConnectionPool * pool = nullptr;
ConnectionPool::Entry pool_entry;
Connection * connection = nullptr;
2012-05-23 19:51:30 +00:00
const String query;
bool send_settings;
Settings settings;
/// Временные таблицы, которые необходимо переслать на удаленные сервера.
Tables external_tables;
2012-05-17 19:15:53 +00:00
QueryProcessingStage::Enum stage;
2012-10-18 19:00:49 +00:00
/// Отправили запрос (это делается перед получением первого блока).
bool sent_query = false;
2012-10-18 19:00:49 +00:00
/** Получили все данные от сервера, до пакета EndOfStream.
* Если при уничтожении объекта, ещё не все данные считаны,
* то для того, чтобы не было рассинхронизации, на сервер отправляется просьба прервать выполнение запроса,
* и после этого считываются все пакеты до EndOfStream.
*/
bool finished = false;
2012-10-18 19:00:49 +00:00
/** На сервер была отправлена просьба прервать выполенение запроса, так как данные больше не нужны.
* Это может быть из-за того, что данных достаточно (например, при использовании LIMIT),
* или если на стороне клиента произошло исключение.
*/
bool was_cancelled = false;
2012-10-18 19:00:49 +00:00
/// С сервера было получено исключение. В этом случае получать больше пакетов или просить прервать запрос не нужно.
bool got_exception_from_server = false;
2012-10-18 19:00:49 +00:00
Logger * log = &Logger::get("RemoteBlockInputStream");
2012-05-17 19:15:53 +00:00
};
}