1#pragma once
2#include <DataStreams/IBlockOutputStream.h>
3#include <DataStreams/IBlockInputStream.h>
4#include <Interpreters/ProcessList.h>
5
6
7namespace DB
8{
9
10
11/// Proxy class which counts number of written block, rows, bytes
12class CountingBlockOutputStream : public IBlockOutputStream
13{
14public:
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
42protected:
43 BlockOutputStreamPtr stream;
44 Progress progress;
45 ProgressCallback progress_callback;
46 QueryStatus * process_elem = nullptr;
47};
48
49}
50