1#include <IO/LimitReadBuffer.h>
2#include <Common/Exception.h>
3
4
5namespace DB
6{
7
8namespace ErrorCodes
9{
10 extern const int LIMIT_EXCEEDED;
11}
12
13
14bool LimitReadBuffer::nextImpl()
15{
16 /// Let underlying buffer calculate read bytes in `next()` call.
17 in.position() = position();
18
19 if (bytes >= limit)
20 {
21 if (throw_exception)
22 throw Exception("Limit for LimitReadBuffer exceeded: " + exception_message, ErrorCodes::LIMIT_EXCEEDED);
23 else
24 return false;
25 }
26
27 if (!in.next())
28 return false;
29
30 working_buffer = in.buffer();
31
32 if (limit - bytes < working_buffer.size())
33 working_buffer.resize(limit - bytes);
34
35 return true;
36}
37
38
39LimitReadBuffer::LimitReadBuffer(ReadBuffer & in_, UInt64 limit_, bool throw_exception_, std::string exception_message_)
40 : ReadBuffer(in_.position(), 0), in(in_), limit(limit_), throw_exception(throw_exception_), exception_message(std::move(exception_message_))
41{
42 size_t remaining_bytes_in_buffer = in.buffer().end() - in.position();
43 if (remaining_bytes_in_buffer > limit)
44 remaining_bytes_in_buffer = limit;
45
46 working_buffer = Buffer(in.position(), in.position() + remaining_bytes_in_buffer);
47}
48
49
50LimitReadBuffer::~LimitReadBuffer()
51{
52 /// Update underlying buffer's position in case when limit wasn't reached.
53 if (working_buffer.size() != 0)
54 in.position() = position();
55}
56
57}
58