1 | /** |
2 | * Copyright 2021 Alibaba, Inc. and its affiliates. All Rights Reserved. |
3 | * |
4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
5 | * you may not use this file except in compliance with the License. |
6 | * You may obtain a copy of the License at |
7 | * |
8 | * http://www.apache.org/licenses/LICENSE-2.0 |
9 | * |
10 | * Unless required by applicable law or agreed to in writing, software |
11 | * distributed under the License is distributed on an "AS IS" BASIS, |
12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
13 | * See the License for the specific language governing permissions and |
14 | * limitations under the License. |
15 | * |
16 | * \author guonix |
17 | * \date Nov 2020 |
18 | * \brief |
19 | */ |
20 | |
21 | #include "collection_query.h" |
22 | #include "common/error_code.h" |
23 | #include "common/logger.h" |
24 | #include "forward_serializer.h" |
25 | |
26 | namespace proxima { |
27 | namespace be { |
28 | namespace query { |
29 | |
30 | ContextImpl::ContextImpl(uint64_t traceID, index::IndexServicePtr index_service, |
31 | MetaWrapperPtr meta_wrapper, ProfilerPtr profiler_ptr, |
32 | ExecutorPtr executor_ptr) |
33 | : trace_id_(traceID), |
34 | index_service_(std::move(index_service)), |
35 | meta_(std::move(meta_wrapper)), |
36 | executor_(std::move(executor_ptr)), |
37 | profiler_(std::move(profiler_ptr)) {} |
38 | |
39 | //! Unique request id, using trace all relevant information |
40 | uint64_t ContextImpl::id() const { |
41 | return trace_id_; |
42 | } |
43 | |
44 | //! Validate query object, 0 for valid, otherwise non zero returned |
45 | int ContextImpl::validate() const { |
46 | // Not Every query should have response |
47 | if (!index_service_ || !meta_ || !profiler_) { |
48 | return PROXIMA_BE_ERROR_CODE(RuntimeError); |
49 | } |
50 | return 0; |
51 | } |
52 | |
53 | ProfilerPtr ContextImpl::profiler() const { |
54 | return profiler_; |
55 | } |
56 | |
57 | MetaWrapperPtr ContextImpl::meta() const { |
58 | return meta_; |
59 | } |
60 | |
61 | ExecutorPtr ContextImpl::executor() const { |
62 | return executor_; |
63 | } |
64 | |
65 | int ContextImpl::list_segments(index::SegmentPtrList *segments) { |
66 | int code = index_service_->list_segments(collection(), segments); |
67 | if (code != 0) { |
68 | LOG_ERROR("Can't get the segments. collection[%s] code[%d]" , |
69 | collection().c_str(), code); |
70 | return code; |
71 | } |
72 | |
73 | return segments->empty() ? PROXIMA_BE_ERROR_CODE(UnavailableSegment) : 0; |
74 | } |
75 | |
76 | ColumnNameList *ContextImpl::get_forward_columns( |
77 | const index::QueryResult &forward) { |
78 | auto it = revision_to_forward_columns_.find(forward.revision); |
79 | if (it != revision_to_forward_columns_.end()) { |
80 | return &it->second; |
81 | } |
82 | ColumnNameList columns; |
83 | int code = meta_->list_columns(collection(), forward.revision, &columns); |
84 | if (code != 0) { |
85 | LOG_ERROR("Can't get the collection meta with specified revision[%zu]" , |
86 | (size_t)forward.revision); |
87 | return nullptr; |
88 | } |
89 | LOG_DEBUG("Get the collection meta with specified revision[%zu]" , |
90 | (size_t)forward.revision); |
91 | auto iit = revision_to_forward_columns_.emplace(forward.revision, |
92 | std::move(columns)); |
93 | return &iit.first->second; |
94 | } |
95 | |
96 | int ContextImpl::fill_forward(const index::QueryResult &forward, |
97 | proto::Document *doc) { |
98 | auto *columns = get_forward_columns(forward); |
99 | int code = PROXIMA_BE_ERROR_CODE(InvalidRevision); |
100 | if (columns) { |
101 | code = ForwardSerializer::FillForward(forward, *columns, doc); |
102 | } |
103 | return code; |
104 | } |
105 | |
106 | //! Constructor |
107 | CollectionQuery::CollectionQuery(uint64_t traceID, |
108 | const proto::QueryRequest *pb_request, |
109 | index::IndexServicePtr index_service_ptr, |
110 | MetaWrapperPtr meta_wrapper, |
111 | ExecutorPtr executor_ptr, |
112 | ProfilerPtr profiler_ptr, |
113 | proto::QueryResponse *pb_response) |
114 | : ContextImpl(traceID, std::move(index_service_ptr), |
115 | std::move(meta_wrapper), std::move(profiler_ptr), |
116 | std::move(executor_ptr)), |
117 | request_(pb_request), |
118 | response_(pb_response) {} |
119 | |
120 | //! Destructor |
121 | CollectionQuery::~CollectionQuery() = default; |
122 | |
123 | //! Validate query object, 0 for valid, otherwise non zero returned |
124 | int CollectionQuery::validate() const { |
125 | // Not Every query should have response |
126 | int code = ContextImpl::validate(); |
127 | if (code != 0) { |
128 | return code; |
129 | } |
130 | return !request_ || !response_ ? PROXIMA_BE_ERROR_CODE(RuntimeError) : 0; |
131 | } |
132 | |
133 | //! Retrieve the PB request of query |
134 | const proto::QueryRequest *CollectionQuery::request() const { |
135 | return request_; |
136 | } |
137 | |
138 | //! Retrieve the PB response of query |
139 | const proto::QueryResponse *CollectionQuery::response() const { |
140 | return response_; |
141 | } |
142 | |
143 | //! Retrieve the mutable PB response of query |
144 | proto::QueryResponse *CollectionQuery::mutable_response() { |
145 | return response_; |
146 | } |
147 | |
148 | //! Retrieve collection name |
149 | const std::string &CollectionQuery::collection() const { |
150 | return request_->collection_name(); |
151 | } |
152 | |
153 | |
154 | } // namespace query |
155 | } // namespace be |
156 | } // namespace proxima |
157 | |