1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
|
#include <Parsers/ASTInsertQuery.h>
#include <Interpreters/Context.h>
#include <Interpreters/InterpreterSetQuery.h>
#include <IO/ConcatReadBuffer.h>
#include <IO/ReadBufferFromMemory.h>
#include <IO/ReadBufferFromFile.h>
#include <IO/EmptyReadBuffer.h>
#include <QueryPipeline/BlockIO.h>
#include <Processors/Transforms/getSourceFromASTInsertQuery.h>
#include <Processors/Transforms/AddingDefaultsTransform.h>
#include <Storages/ColumnsDescription.h>
#include <Storages/IStorage.h>
#include <QueryPipeline/Pipe.h>
#include <Processors/Formats/IInputFormat.h>
#include "IO/CompressionMethod.h"
#include "Parsers/ASTLiteral.h"
namespace DB
{
namespace ErrorCodes
{
extern const int LOGICAL_ERROR;
extern const int INVALID_USAGE_OF_INPUT;
extern const int UNKNOWN_TYPE_OF_QUERY;
}
InputFormatPtr getInputFormatFromASTInsertQuery(
const ASTPtr & ast,
bool with_buffers,
const Block & header,
ContextPtr context,
const ASTPtr & input_function)
{
/// get ast query
const auto * ast_insert_query = ast->as<ASTInsertQuery>();
if (!ast_insert_query)
throw Exception(ErrorCodes::LOGICAL_ERROR, "Logical error: query requires data to insert, but it is not INSERT query");
if (ast_insert_query->infile && context->getApplicationType() == Context::ApplicationType::SERVER)
throw Exception(ErrorCodes::UNKNOWN_TYPE_OF_QUERY, "Query has infile and was send directly to server");
if (ast_insert_query->format.empty())
{
if (input_function)
throw Exception(ErrorCodes::INVALID_USAGE_OF_INPUT, "FORMAT must be specified for function input()");
else
throw Exception(ErrorCodes::LOGICAL_ERROR, "Logical error: INSERT query requires format to be set");
}
/// Data could be in parsed (ast_insert_query.data) and in not parsed yet (input_buffer_tail_part) part of query.
auto input_buffer_ast_part = std::make_unique<ReadBufferFromMemory>(
ast_insert_query->data, ast_insert_query->data ? ast_insert_query->end - ast_insert_query->data : 0);
std::unique_ptr<ReadBuffer> input_buffer = with_buffers
? getReadBufferFromASTInsertQuery(ast)
: std::make_unique<EmptyReadBuffer>();
/// Create a source from input buffer using format from query
auto source = context->getInputFormat(ast_insert_query->format, *input_buffer, header, context->getSettingsRef().max_insert_block_size);
source->addBuffer(std::move(input_buffer));
return source;
}
Pipe getSourceFromInputFormat(
const ASTPtr & ast,
InputFormatPtr format,
ContextPtr context,
const ASTPtr & input_function)
{
Pipe pipe(format);
const auto * ast_insert_query = ast->as<ASTInsertQuery>();
if (context->getSettingsRef().input_format_defaults_for_omitted_fields && ast_insert_query->table_id && !input_function)
{
StoragePtr storage = DatabaseCatalog::instance().getTable(ast_insert_query->table_id, context);
auto metadata_snapshot = storage->getInMemoryMetadataPtr();
const auto & columns = metadata_snapshot->getColumns();
if (columns.hasDefaults())
{
pipe.addSimpleTransform([&](const Block & cur_header)
{
return std::make_shared<AddingDefaultsTransform>(cur_header, columns, *format, context);
});
}
}
return pipe;
}
Pipe getSourceFromASTInsertQuery(
const ASTPtr & ast,
bool with_buffers,
const Block & header,
ContextPtr context,
const ASTPtr & input_function)
{
auto format = getInputFormatFromASTInsertQuery(ast, with_buffers, header, context, input_function);
return getSourceFromInputFormat(ast, std::move(format), std::move(context), input_function);
}
std::unique_ptr<ReadBuffer> getReadBufferFromASTInsertQuery(const ASTPtr & ast)
{
const auto * insert_query = ast->as<ASTInsertQuery>();
if (!insert_query)
throw Exception(ErrorCodes::LOGICAL_ERROR, "Logical error: query requires data to insert, but it is not INSERT query");
if (insert_query->infile)
{
/// Data can be from infile
const auto & in_file_node = insert_query->infile->as<ASTLiteral &>();
const auto in_file = in_file_node.value.safeGet<std::string>();
/// It can be compressed and compression method maybe specified in query
std::string compression_method;
if (insert_query->compression)
{
const auto & compression_method_node = insert_query->compression->as<ASTLiteral &>();
compression_method = compression_method_node.value.safeGet<std::string>();
}
/// Otherwise, it will be detected from file name automatically (by chooseCompressionMethod)
/// Buffer for reading from file is created and wrapped with appropriate compression method
return wrapReadBufferWithCompressionMethod(std::make_unique<ReadBufferFromFile>(in_file), chooseCompressionMethod(in_file, compression_method));
}
std::vector<std::unique_ptr<ReadBuffer>> buffers;
if (insert_query->data)
{
/// Data could be in parsed (ast_insert_query.data) and in not parsed yet (input_buffer_tail_part) part of query.
auto ast_buffer = std::make_unique<ReadBufferFromMemory>(
insert_query->data, insert_query->end - insert_query->data);
buffers.emplace_back(std::move(ast_buffer));
}
if (insert_query->tail)
buffers.emplace_back(wrapReadBufferReference(*insert_query->tail));
return std::make_unique<ConcatReadBuffer>(std::move(buffers));
}
}
|
