github.com/rohankumardubey/aresdb@v0.0.2-0.20190517170215-e54e3ca06b9c/api/query_handler.go (about) 1 // Copyright (c) 2017-2018 Uber Technologies, Inc. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package api 16 17 import ( 18 "encoding/json" 19 "net/http" 20 21 "github.com/uber/aresdb/memstore" 22 "github.com/uber/aresdb/query" 23 queryCom "github.com/uber/aresdb/query/common" 24 "github.com/uber/aresdb/utils" 25 26 "time" 27 28 "github.com/gorilla/mux" 29 "github.com/uber/aresdb/common" 30 ) 31 32 // QueryHandler handles query execution. 33 type QueryHandler struct { 34 memStore memstore.MemStore 35 deviceManager *query.DeviceManager 36 } 37 38 // NewQueryHandler creates a new QueryHandler. 39 func NewQueryHandler(memStore memstore.MemStore, cfg common.QueryConfig) *QueryHandler { 40 return &QueryHandler{ 41 memStore: memStore, 42 deviceManager: query.NewDeviceManager(cfg), 43 } 44 } 45 46 // GetDeviceManager returns the device manager of query handler. 47 func (handler *QueryHandler) GetDeviceManager() *query.DeviceManager { 48 return handler.deviceManager 49 } 50 51 // Register registers http handlers. 52 func (handler *QueryHandler) Register(router *mux.Router, wrappers ...utils.HTTPHandlerWrapper) { 53 router.HandleFunc("/aql", utils.ApplyHTTPWrappers(handler.HandleAQL, wrappers)).Methods(http.MethodGet, http.MethodPost) 54 router.HandleFunc("/sql", utils.ApplyHTTPWrappers(handler.HandleSQL, wrappers)).Methods(http.MethodGet, http.MethodPost) 55 } 56 57 // HandleAQL swagger:route POST /query/aql queryAQL 58 // query in AQL 59 // 60 // Consumes: 61 // - application/json 62 // - application/hll 63 // 64 // Produces: 65 // - application/json 66 // 67 // Responses: 68 // default: errorResponse 69 // 200: aqlResponse 70 // 400: aqlResponse 71 func (handler *QueryHandler) HandleAQL(w http.ResponseWriter, r *http.Request) { 72 // default device to negative value to differentiate 0 from empty 73 aqlRequest := AQLRequest{Device: -1} 74 75 if err := ReadRequest(r, &aqlRequest); err != nil { 76 RespondWithBadRequest(w, err) 77 utils.GetLogger().With( 78 "error", err, 79 "statusCode", http.StatusBadRequest, 80 ).Error("failed to parse query") 81 return 82 } 83 84 handler.handleAQLInternal(aqlRequest, w, r) 85 } 86 87 func (handler *QueryHandler) handleAQLInternal(aqlRequest AQLRequest, w http.ResponseWriter, r *http.Request) { 88 var err error 89 var duration time.Duration 90 var qcs []*query.AQLQueryContext 91 var statusCode int 92 93 defer func() { 94 var errStr string 95 if err != nil { 96 errStr = err.Error() 97 } 98 99 l := utils.GetQueryLogger().With( 100 "error", errStr, 101 "request", aqlRequest, 102 "queries_enabled_", aqlRequest.Body.Queries, 103 "duration", duration, 104 "statusCode", statusCode, 105 "contexts_enabled_", qcs, 106 "headers", r.Header, 107 ) 108 109 if statusCode == http.StatusOK { 110 l.Info("All queries succeeded") 111 } else { 112 l.Error("Some of the queries finished with error") 113 } 114 115 }() 116 117 if aqlRequest.Query != "" { 118 // Override from query parameter 119 err = json.Unmarshal([]byte(aqlRequest.Query), &aqlRequest.Body) 120 if err != nil { 121 statusCode = http.StatusBadRequest 122 RespondWithBadRequest(w, utils.APIError{ 123 Code: http.StatusBadRequest, 124 Message: ErrMsgFailedToUnmarshalRequest, 125 Cause: err, 126 }) 127 return 128 } 129 } 130 131 if aqlRequest.Body.Queries == nil { 132 statusCode = http.StatusBadRequest 133 RespondWithBadRequest(w, utils.APIError{ 134 Code: http.StatusBadRequest, 135 Message: ErrMsgMissingParameter, 136 }) 137 return 138 } 139 140 returnHLL := aqlRequest.Accept == ContentTypeHyperLogLog 141 requestResponseWriter := getReponseWriter(returnHLL, len(aqlRequest.Body.Queries)) 142 143 queryTimer := utils.GetRootReporter().GetTimer(utils.QueryLatency) 144 start := utils.Now() 145 var qc *query.AQLQueryContext 146 for i, aqlQuery := range aqlRequest.Body.Queries { 147 qc, statusCode = handleQuery(handler.memStore, handler.deviceManager, aqlRequest, aqlQuery) 148 if aqlRequest.Verbose > 0 { 149 requestResponseWriter.ReportQueryContext(qc) 150 } 151 if qc.Error != nil { 152 requestResponseWriter.ReportError(i, aqlQuery.Table, qc.Error, statusCode) 153 } else { 154 requestResponseWriter.ReportResult(i, qc) 155 qc.ReleaseHostResultsBuffers() 156 utils.GetRootReporter().GetChildCounter(map[string]string{ 157 "table": aqlQuery.Table, 158 }, utils.QuerySucceeded).Inc(1) 159 } 160 161 qcs = append(qcs, qc) 162 } 163 duration = utils.Now().Sub(start) 164 queryTimer.Record(duration) 165 requestResponseWriter.Respond(w) 166 statusCode = requestResponseWriter.GetStatusCode() 167 return 168 } 169 170 func handleQuery(memStore memstore.MemStore, deviceManager *query.DeviceManager, aqlRequest AQLRequest, aqlQuery query.AQLQuery) (qc *query.AQLQueryContext, statusCode int) { 171 qc = aqlQuery.Compile(memStore, aqlRequest.Accept == ContentTypeHyperLogLog) 172 173 for tableName := range qc.TableSchemaByName { 174 utils.GetRootReporter().GetChildCounter(map[string]string{ 175 "table": tableName, 176 }, utils.QueryReceived).Inc(1) 177 } 178 179 if aqlRequest.Debug > 0 || aqlRequest.Profiling != "" { 180 qc.Debug = true 181 } 182 qc.Profiling = aqlRequest.Profiling 183 184 // Compilation error, should be bad request 185 if qc.Error != nil { 186 statusCode = http.StatusBadRequest 187 return 188 } 189 190 deviceChoosingTimeout := -1 191 if aqlRequest.DeviceChoosingTimeout > 0 { 192 deviceChoosingTimeout = aqlRequest.DeviceChoosingTimeout 193 } 194 // Find a device that meets the resource requirement of this query 195 // Use query specified device as hint 196 qc.FindDeviceForQuery(memStore, aqlRequest.Device, deviceManager, int(deviceChoosingTimeout)) 197 // Unable to find a device for the query. 198 if qc.Error != nil { 199 // Unable to fulfill this request due to resource not available, clients need to try sometimes later. 200 statusCode = http.StatusServiceUnavailable 201 return 202 } 203 defer deviceManager.ReleaseReservedMemory(qc.Device, qc.Query) 204 // Execute. 205 qc.ProcessQuery(memStore) 206 if qc.Error != nil { 207 utils.GetQueryLogger().With( 208 "error", qc.Error, 209 "query", aqlQuery, 210 "context", qc, 211 ).Error("Error happened when processing query") 212 statusCode = http.StatusInternalServerError 213 } else { 214 // Report 215 utils.GetRootReporter().GetChildCounter(map[string]string{ 216 "table": aqlQuery.Table, 217 }, utils.QueryRowsReturned).Inc(int64(qc.OOPK.ResultSize)) 218 } 219 return 220 } 221 222 func getReponseWriter(returnHLL bool, nQueries int) QueryResponseWriter { 223 if returnHLL { 224 return NewHLLQueryResponseWriter() 225 } 226 return NewJSONQueryResponseWriter(nQueries) 227 } 228 229 // QueryResponseWriter defines the interface to write query result and error to final response. 230 type QueryResponseWriter interface { 231 ReportError(queryIndex int, table string, err error, statusCode int) 232 ReportQueryContext(*query.AQLQueryContext) 233 ReportResult(int, *query.AQLQueryContext) 234 Respond(w http.ResponseWriter) 235 GetStatusCode() int 236 } 237 238 // JSONQueryResponseWriter writes query result as json. 239 type JSONQueryResponseWriter struct { 240 response query.AQLResponse 241 statusCode int 242 } 243 244 // NewJSONQueryResponseWriter creates a new JSONQueryResponseWriter. 245 func NewJSONQueryResponseWriter(nQueries int) QueryResponseWriter { 246 return &JSONQueryResponseWriter{ 247 response: query.AQLResponse{ 248 Results: make([]queryCom.AQLQueryResult, nQueries), 249 }, 250 statusCode: http.StatusOK, 251 } 252 } 253 254 // ReportError writes the error of the query to the response. 255 func (w *JSONQueryResponseWriter) ReportError(queryIndex int, table string, err error, statusCode int) { 256 // Usually larger status code means more severe problem. 257 if statusCode > w.statusCode { 258 w.statusCode = statusCode 259 } 260 if w.response.Errors == nil { 261 w.response.Errors = make([]error, len(w.response.Results)) 262 } 263 w.response.Errors[queryIndex] = err 264 utils.GetRootReporter().GetChildCounter(map[string]string{ 265 "table": table, 266 }, utils.QueryFailed).Inc(1) 267 } 268 269 // ReportQueryContext writes the query context to the response. 270 func (w *JSONQueryResponseWriter) ReportQueryContext(qc *query.AQLQueryContext) { 271 w.response.QueryContext = append(w.response.QueryContext, qc) 272 } 273 274 // ReportResult writes the query result to the response. 275 func (w *JSONQueryResponseWriter) ReportResult(queryIndex int, qc *query.AQLQueryContext) { 276 qc.Results = qc.Postprocess() 277 if qc.Error != nil { 278 w.ReportError(queryIndex, qc.Query.Table, qc.Error, http.StatusInternalServerError) 279 } 280 w.response.Results[queryIndex] = qc.Results 281 } 282 283 // Respond writes the final response into ResponseWriter. 284 func (w *JSONQueryResponseWriter) Respond(rw http.ResponseWriter) { 285 RespondJSONObjectWithCode(rw, w.statusCode, w.response) 286 } 287 288 // GetStatusCode returns the status code written into response. 289 func (w *JSONQueryResponseWriter) GetStatusCode() int { 290 return w.statusCode 291 } 292 293 // HLLQueryResponseWriter writes query result as application/hll. For more inforamtion, please refer to 294 // https://github.com/uber/aresdb/wiki/HyperLogLog. 295 type HLLQueryResponseWriter struct { 296 response *query.HLLQueryResults 297 statusCode int 298 } 299 300 // NewHLLQueryResponseWriter creates a new HLLQueryResponseWriter. 301 func NewHLLQueryResponseWriter() QueryResponseWriter { 302 w := HLLQueryResponseWriter{ 303 response: query.NewHLLQueryResults(), 304 statusCode: http.StatusOK, 305 } 306 return &w 307 } 308 309 // ReportError writes the error of the query to the response. 310 func (w *HLLQueryResponseWriter) ReportError(queryIndex int, table string, err error, statusCode int) { 311 if statusCode > w.statusCode { 312 w.statusCode = statusCode 313 } 314 w.response.WriteError(err) 315 } 316 317 // ReportQueryContext writes the query context to the response. Since the format of application/hll is not 318 // designed for human reading, we will ignore storing query context in response for now. 319 func (w *HLLQueryResponseWriter) ReportQueryContext(qc *query.AQLQueryContext) { 320 } 321 322 // ReportResult writes the query result to the response. 323 func (w *HLLQueryResponseWriter) ReportResult(queryIndex int, qc *query.AQLQueryContext) { 324 w.response.WriteResult(qc.HLLQueryResult) 325 } 326 327 // Respond writes the final response into ResponseWriter. 328 func (w *HLLQueryResponseWriter) Respond(rw http.ResponseWriter) { 329 rw.Header().Set("Content-Type", ContentTypeHyperLogLog) 330 RespondBytesWithCode(rw, w.statusCode, w.response.GetBytes()) 331 } 332 333 // GetStatusCode returns the status code written into response. 334 func (w *HLLQueryResponseWriter) GetStatusCode() int { 335 return w.statusCode 336 }