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
26namespace proxima {
27namespace be {
28namespace query {
29
30ContextImpl::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
40uint64_t ContextImpl::id() const {
41 return trace_id_;
42}
43
44//! Validate query object, 0 for valid, otherwise non zero returned
45int 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
53ProfilerPtr ContextImpl::profiler() const {
54 return profiler_;
55}
56
57MetaWrapperPtr ContextImpl::meta() const {
58 return meta_;
59}
60
61ExecutorPtr ContextImpl::executor() const {
62 return executor_;
63}
64
65int 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
76ColumnNameList *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
96int 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
107CollectionQuery::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
121CollectionQuery::~CollectionQuery() = default;
122
123//! Validate query object, 0 for valid, otherwise non zero returned
124int 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
134const proto::QueryRequest *CollectionQuery::request() const {
135 return request_;
136}
137
138//! Retrieve the PB response of query
139const proto::QueryResponse *CollectionQuery::response() const {
140 return response_;
141}
142
143//! Retrieve the mutable PB response of query
144proto::QueryResponse *CollectionQuery::mutable_response() {
145 return response_;
146}
147
148//! Retrieve collection name
149const std::string &CollectionQuery::collection() const {
150 return request_->collection_name();
151}
152
153
154} // namespace query
155} // namespace be
156} // namespace proxima
157