1 | #pragma once |
---|---|
2 | #include <DataStreams/IBlockOutputStream.h> |
3 | #include <DataStreams/IBlockInputStream.h> |
4 | #include <Interpreters/ProcessList.h> |
5 | |
6 | |
7 | namespace DB |
8 | { |
9 | |
10 | |
11 | /// Proxy class which counts number of written block, rows, bytes |
12 | class CountingBlockOutputStream : public IBlockOutputStream |
13 | { |
14 | public: |
15 | CountingBlockOutputStream(const BlockOutputStreamPtr & stream_) |
16 | : stream(stream_) {} |
17 | |
18 | void setProgressCallback(const ProgressCallback & callback) |
19 | { |
20 | progress_callback = callback; |
21 | } |
22 | |
23 | void setProcessListElement(QueryStatus * elem) |
24 | { |
25 | process_elem = elem; |
26 | } |
27 | |
28 | const Progress & getProgress() const |
29 | { |
30 | return progress; |
31 | } |
32 | |
33 | Block getHeader() const override { return stream->getHeader(); } |
34 | void write(const Block & block) override; |
35 | |
36 | void writePrefix() override { stream->writePrefix(); } |
37 | void writeSuffix() override { stream->writeSuffix(); } |
38 | void flush() override { stream->flush(); } |
39 | void onProgress(const Progress & current_progress) override { stream->onProgress(current_progress); } |
40 | String getContentType() const override { return stream->getContentType(); } |
41 | |
42 | protected: |
43 | BlockOutputStreamPtr stream; |
44 | Progress progress; |
45 | ProgressCallback progress_callback; |
46 | QueryStatus * process_elem = nullptr; |
47 | }; |
48 | |
49 | } |
50 |