mirror of
https://github.com/ClickHouse/ClickHouse.git
synced 2024-11-18 05:32:52 +00:00
123 lines
4.1 KiB
C++
123 lines
4.1 KiB
C++
#pragma once
|
||
|
||
#include <experimental/optional>
|
||
|
||
#include <Poco/Net/HTTPServerResponse.h>
|
||
#include <Poco/DeflatingStream.h>
|
||
|
||
#include <DB/Common/Exception.h>
|
||
|
||
#include <DB/IO/WriteBuffer.h>
|
||
#include <DB/IO/BufferWithOwnMemory.h>
|
||
#include <DB/Common/NetException.h>
|
||
|
||
|
||
namespace DB
|
||
{
|
||
|
||
namespace ErrorCodes
|
||
{
|
||
extern const int CANNOT_WRITE_TO_OSTREAM;
|
||
extern const int LOGICAL_ERROR;
|
||
}
|
||
|
||
|
||
/** Отличается от WriteBufferFromOStream тем, что инициализируется не std::ostream, а Poco::Net::HTTPServerResponse.
|
||
* При первом сбросе данных, получает из него std::ostream (с помощью метода send).
|
||
* Это нужно в HTTP серверах, чтобы после передачи в какой-нибудь метод WriteBuffer-а,
|
||
* но до вывода первых данных клиенту, можно было изменить какие-нибудь HTTP заголовки (например, код ответа).
|
||
* (После вызова Poco::Net::HTTPServerResponse::send() изменить заголовки уже нельзя.)
|
||
* То есть, суть в том, чтобы вызывать метод Poco::Net::HTTPServerResponse::send() не сразу.
|
||
*
|
||
* Дополнительно, позволяет сжимать тело HTTP-ответа, выставив соответствующий заголовок Content-Encoding.
|
||
*/
|
||
class WriteBufferFromHTTPServerResponse : public BufferWithOwnMemory<WriteBuffer>
|
||
{
|
||
private:
|
||
Poco::Net::HTTPServerResponse & response;
|
||
bool compress;
|
||
Poco::DeflatingStreamBuf::StreamType compression_method;
|
||
|
||
std::ostream * response_ostr = nullptr; /// Сюда записывается тело HTTP ответа, возможно, сжатое.
|
||
std::experimental::optional<Poco::DeflatingOutputStream> deflating_stream;
|
||
std::ostream * ostr = nullptr; /// Куда записывать несжатое тело HTTP ответа. Указывает туда же, куда response_ostr или на deflating_stream.
|
||
|
||
void sendHeaders()
|
||
{
|
||
if (!ostr)
|
||
{
|
||
if (compress)
|
||
{
|
||
if (compression_method == Poco::DeflatingStreamBuf::STREAM_GZIP)
|
||
response.set("Content-Encoding", "gzip");
|
||
else if (compression_method == Poco::DeflatingStreamBuf::STREAM_ZLIB)
|
||
response.set("Content-Encoding", "deflate");
|
||
else
|
||
throw Exception("Logical error: unknown compression method passed to WriteBufferFromHTTPServerResponse",
|
||
ErrorCodes::LOGICAL_ERROR);
|
||
|
||
response_ostr = &response.send();
|
||
deflating_stream.emplace(*response_ostr, compression_method);
|
||
ostr = &deflating_stream.value();
|
||
}
|
||
else
|
||
{
|
||
response_ostr = &response.send();
|
||
ostr = response_ostr;
|
||
}
|
||
}
|
||
}
|
||
|
||
void nextImpl()
|
||
{
|
||
sendHeaders();
|
||
|
||
if (!offset())
|
||
return;
|
||
|
||
ostr->write(working_buffer.begin(), offset());
|
||
ostr->flush();
|
||
|
||
if (!ostr->good())
|
||
throw NetException("Cannot write to ostream", ErrorCodes::CANNOT_WRITE_TO_OSTREAM);
|
||
}
|
||
|
||
public:
|
||
WriteBufferFromHTTPServerResponse(
|
||
Poco::Net::HTTPServerResponse & response_,
|
||
bool compress_ = false, /// Если true - выставить заголовок Content-Encoding и сжимать результат.
|
||
Poco::DeflatingStreamBuf::StreamType compression_method_ = Poco::DeflatingStreamBuf::STREAM_GZIP, /// Как сжимать результат (gzip, deflate).
|
||
size_t size = DBMS_DEFAULT_BUFFER_SIZE)
|
||
: BufferWithOwnMemory<WriteBuffer>(size), response(response_),
|
||
compress(compress_), compression_method(compression_method_) {}
|
||
|
||
/** Если данные ещё не были отправлены - отправить хотя бы HTTP заголовки.
|
||
* Используйте эту функцию после того, как данные, возможно, были отправлены,
|
||
* и не было ошибок (вы не планируете поменять код ответа).
|
||
*/
|
||
void finalize()
|
||
{
|
||
sendHeaders();
|
||
}
|
||
|
||
~WriteBufferFromHTTPServerResponse()
|
||
{
|
||
if (!offset())
|
||
return;
|
||
|
||
try
|
||
{
|
||
next();
|
||
|
||
if (deflating_stream)
|
||
deflating_stream->close();
|
||
}
|
||
catch (...)
|
||
{
|
||
tryLogCurrentException(__PRETTY_FUNCTION__);
|
||
}
|
||
}
|
||
};
|
||
|
||
}
|