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

194 lines
6.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/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_,
QueryProcessingStage::Enum stage_ = QueryProcessingStage::Complete)
2012-10-18 19:00:49 +00:00
: connection(connection_), query(query_), stage(stage_),
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
{
}
2012-10-18 19:24:46 +00:00
/// Захватывает владение соединением из пула.
RemoteBlockInputStream(ConnectionPool::Entry pool_entry_, const String & query_,
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() + ")"))
{
}
2012-10-18 19:24:46 +00:00
2012-05-17 19:15:53 +00:00
String getName() const { return "RemoteBlockInputStream"; }
BlockInputStreamPtr clone()
{
return pool_entry.isNull()
? new RemoteBlockInputStream(connection, query, stage)
: new RemoteBlockInputStream(pool_entry, query, stage);
}
2012-05-17 19:15:53 +00:00
2012-10-18 19:24:46 +00:00
2012-05-17 19:15:53 +00:00
/** Отменяем умолчальное уведомление о прогрессе,
* так как колбэк прогресса вызывается самостоятельно.
*/
void progress(Block & block) {}
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()
{
2012-10-18 19:00:49 +00:00
/** Если одно из:
* - ничего не начинали делать;
* - получили все пакеты до EndOfStream;
* - получили с сервера эксепшен;
2012-10-19 00:10:17 +00:00
* - объект уничтожается из-за эксепшена;
2012-10-18 19:00:49 +00:00
* - то больше читать ничего не нужно.
*/
2012-10-19 00:10:17 +00:00
if (!sent_query || finished || got_exception_from_server || std::uncaught_exception())
2012-10-18 19:00:49 +00:00
return;
2012-10-19 00:10:17 +00:00
2012-10-18 19:00:49 +00:00
/** Если ещё прочитали не все данные, но они больше не нужны.
2012-10-23 18:51:29 +00:00
* Это может быть из-за того, что данных достаточно (например, при использовании LIMIT).
2012-10-18 19:00:49 +00:00
*/
/// Отправим просьбу прервать выполнение запроса, если ещё не отправляли.
2012-10-19 00:10:17 +00:00
if (!was_cancelled)
{
2012-10-18 23:52:31 +00:00
LOG_TRACE(log, "Cancelling query because enough data has been read");
2012-10-18 19:00:49 +00:00
was_cancelled = true;
connection.sendCancel();
}
2012-10-18 19:00:49 +00:00
/// Получим оставшиеся пакеты, чтобы не было рассинхронизации в соединении с сервером.
while (true)
{
Connection::Packet packet = connection.receivePacket();
switch (packet.type)
{
2012-10-18 19:00:49 +00:00
case Protocol::Server::Data:
case Protocol::Server::Progress:
break;
case Protocol::Server::EndOfStream:
return;
case Protocol::Server::Exception:
2012-10-19 00:10:17 +00:00
got_exception_from_server = true;
packet.exception->rethrow();
2012-10-18 19:00:49 +00:00
break;
default:
2012-10-19 00:10:17 +00:00
throw Exception("Unknown packet from server", ErrorCodes::UNKNOWN_PACKET_FROM_SERVER);
}
}
}
2012-10-20 02:10:47 +00:00
protected:
Block readImpl()
{
if (!sent_query)
{
connection.sendQuery(query, 1, stage);
sent_query = true;
}
while (true)
{
Connection::Packet packet = connection.receivePacket();
switch (packet.type)
{
case Protocol::Server::Data:
if (packet.block)
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:
if (progress_callback)
progress_callback(packet.progress.rows, packet.progress.bytes);
if (!was_cancelled && !finished && isCancelled())
cancel();
break;
default:
throw Exception("Unknown packet from server", ErrorCodes::UNKNOWN_PACKET_FROM_SERVER);
}
}
}
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;
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
};
}