ClickHouse/dbms/include/DB/IO/WriteBufferFromHTTPServerResponse.h

157 lines
5.3 KiB
C++
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#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 add_cors_header;
bool compress;
Poco::DeflatingStreamBuf::StreamType compression_method;
int compression_level = Z_DEFAULT_COMPRESSION;
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 (add_cors_header)
{
response.set("Access-Control-Allow-Origin","*");
}
if (compress && offset()) /// Пустой ответ сжимать не нужно.
{
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, compression_level);
ostr = &deflating_stream.value();
}
else
{
response_ostr = &response.send();
ostr = response_ostr;
}
}
}
void nextImpl() override
{
if (!offset())
return;
sendHeaders();
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();
}
/** Включить или отключить сжатие.
* Работает только перед тем, как были отправлены HTTP заголовки.
* Иначе - не имеет эффекта.
*/
void setCompression(bool enable_compression)
{
compress = enable_compression;
}
/** Установить уровень сжатия, если данные будут сжиматься.
* Работает только перед тем, как были отправлены HTTP заголовки.
* Иначе - не имеет эффекта.
*/
void setCompressionLevel(int level)
{
compression_level = level;
}
/** Включить или отключить CORS.
* Работает только перед тем, как были отправлены HTTP заголовки.
* Иначе - не имеет эффекта.
*/
void addHeaderCORS(bool enable_cors)
{
add_cors_header = enable_cors;
}
~WriteBufferFromHTTPServerResponse()
{
if (!offset())
return;
try
{
next();
if (deflating_stream)
deflating_stream->close();
}
catch (...)
{
tryLogCurrentException(__PRETTY_FUNCTION__);
}
}
};
}