Added session, statement, emrjob metrics to sql stats api

Signed-off-by: Vamsi Manohar <[email protected]>

datasources/build.gradle

@@ -17,6 +17,7 @@ dependencies {
     implementation project(':core')
     implementation project(':protocol')
     implementation project(':opensearch')
+    implementation project(':legacy')
     implementation group: 'org.opensearch', name: 'opensearch', version: "${opensearch_version}"
     implementation group: 'org.opensearch', name: 'opensearch-x-content', version: "${opensearch_version}"
     implementation group: 'org.opensearch', name: 'common-utils', version: "${opensearch_build}"

datasources/src/main/java/org/opensearch/sql/datasources/rest/RestDataSourceQueryAction.java

@@ -34,6 +34,8 @@
 import org.opensearch.sql.datasources.transport.*;
 import org.opensearch.sql.datasources.utils.Scheduler;
 import org.opensearch.sql.datasources.utils.XContentParserUtils;
+import org.opensearch.sql.legacy.metrics.MetricName;
+import org.opensearch.sql.legacy.utils.MetricUtils;
 
 public class RestDataSourceQueryAction extends BaseRestHandler {
 
@@ -133,7 +135,6 @@ protected RestChannelConsumer prepareRequest(RestRequest restRequest, NodeClient
 
   private RestChannelConsumer executePostRequest(RestRequest restRequest, NodeClient nodeClient)
       throws IOException {
-
     DataSourceMetadata dataSourceMetadata =
         XContentParserUtils.toDataSourceMetadata(restRequest.contentParser());
     return restChannel ->
@@ -282,8 +283,10 @@ private void handleException(Exception e, RestChannel restChannel) {
     } else {
       LOG.error("Error happened during request handling", e);
       if (isClientError(e)) {
+        MetricUtils.incrementNumericalMetric(MetricName.DATASOURCE_FAILED_REQ_COUNT_CUS);
         reportError(restChannel, e, BAD_REQUEST);
       } else {
+        MetricUtils.incrementNumericalMetric(MetricName.DATASOURCE_FAILED_REQ_COUNT_SYS);
         reportError(restChannel, e, SERVICE_UNAVAILABLE);
       }
     }

datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceAction.java

@@ -20,6 +20,8 @@
 import org.opensearch.sql.datasources.model.transport.CreateDataSourceActionRequest;
 import org.opensearch.sql.datasources.model.transport.CreateDataSourceActionResponse;
 import org.opensearch.sql.datasources.service.DataSourceServiceImpl;
+import org.opensearch.sql.legacy.metrics.MetricName;
+import org.opensearch.sql.legacy.utils.MetricUtils;
 import org.opensearch.sql.protocol.response.format.JsonResponseFormatter;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
@@ -60,6 +62,7 @@ protected void doExecute(
       Task task,
       CreateDataSourceActionRequest request,
       ActionListener<CreateDataSourceActionResponse> actionListener) {
+    MetricUtils.incrementNumericalMetric(MetricName.DATASOURCE_CREATION_REQ_COUNT);
     int dataSourceLimit = settings.getSettingValue(DATASOURCES_LIMIT);
     if (dataSourceService.getDataSourceMetadata(false).size() >= dataSourceLimit) {
       actionListener.onFailure(

datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportDeleteDataSourceAction.java

@@ -16,6 +16,8 @@
 import org.opensearch.sql.datasources.model.transport.DeleteDataSourceActionRequest;
 import org.opensearch.sql.datasources.model.transport.DeleteDataSourceActionResponse;
 import org.opensearch.sql.datasources.service.DataSourceServiceImpl;
+import org.opensearch.sql.legacy.metrics.MetricName;
+import org.opensearch.sql.legacy.utils.MetricUtils;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
 
@@ -53,6 +55,7 @@ protected void doExecute(
       Task task,
       DeleteDataSourceActionRequest request,
       ActionListener<DeleteDataSourceActionResponse> actionListener) {
+    MetricUtils.incrementNumericalMetric(MetricName.DATASOURCE_DELETE_REQ_COUNT);
     try {
       dataSourceService.deleteDataSource(request.getDataSourceName());
       actionListener.onResponse(

datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportGetDataSourceAction.java

@@ -18,6 +18,8 @@
 import org.opensearch.sql.datasources.model.transport.GetDataSourceActionRequest;
 import org.opensearch.sql.datasources.model.transport.GetDataSourceActionResponse;
 import org.opensearch.sql.datasources.service.DataSourceServiceImpl;
+import org.opensearch.sql.legacy.metrics.MetricName;
+import org.opensearch.sql.legacy.utils.MetricUtils;
 import org.opensearch.sql.protocol.response.format.JsonResponseFormatter;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
@@ -56,6 +58,7 @@ protected void doExecute(
       Task task,
       GetDataSourceActionRequest request,
       ActionListener<GetDataSourceActionResponse> actionListener) {
+    MetricUtils.incrementNumericalMetric(MetricName.DATASOURCE_GET_REQ_COUNT);
     try {
       String responseContent;
       if (request.getDataSourceName() == null) {

datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportPatchDataSourceAction.java

@@ -19,6 +19,8 @@
 import org.opensearch.sql.datasources.model.transport.PatchDataSourceActionRequest;
 import org.opensearch.sql.datasources.model.transport.PatchDataSourceActionResponse;
 import org.opensearch.sql.datasources.service.DataSourceServiceImpl;
+import org.opensearch.sql.legacy.metrics.MetricName;
+import org.opensearch.sql.legacy.utils.MetricUtils;
 import org.opensearch.sql.protocol.response.format.JsonResponseFormatter;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
@@ -57,6 +59,7 @@ protected void doExecute(
       Task task,
       PatchDataSourceActionRequest request,
       ActionListener<PatchDataSourceActionResponse> actionListener) {
+    MetricUtils.incrementNumericalMetric(MetricName.DATASOURCE_PATCH_REQ_COUNT);
     try {
       dataSourceService.patchDataSource(request.getDataSourceData());
       String responseContent =

datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportUpdateDataSourceAction.java

@@ -18,6 +18,8 @@
 import org.opensearch.sql.datasources.model.transport.UpdateDataSourceActionRequest;
 import org.opensearch.sql.datasources.model.transport.UpdateDataSourceActionResponse;
 import org.opensearch.sql.datasources.service.DataSourceServiceImpl;
+import org.opensearch.sql.legacy.metrics.MetricName;
+import org.opensearch.sql.legacy.utils.MetricUtils;
 import org.opensearch.sql.protocol.response.format.JsonResponseFormatter;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
@@ -56,6 +58,7 @@ protected void doExecute(
       Task task,
       UpdateDataSourceActionRequest request,
       ActionListener<UpdateDataSourceActionResponse> actionListener) {
+    MetricUtils.incrementNumericalMetric(MetricName.DATASOURCE_PUT_REQ_COUNT);
     try {
       dataSourceService.updateDataSource(request.getDataSourceMetadata());
       String responseContent =

datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceActionTest.java

@@ -1,7 +1,9 @@
 package org.opensearch.sql.datasources.transport;
 
+import static java.util.Collections.emptyList;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.fail;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -21,11 +23,14 @@
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
 import org.opensearch.sql.datasource.model.DataSourceType;
 import org.opensearch.sql.datasources.model.transport.CreateDataSourceActionRequest;
 import org.opensearch.sql.datasources.model.transport.CreateDataSourceActionResponse;
 import org.opensearch.sql.datasources.service.DataSourceServiceImpl;
+import org.opensearch.sql.legacy.esdomain.LocalClusterState;
+import org.opensearch.sql.legacy.metrics.Metrics;
 import org.opensearch.sql.opensearch.setting.OpenSearchSettings;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
@@ -56,6 +61,12 @@ public void setUp() {
             transportService, new ActionFilters(new HashSet<>()), dataSourceService, settings);
     when(dataSourceService.getDataSourceMetadata(false).size()).thenReturn(1);
     when(settings.getSettingValue(DATASOURCES_LIMIT)).thenReturn(20);
+    // Required for metrics initialization
+    doReturn(emptyList()).when(settings).getSettings();
+    when(settings.getSettingValue(Settings.Key.METRICS_ROLLING_INTERVAL)).thenReturn(3600L);
+    when(settings.getSettingValue(Settings.Key.METRICS_ROLLING_WINDOW)).thenReturn(600L);
+    LocalClusterState.state().setPluginSettings(settings);
+    Metrics.getInstance().registerDefaultMetrics();
   }
 
   @Test

datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportDeleteDataSourceActionTest.java

@@ -1,8 +1,11 @@
 package org.opensearch.sql.datasources.transport;
 
+import static java.util.Collections.emptyList;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import java.util.HashSet;
 import org.junit.jupiter.api.Assertions;
@@ -16,9 +19,13 @@
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.datasources.model.transport.DeleteDataSourceActionRequest;
 import org.opensearch.sql.datasources.model.transport.DeleteDataSourceActionResponse;
 import org.opensearch.sql.datasources.service.DataSourceServiceImpl;
+import org.opensearch.sql.legacy.esdomain.LocalClusterState;
+import org.opensearch.sql.legacy.metrics.Metrics;
+import org.opensearch.sql.opensearch.setting.OpenSearchSettings;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
 
@@ -31,6 +38,8 @@ public class TransportDeleteDataSourceActionTest {
   @Mock private Task task;
   @Mock private ActionListener<DeleteDataSourceActionResponse> actionListener;
 
+  @Mock private OpenSearchSettings settings;
+
   @Captor
   private ArgumentCaptor<DeleteDataSourceActionResponse>
       deleteDataSourceActionResponseArgumentCaptor;
@@ -42,6 +51,12 @@ public void setUp() {
     action =
         new TransportDeleteDataSourceAction(
             transportService, new ActionFilters(new HashSet<>()), dataSourceService);
+    // Required for metrics initialization
+    doReturn(emptyList()).when(settings).getSettings();
+    when(settings.getSettingValue(Settings.Key.METRICS_ROLLING_INTERVAL)).thenReturn(3600L);
+    when(settings.getSettingValue(Settings.Key.METRICS_ROLLING_WINDOW)).thenReturn(600L);
+    LocalClusterState.state().setPluginSettings(settings);
+    Metrics.getInstance().registerDefaultMetrics();
   }
 
   @Test

datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportGetDataSourceActionTest.java

@@ -1,5 +1,7 @@
 package org.opensearch.sql.datasources.transport;
 
+import static java.util.Collections.emptyList;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -22,11 +24,16 @@
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
 import org.opensearch.sql.datasource.model.DataSourceType;
 import org.opensearch.sql.datasources.model.transport.GetDataSourceActionRequest;
 import org.opensearch.sql.datasources.model.transport.GetDataSourceActionResponse;
 import org.opensearch.sql.datasources.service.DataSourceServiceImpl;
+import org.opensearch.sql.legacy.esdomain.LocalClusterState;
+import org.opensearch.sql.legacy.metrics.MetricName;
+import org.opensearch.sql.legacy.metrics.Metrics;
+import org.opensearch.sql.opensearch.setting.OpenSearchSettings;
 import org.opensearch.sql.protocol.response.format.JsonResponseFormatter;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
@@ -45,11 +52,19 @@ public class TransportGetDataSourceActionTest {
 
   @Captor private ArgumentCaptor<Exception> exceptionArgumentCaptor;
 
+  @Mock private OpenSearchSettings settings;
+
   @BeforeEach
   public void setUp() {
     action =
         new TransportGetDataSourceAction(
             transportService, new ActionFilters(new HashSet<>()), dataSourceService);
+    // Required for metrics initialization
+    doReturn(emptyList()).when(settings).getSettings();
+    when(settings.getSettingValue(Settings.Key.METRICS_ROLLING_INTERVAL)).thenReturn(3600L);
+    when(settings.getSettingValue(Settings.Key.METRICS_ROLLING_WINDOW)).thenReturn(600L);
+    LocalClusterState.state().setPluginSettings(settings);
+    Metrics.getInstance().registerDefaultMetrics();
   }
 
   @Test
@@ -79,6 +94,9 @@ protected Object buildJsonObject(DataSourceMetadata response) {
         new Gson().fromJson(getDataSourceActionResponse.getResult(), DataSourceMetadata.class);
     Assertions.assertEquals("test_datasource", result.getName());
     Assertions.assertEquals(DataSourceType.PROMETHEUS, result.getConnector());
+    Assertions.assertEquals(
+        1L,
+        Metrics.getInstance().getNumericalMetric(MetricName.DATASOURCE_GET_REQ_COUNT).getValue());
   }
 
   @Test

datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportPatchDataSourceActionTest.java

@@ -1,5 +1,6 @@
 package org.opensearch.sql.datasources.transport;
 
+import static java.util.Collections.emptyList;
 import static org.mockito.Mockito.*;
 import static org.opensearch.sql.datasources.utils.XContentParserUtils.*;
 
@@ -17,9 +18,13 @@
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.datasources.model.transport.PatchDataSourceActionRequest;
 import org.opensearch.sql.datasources.model.transport.PatchDataSourceActionResponse;
 import org.opensearch.sql.datasources.service.DataSourceServiceImpl;
+import org.opensearch.sql.legacy.esdomain.LocalClusterState;
+import org.opensearch.sql.legacy.metrics.Metrics;
+import org.opensearch.sql.opensearch.setting.OpenSearchSettings;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
 
@@ -36,12 +41,19 @@ public class TransportPatchDataSourceActionTest {
   private ArgumentCaptor<PatchDataSourceActionResponse> patchDataSourceActionResponseArgumentCaptor;
 
   @Captor private ArgumentCaptor<Exception> exceptionArgumentCaptor;
+  @Mock private OpenSearchSettings settings;
 
   @BeforeEach
   public void setUp() {
     action =
         new TransportPatchDataSourceAction(
             transportService, new ActionFilters(new HashSet<>()), dataSourceService);
+    // Required for metrics initialization
+    doReturn(emptyList()).when(settings).getSettings();
+    when(settings.getSettingValue(Settings.Key.METRICS_ROLLING_INTERVAL)).thenReturn(3600L);
+    when(settings.getSettingValue(Settings.Key.METRICS_ROLLING_WINDOW)).thenReturn(600L);
+    LocalClusterState.state().setPluginSettings(settings);
+    Metrics.getInstance().registerDefaultMetrics();
   }
 
   @Test

datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportUpdateDataSourceActionTest.java

@@ -1,8 +1,11 @@
 package org.opensearch.sql.datasources.transport;
 
+import static java.util.Collections.emptyList;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import java.util.HashSet;
 import org.junit.jupiter.api.Assertions;
@@ -16,11 +19,15 @@
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
 import org.opensearch.sql.datasource.model.DataSourceType;
 import org.opensearch.sql.datasources.model.transport.UpdateDataSourceActionRequest;
 import org.opensearch.sql.datasources.model.transport.UpdateDataSourceActionResponse;
 import org.opensearch.sql.datasources.service.DataSourceServiceImpl;
+import org.opensearch.sql.legacy.esdomain.LocalClusterState;
+import org.opensearch.sql.legacy.metrics.Metrics;
+import org.opensearch.sql.opensearch.setting.OpenSearchSettings;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
 
@@ -32,6 +39,7 @@ public class TransportUpdateDataSourceActionTest {
   @Mock private DataSourceServiceImpl dataSourceService;
   @Mock private Task task;
   @Mock private ActionListener<UpdateDataSourceActionResponse> actionListener;
+  @Mock private OpenSearchSettings settings;
 
   @Captor
   private ArgumentCaptor<UpdateDataSourceActionResponse>
@@ -44,6 +52,12 @@ public void setUp() {
     action =
         new TransportUpdateDataSourceAction(
             transportService, new ActionFilters(new HashSet<>()), dataSourceService);
+    // Required for metrics initialization
+    doReturn(emptyList()).when(settings).getSettings();
+    when(settings.getSettingValue(Settings.Key.METRICS_ROLLING_INTERVAL)).thenReturn(3600L);
+    when(settings.getSettingValue(Settings.Key.METRICS_ROLLING_WINDOW)).thenReturn(600L);
+    LocalClusterState.state().setPluginSettings(settings);
+    Metrics.getInstance().registerDefaultMetrics();
   }
 
   @Test

legacy/src/main/java/org/opensearch/sql/legacy/metrics/MetricFactory.java

@@ -27,6 +27,19 @@ public static Metric createMetric(MetricName name) {
       case PPL_REQ_COUNT_TOTAL:
       case PPL_FAILED_REQ_COUNT_CUS:
       case PPL_FAILED_REQ_COUNT_SYS:
+      case DATASOURCE_CREATION_REQ_COUNT:
+      case DATASOURCE_GET_REQ_COUNT:
+      case DATASOURCE_PUT_REQ_COUNT:
+      case DATASOURCE_PATCH_REQ_COUNT:
+      case DATASOURCE_DELETE_REQ_COUNT:
+      case DATASOURCE_FAILED_REQ_COUNT_SYS:
+      case DATASOURCE_FAILED_REQ_COUNT_CUS:
+      case EMR_GET_JOB_RESULT_FAILURE_COUNT:
+      case EMR_START_JOB_REQUEST_FAILURE_COUNT:
+      case EMR_CANCEL_JOB_REQUEST_FAILURE_COUNT:
+      case EMR_BATCH_QUERY_JOBS_CREATION_COUNT:
+      case EMR_STREAMING_QUERY_JOBS_CREATION_COUNT:
+      case EMR_INTERACTIVE_QUERY_JOBS_CREATION_COUNT:
         return new NumericMetric<>(name.getName(), new RollingCounter());
       default:
         return new NumericMetric<>(name.getName(), new BasicCounter());