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  }