1 | #include "config_formats.h" |
2 | #if USE_PROTOBUF |
3 | |
4 | #include <Formats/FormatSchemaInfo.h> |
5 | #include <Formats/ProtobufSchemas.h> |
6 | #include <google/protobuf/compiler/importer.h> |
7 | #include <Common/Exception.h> |
8 | |
9 | |
10 | namespace DB |
11 | { |
12 | namespace ErrorCodes |
13 | { |
14 | extern const int BAD_ARGUMENTS; |
15 | extern const int CANNOT_PARSE_PROTOBUF_SCHEMA; |
16 | } |
17 | |
18 | ProtobufSchemas & ProtobufSchemas::instance() |
19 | { |
20 | static ProtobufSchemas instance; |
21 | return instance; |
22 | } |
23 | |
24 | class ProtobufSchemas::ImporterWithSourceTree : public google::protobuf::compiler::MultiFileErrorCollector |
25 | { |
26 | public: |
27 | ImporterWithSourceTree(const String & schema_directory) : importer(&disk_source_tree, this) |
28 | { |
29 | disk_source_tree.MapPath("" , schema_directory); |
30 | } |
31 | |
32 | ~ImporterWithSourceTree() override = default; |
33 | |
34 | const google::protobuf::Descriptor * import(const String & schema_path, const String & message_name) |
35 | { |
36 | // Search the message type among already imported ones. |
37 | const auto * descriptor = importer.pool()->FindMessageTypeByName(message_name); |
38 | if (descriptor) |
39 | return descriptor; |
40 | |
41 | const auto * file_descriptor = importer.Import(schema_path); |
42 | // If there are parsing errors AddError() throws an exception and in this case the following line |
43 | // isn't executed. |
44 | assert(file_descriptor); |
45 | |
46 | descriptor = file_descriptor->FindMessageTypeByName(message_name); |
47 | if (!descriptor) |
48 | throw Exception( |
49 | "Not found a message named '" + message_name + "' in the schema file '" + schema_path + "'" , ErrorCodes::BAD_ARGUMENTS); |
50 | |
51 | return descriptor; |
52 | } |
53 | |
54 | private: |
55 | // Overrides google::protobuf::compiler::MultiFileErrorCollector: |
56 | void AddError(const String & filename, int line, int column, const String & message) override |
57 | { |
58 | throw Exception( |
59 | "Cannot parse '" + filename + "' file, found an error at line " + std::to_string(line) + ", column " + std::to_string(column) |
60 | + ", " + message, |
61 | ErrorCodes::CANNOT_PARSE_PROTOBUF_SCHEMA); |
62 | } |
63 | |
64 | google::protobuf::compiler::DiskSourceTree disk_source_tree; |
65 | google::protobuf::compiler::Importer importer; |
66 | }; |
67 | |
68 | |
69 | ProtobufSchemas::ProtobufSchemas() = default; |
70 | ProtobufSchemas::~ProtobufSchemas() = default; |
71 | |
72 | const google::protobuf::Descriptor * ProtobufSchemas::getMessageTypeForFormatSchema(const FormatSchemaInfo & info) |
73 | { |
74 | auto it = importers.find(info.schemaDirectory()); |
75 | if (it == importers.end()) |
76 | it = importers.emplace(info.schemaDirectory(), std::make_unique<ImporterWithSourceTree>(info.schemaDirectory())).first; |
77 | auto * importer = it->second.get(); |
78 | return importer->import(info.schemaPath(), info.messageName()); |
79 | } |
80 | |
81 | } |
82 | |
83 | #endif |
84 | |