1 | #pragma once |
2 | |
3 | #include <DataStreams/IBlockInputStream.h> |
4 | |
5 | |
6 | namespace DB |
7 | { |
8 | |
9 | /** Initialize another source on the first `read` call, and then use it. |
10 | * This is needed, for example, to read from a table that will be populated |
11 | * after creation of LazyBlockInputStream object, but before the first `read` call. |
12 | */ |
13 | class LazyBlockInputStream : public IBlockInputStream |
14 | { |
15 | public: |
16 | using Generator = std::function<BlockInputStreamPtr()>; |
17 | |
18 | LazyBlockInputStream(const Block & , Generator generator_) |
19 | : header(header_), generator(std::move(generator_)) |
20 | { |
21 | } |
22 | |
23 | LazyBlockInputStream(const char * name_, const Block & , Generator generator_) |
24 | : name(name_), header(header_), generator(std::move(generator_)) |
25 | { |
26 | } |
27 | |
28 | String getName() const override { return name; } |
29 | |
30 | Block () const override |
31 | { |
32 | return header; |
33 | } |
34 | |
35 | /// We call readPrefix lazily. Suppress default behaviour. |
36 | void readPrefix() override {} |
37 | |
38 | protected: |
39 | Block readImpl() override |
40 | { |
41 | if (!input) |
42 | { |
43 | input = generator(); |
44 | |
45 | if (!input) |
46 | return Block(); |
47 | |
48 | auto * p_input = dynamic_cast<IBlockInputStream *>(input.get()); |
49 | |
50 | if (p_input) |
51 | { |
52 | /// They could have been set before, but were not passed into the `input`. |
53 | if (progress_callback) |
54 | p_input->setProgressCallback(progress_callback); |
55 | if (process_list_elem) |
56 | p_input->setProcessListElement(process_list_elem); |
57 | } |
58 | |
59 | input->readPrefix(); |
60 | |
61 | { |
62 | addChild(input); |
63 | |
64 | if (isCancelled() && p_input) |
65 | p_input->cancel(is_killed); |
66 | } |
67 | } |
68 | |
69 | return input->read(); |
70 | } |
71 | |
72 | private: |
73 | const char * name = "Lazy" ; |
74 | Block ; |
75 | Generator generator; |
76 | |
77 | BlockInputStreamPtr input; |
78 | }; |
79 | |
80 | } |
81 | |