Added metrics for SQL query requests in new engine

integ-test/src/test/java/com/amazon/opendistroforelasticsearch/sql/sql/MetricsIT.java

@@ -0,0 +1,78 @@
+/*
+ *    Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *    Licensed under the Apache License, Version 2.0 (the "License").
+ *    You may not use this file except in compliance with the License.
+ *    A copy of the License is located at
+ *
+ *        http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    or in the "license" file accompanying this file. This file is distributed
+ *    on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *    express or implied. See the License for the specific language governing
+ *    permissions and limitations under the License.
+ *
+ */
+
+package com.amazon.opendistroforelasticsearch.sql.sql;
+
+import static com.amazon.opendistroforelasticsearch.sql.legacy.TestsConstants.TEST_INDEX_BANK;
+
+import com.amazon.opendistroforelasticsearch.sql.legacy.SQLIntegTestCase;
+import com.amazon.opendistroforelasticsearch.sql.legacy.metrics.MetricName;
+import com.amazon.opendistroforelasticsearch.sql.util.TestUtils;
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.util.Locale;
+import java.util.concurrent.TimeUnit;
+import org.elasticsearch.client.Request;
+import org.elasticsearch.client.Response;
+import org.json.JSONObject;
+import org.junit.Assert;
+import org.junit.Test;
+
+public class MetricsIT extends SQLIntegTestCase {
+
+  @Override
+  protected void init() throws Exception {
+    loadIndex(Index.BANK);
+    TestUtils.enableNewQueryEngine(client());
+  }
+
+  @Test
+  public void requestCount() throws IOException, InterruptedException {
+    int beforeQueries = requestTotal();
+    executeQuery(String.format(Locale.ROOT, "select age from %s", TEST_INDEX_BANK));
+    TimeUnit.SECONDS.sleep(2L);
+
+    assertEquals(beforeQueries + 1, requestTotal());
+  }
+
+  private Request makeStatRequest() {
+    return new Request(
+        "GET", "/_opendistro/_sql/stats"
+    );
+  }
+
+  private int requestTotal() throws IOException {
+    JSONObject jsonObject = new JSONObject(executeStatRequest(makeStatRequest()));
+    return jsonObject.getInt(MetricName.REQ_TOTAL.getName());
+  }
+
+  private String executeStatRequest(final Request request) throws IOException {
+    Response response = client().performRequest(request);
+    Assert.assertEquals(200, response.getStatusLine().getStatusCode());
+
+    InputStream is = response.getEntity().getContent();
+    StringBuilder sb = new StringBuilder();
+    try (BufferedReader br = new BufferedReader(new InputStreamReader(is))) {
+      String line;
+      while ((line = br.readLine()) != null) {
+        sb.append(line);
+      }
+    }
+    return sb.toString();
+  }
+}

legacy/src/main/java/com/amazon/opendistroforelasticsearch/sql/legacy/plugin/RestSQLQueryAction.java

@@ -18,14 +18,28 @@
 
 import static com.amazon.opendistroforelasticsearch.sql.executor.ExecutionEngine.QueryResponse;
 import static com.amazon.opendistroforelasticsearch.sql.protocol.response.format.JsonResponseFormatter.Style.PRETTY;
+import static org.elasticsearch.rest.RestStatus.BAD_REQUEST;
 import static org.elasticsearch.rest.RestStatus.INTERNAL_SERVER_ERROR;
 import static org.elasticsearch.rest.RestStatus.OK;
+import static org.elasticsearch.rest.RestStatus.SERVICE_UNAVAILABLE;
 
+
+import com.alibaba.druid.sql.parser.ParserException;
 import com.amazon.opendistroforelasticsearch.sql.common.antlr.SyntaxCheckException;
 import com.amazon.opendistroforelasticsearch.sql.common.response.ResponseListener;
 import com.amazon.opendistroforelasticsearch.sql.common.setting.Settings;
 import com.amazon.opendistroforelasticsearch.sql.elasticsearch.security.SecurityAccess;
+import com.amazon.opendistroforelasticsearch.sql.exception.QueryEngineException;
+import com.amazon.opendistroforelasticsearch.sql.exception.SemanticCheckException;
 import com.amazon.opendistroforelasticsearch.sql.executor.ExecutionEngine.ExplainResponse;
+import com.amazon.opendistroforelasticsearch.sql.legacy.antlr.SqlAnalysisException;
+import com.amazon.opendistroforelasticsearch.sql.legacy.exception.SQLFeatureDisabledException;
+import com.amazon.opendistroforelasticsearch.sql.legacy.exception.SqlParseException;
+import com.amazon.opendistroforelasticsearch.sql.legacy.executor.format.ErrorMessageFactory;
+import com.amazon.opendistroforelasticsearch.sql.legacy.metrics.MetricName;
+import com.amazon.opendistroforelasticsearch.sql.legacy.metrics.Metrics;
+import com.amazon.opendistroforelasticsearch.sql.legacy.rewriter.matchtoterm.VerificationException;
+import com.amazon.opendistroforelasticsearch.sql.legacy.utils.LogUtils;
 import com.amazon.opendistroforelasticsearch.sql.planner.physical.PhysicalPlan;
 import com.amazon.opendistroforelasticsearch.sql.protocol.response.QueryResult;
 import com.amazon.opendistroforelasticsearch.sql.protocol.response.format.JdbcResponseFormatter;
@@ -35,11 +49,13 @@
 import com.amazon.opendistroforelasticsearch.sql.sql.domain.SQLQueryRequest;
 import java.io.IOException;
 import java.security.PrivilegedExceptionAction;
+import java.sql.SQLFeatureNotSupportedException;
 import java.util.List;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.client.node.NodeClient;
 import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.index.IndexNotFoundException;
 import org.elasticsearch.rest.BaseRestHandler;
 import org.elasticsearch.rest.BytesRestResponse;
 import org.elasticsearch.rest.RestChannel;
@@ -143,6 +159,7 @@ protected Object buildJsonObject(ExplainResponse response) {
       @Override
       public void onFailure(Exception e) {
         LOG.error("Error happened during explain", e);
+        logAndPublishMetrics(e);
         sendResponse(channel, INTERNAL_SERVER_ERROR,
             "Failed to explain the query due to error: " + e.getMessage());
       }
@@ -161,6 +178,7 @@ public void onResponse(QueryResponse response) {
       @Override
       public void onFailure(Exception e) {
         LOG.error("Error happened during query handling", e);
+        logAndPublishMetrics(e);
         sendResponse(channel, INTERNAL_SERVER_ERROR, formatter.format(e));
       }
     };
@@ -179,4 +197,17 @@ private void sendResponse(RestChannel channel, RestStatus status, String content
         status, "application/json; charset=UTF-8", content));
   }
 
+  private static void logAndPublishMetrics(Exception e) {
+    if (isServerError(e)) {
+      LOG.error(LogUtils.getRequestId() + " Server side error during query execution", e);
+      Metrics.getInstance().getNumericalMetric(MetricName.FAILED_REQ_COUNT_SYS).increment();
+    } else {
+      LOG.error(LogUtils.getRequestId() + " Client side error during query execution", e);
+      Metrics.getInstance().getNumericalMetric(MetricName.FAILED_REQ_COUNT_CUS).increment();
+    }
+  }
+
+  private static boolean isServerError(Exception e) {
+    return e instanceof QueryEngineException;
+  }
 }