1 | #include <IO/LimitReadBuffer.h> |
2 | #include <Common/Exception.h> |
3 | |
4 | |
5 | namespace DB |
6 | { |
7 | |
8 | namespace ErrorCodes |
9 | { |
10 | extern const int LIMIT_EXCEEDED; |
11 | } |
12 | |
13 | |
14 | bool 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 | |
39 | LimitReadBuffer::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 | |
50 | LimitReadBuffer::~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 | |