Add data source stats by type

Signed-off-by: Louis Chu <[email protected]>

datasources/src/main/java/org/opensearch/sql/datasources/service/DataSourceMetadataStorage.java

@@ -8,6 +8,7 @@
 package org.opensearch.sql.datasources.service;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import org.opensearch.sql.datasource.model.DataSource;
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
@@ -26,6 +27,8 @@ public interface DataSourceMetadataStorage {
    */
   List<DataSourceMetadata> getDataSourceMetadata();
 
+  Map<String, Long> countDataSourcesPerConnector();
+
   /**
    * Gets {@link DataSourceMetadata} corresponding to the datasourceName from underlying storage.
    *

datasources/src/main/java/org/opensearch/sql/datasources/service/DataSourceServiceImpl.java

@@ -17,6 +17,8 @@
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
 import org.opensearch.sql.datasources.auth.DataSourceUserAuthorizationHelper;
 import org.opensearch.sql.datasources.exceptions.DataSourceNotFoundException;
+import org.opensearch.sql.legacy.metrics.GaugeMetric;
+import org.opensearch.sql.legacy.metrics.Metrics;
 import org.opensearch.sql.storage.DataSourceFactory;
 
 /**
@@ -47,6 +49,7 @@ public DataSourceServiceImpl(
     this.dataSourceMetadataStorage = dataSourceMetadataStorage;
     this.dataSourceUserAuthorizationHelper = dataSourceUserAuthorizationHelper;
     this.dataSourceLoaderCache = new DataSourceLoaderCacheImpl(dataSourceFactories);
+    registerDataSourceTypeMetrics();
   }
 
   @Override
@@ -207,4 +210,12 @@ private void removeAuthInfo(DataSourceMetadata dataSourceMetadata) {
                     .anyMatch(confidentialKey -> entry.getKey().endsWith(confidentialKey)));
     dataSourceMetadata.setProperties(safeProperties);
   }
+
+  private void registerDataSourceTypeMetrics() {
+    GaugeMetric<Map<String, Long>> activeDataSourcesMetric =
+        new GaugeMetric<>(
+            "active_data_sources_count",
+            this.dataSourceMetadataStorage::countDataSourcesPerConnector);
+    Metrics.getInstance().registerMetric(activeDataSourcesMetric);
+  }
 }

datasources/src/main/java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorage.java

@@ -12,6 +12,7 @@
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -41,6 +42,8 @@
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.search.SearchHit;
+import org.opensearch.search.aggregations.AggregationBuilders;
+import org.opensearch.search.aggregations.bucket.terms.Terms;
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
 import org.opensearch.sql.datasources.encryptor.Encryptor;
@@ -86,6 +89,39 @@ public List<DataSourceMetadata> getDataSourceMetadata() {
     return searchInDataSourcesIndex(QueryBuilders.matchAllQuery());
   }
 
+  @Override
+  public Map<String, Long> countDataSourcesPerConnector() {
+    if (!this.clusterService.state().routingTable().hasIndex(DATASOURCE_INDEX_NAME)) {
+      createDataSourcesIndex();
+      return Collections.emptyMap();
+    }
+
+    SearchRequest searchRequest = new SearchRequest(DATASOURCE_INDEX_NAME);
+    SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
+    searchSourceBuilder.aggregation(
+        AggregationBuilders.terms("data_counts_by_connector_type").field("connector.keyword"));
+    searchRequest.source(searchSourceBuilder);
+    searchRequest.preference("_primary_first");
+    ActionFuture<SearchResponse> searchResponseActionFuture;
+    try (ThreadContext.StoredContext ignored =
+        client.threadPool().getThreadContext().stashContext()) {
+      searchResponseActionFuture = client.search(searchRequest);
+    }
+    SearchResponse searchResponse = searchResponseActionFuture.actionGet();
+    if (searchResponse.status().getStatus() != 200) {
+      throw new RuntimeException(
+          "Fetching dataSource metadata information failed with status : "
+              + searchResponse.status());
+    } else {
+      Terms connectorsAgg = searchResponse.getAggregations().get("data_counts_by_connector_type");
+      Map<String, Long> connectorCounts = new HashMap<>();
+      for (Terms.Bucket bucket : connectorsAgg.getBuckets()) {
+        connectorCounts.put(bucket.getKeyAsString(), bucket.getDocCount());
+      }
+      return connectorCounts;
+    }
+  }
+
   @Override
   public Optional<DataSourceMetadata> getDataSourceMetadata(String datasourceName) {
     if (!this.clusterService.state().routingTable().hasIndex(DATASOURCE_INDEX_NAME)) {

datasources/src/test/java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorageTest.java

@@ -5,10 +5,12 @@
 
 package org.opensearch.sql.datasources.storage;
 
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.opensearch.sql.datasources.storage.OpenSearchDataSourceMetadataStorage.DATASOURCE_INDEX_NAME;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -40,6 +42,8 @@
 import org.opensearch.index.engine.VersionConflictEngineException;
 import org.opensearch.search.SearchHit;
 import org.opensearch.search.SearchHits;
+import org.opensearch.search.aggregations.Aggregations;
+import org.opensearch.search.aggregations.bucket.terms.Terms;
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
 import org.opensearch.sql.datasource.model.DataSourceType;
 import org.opensearch.sql.datasources.encryptor.Encryptor;
@@ -247,6 +251,65 @@ public void testGetDataSourceMetadataListWithNoIndex() {
     Assertions.assertEquals(0, dataSourceMetadataList.size());
   }
 
+  @SneakyThrows
+  @Test
+  public void testCountDataSourcesPerConnector() {
+    Mockito.when(clusterService.state().routingTable().hasIndex(DATASOURCE_INDEX_NAME))
+        .thenReturn(true);
+    Mockito.when(client.search(ArgumentMatchers.any())).thenReturn(searchResponseActionFuture);
+    Mockito.when(searchResponseActionFuture.actionGet()).thenReturn(searchResponse);
+    Mockito.when(searchResponse.status()).thenReturn(RestStatus.OK);
+    // Mocking the aggregation response
+    Terms.Bucket bucketA = Mockito.mock(Terms.Bucket.class);
+    Mockito.when(bucketA.getKeyAsString()).thenReturn("DataSourceTypeA");
+    Mockito.when(bucketA.getDocCount()).thenReturn(1L);
+
+    Terms.Bucket bucketB = Mockito.mock(Terms.Bucket.class);
+    Mockito.when(bucketB.getKeyAsString()).thenReturn("DataSourceTypeB");
+    Mockito.when(bucketB.getDocCount()).thenReturn(2L);
+
+    // Use raw type here and cast appropriately
+    Terms terms = Mockito.mock(Terms.class);
+    List rawBucketsList = Arrays.asList(bucketA, bucketB); // Raw type
+    Mockito.when(terms.getBuckets()).thenReturn((List) rawBucketsList); // Casting to raw type
+
+    Aggregations aggregations = Mockito.mock(Aggregations.class);
+    Mockito.when(aggregations.get(anyString())).thenReturn(terms);
+
+    Mockito.when(searchResponse.getAggregations()).thenReturn(aggregations);
+
+    Map<String, Long> connectorCounts =
+        openSearchDataSourceMetadataStorage.countDataSourcesPerConnector();
+
+    Assertions.assertEquals(
+        2, connectorCounts.size(), "The size of connectorCounts map should be 2.");
+    Assertions.assertEquals(
+        Long.valueOf(1),
+        connectorCounts.get("DataSourceTypeA"),
+        "DataSourceTypeA should have a count of 1.");
+    Assertions.assertEquals(
+        Long.valueOf(2),
+        connectorCounts.get("DataSourceTypeB"),
+        "DataSourceTypeB should have a count of 2.");
+  }
+
+  @SneakyThrows
+  @Test
+  public void testCountDataSourcesPerConnectorWithNoIndex() {
+    Mockito.when(clusterService.state().routingTable().hasIndex(DATASOURCE_INDEX_NAME))
+        .thenReturn(Boolean.FALSE);
+    Mockito.when(client.admin().indices().create(ArgumentMatchers.any()))
+        .thenReturn(createIndexResponseActionFuture);
+    Mockito.when(createIndexResponseActionFuture.actionGet())
+        .thenReturn(new CreateIndexResponse(true, true, DATASOURCE_INDEX_NAME));
+    Mockito.when(client.index(ArgumentMatchers.any())).thenReturn(indexResponseActionFuture);
+
+    Map<String, Long> connectorCounts =
+        openSearchDataSourceMetadataStorage.countDataSourcesPerConnector();
+
+    Assertions.assertEquals(0, connectorCounts.size());
+  }
+
   @SneakyThrows
   @Test
   public void testGetDataSourceMetadataWithNoIndex() {