ATLAS-2578: Update metric queries for faster execution

intg/src/main/java/org/apache/atlas/model/metrics/AtlasMetrics.java

@@ -36,13 +36,13 @@ import static com.fasterxml.jackson.annotation.JsonAutoDetect.Visibility.PUBLIC_
 @JsonSerialize(include=JsonSerialize.Inclusion.NON_NULL)
 @JsonIgnoreProperties(ignoreUnknown=true)
 public class AtlasMetrics {
-    private Map<String, Map<String, Number>> data;
+    private Map<String, Map<String, Object>> data;
 
     public AtlasMetrics() {
         setData(null);
     }
 
-    public AtlasMetrics(Map<String, Map<String, Number>> data) {
+    public AtlasMetrics(Map<String, Map<String, Object>> data) {
         setData(data);
     }
 
@@ -52,41 +52,41 @@ public class AtlasMetrics {
         }
     }
 
-    public Map<String, Map<String, Number>> getData() {
+    public Map<String, Map<String, Object>> getData() {
         return data;
     }
 
-    public void setData(Map<String, Map<String, Number>> data) {
+    public void setData(Map<String, Map<String, Object>> data) {
         this.data = data;
     }
 
     @JsonIgnore
-    public void addData(String groupKey, String key, Number value) {
-        Map<String, Map<String, Number>> data = this.data;
+    public void addMetric(String groupKey, String key, Object value) {
+        Map<String, Map<String, Object>> data = this.data;
         if (data == null) {
             data = new HashMap<>();
         }
-        Map<String, Number> metricMap = data.get(groupKey);
-        if (metricMap == null) {
-            metricMap = new HashMap<>();
-            data.put(groupKey, metricMap);
-        }
+        Map<String, Object> metricMap = data.computeIfAbsent(groupKey, k -> new HashMap<>());
         metricMap.put(key, value);
         setData(data);
     }
 
     @JsonIgnore
-    public Number getMetric(String groupKey, String key) {
-        Map<String, Map<String, Number>> data = this.data;
-        if (data == null) {
-            return null;
-        } else {
-            Map<String, Number> metricMap = data.get(groupKey);
-            if (metricMap == null || metricMap.isEmpty()) {
-                return null;
-            } else {
-                return metricMap.get(key);
+    public Number getNumericMetric(String groupKey, String key) {
+        Object metric = getMetric(groupKey, key);
+        return metric instanceof Number ? (Number) metric : null;
+    }
+
+    @JsonIgnore
+    public Object getMetric(String groupKey, String key) {
+        Map<String, Map<String, Object>> data = this.data;
+        Object ret = null;
+        if (data != null) {
+            Map<String, Object> metricMap = data.get(groupKey);
+            if (metricMap != null && !metricMap.isEmpty()) {
+                ret = metricMap.get(key);
             }
         }
+        return ret;
     }
 }

repository/src/main/java/org/apache/atlas/util/AtlasGremlin2QueryProvider.java

@@ -21,22 +21,11 @@ public class AtlasGremlin2QueryProvider extends AtlasGremlinQueryProvider {
     @Override
     public String getQuery(final AtlasGremlinQuery gremlinQuery) {
         switch (gremlinQuery) {
-            case TYPE_COUNT_METRIC:
-                return "g.V().has('__type', 'typeSystem').filter({!it.'__type.category'.name().matches('TRAIT')}).count()";
-            case TYPE_UNUSED_COUNT_METRIC:
-                return "g.V('__type', 'typeSystem').filter({ !it.getProperty('__type.category').name().matches('TRAIT') && it.inE().count() == 0}).count()";
-            case ENTITY_COUNT_METRIC:
-                return "g.V().has('__superTypeNames', T.in, ['Referenceable']).count()";
-            case TAG_COUNT_METRIC:
-                return "g.V().has('__type', 'typeSystem').filter({it.getProperty('__type.category').name().matches('TRAIT')}).count()";
+            case ENTITY_ACTIVE_METRIC:
+                return "g.V().has('__typeName', T.in, [%s]).has('__state', 'ACTIVE').groupCount{it.getProperty('__typeName')}.cap.toList()";
             case ENTITY_DELETED_METRIC:
-                return "g.V().has('__typeName', T.in, g.V().has('__type', 'typeSystem').filter{it.getProperty('__type.category').name().matches('CLASS')}.'__type.name'.toSet()).has('__status', 'DELETED').count()";
-            case ENTITIES_PER_TYPE_METRIC:
-                return "g.V().has('__typeName', T.in, g.V().has('__type', 'typeSystem').filter{it.getProperty('__type.category').name() == 'CLASS'}.'__type.name'.toSet()).groupCount{it.getProperty('__typeName')}.cap.toList()";
-            case TAGGED_ENTITIES_METRIC:
-                return "g.V().has('__traitNames', T.in, g.V().has('__type', 'typeSystem').filter{it.getProperty('__type.category').name() == 'TRAIT'}.'__type.name'.toSet()).count()";
-            case ENTITIES_FOR_TAG_METRIC:
-                return "g.V().has('__typeName', T.in, g.V().has('__type', 'typeSystem').filter{it.getProperty('__type.category').name() == 'TRAIT'}.'__type.name'.toSet()).groupCount{it.getProperty('__typeName')}.cap.toList()";
+                return "g.V().has('__typeName', T.in, [%s]).has('__state', 'DELETED').groupCount{it.getProperty('__typeName')}.cap.toList()";
+
                 case EXPORT_BY_GUID_FULL:
                 return "g.V('__guid', startGuid).bothE().bothV().has('__guid').transform{[__guid:it.__guid,isProcess:(it.__superTypeNames != null) ? it.__superTypeNames.contains('Process') : false ]}.dedup().toList()";
             case EXPORT_BY_GUID_CONNECTED_IN_EDGE:

repository/src/main/java/org/apache/atlas/util/AtlasGremlin3QueryProvider.java

@@ -23,10 +23,10 @@ public class AtlasGremlin3QueryProvider extends AtlasGremlin2QueryProvider {
         // In case any overrides are necessary, a specific switch case can be added here to
         // return Gremlin 3 specific query otherwise delegate to super.getQuery
         switch (gremlinQuery) {
-            case TYPE_UNUSED_COUNT_METRIC:
-                return "g.V().has('__type', 'typeSystem').filter({ !it.getProperty('__type.category').name().matches('TRAIT') && it.inE().count() == 0}).count()";
-            case ENTITY_COUNT_METRIC:
-                return "g.V().has('__superTypeNames', within(['Referenceable'])).count()";
+            case ENTITY_ACTIVE_METRIC:
+                return "g.V().has('__typeName', within(%s)).has('__state', 'ACTIVE').groupCount().by('__typeName').toList()";
+            case ENTITY_DELETED_METRIC:
+                return "g.V().has('__typeName', within(%s)).has('__state', 'DELETED').groupCount().by('__typeName').toList()";
             case EXPORT_TYPE_STARTS_WITH:
                 return "g.V().has('__typeName',typeName).filter({it.get().value(attrName).startsWith(attrValue)}).has('__guid').values('__guid').toList()";
             case EXPORT_TYPE_ENDS_WITH:

repository/src/main/java/org/apache/atlas/util/AtlasGremlinQueryProvider.java

@@ -34,8 +34,8 @@ public abstract class AtlasGremlinQueryProvider {
         // Metrics related Queries
         TYPE_COUNT_METRIC,
         TYPE_UNUSED_COUNT_METRIC,
-        ENTITY_COUNT_METRIC,
         TAG_COUNT_METRIC,
+        ENTITY_ACTIVE_METRIC,
         ENTITY_DELETED_METRIC,
         ENTITIES_PER_TYPE_METRIC,
         TAGGED_ENTITIES_METRIC,

repository/src/test/java/org/apache/atlas/services/MetricsServiceTest.java

@@ -17,19 +17,16 @@
  */
 package org.apache.atlas.services;
 
-import org.apache.atlas.TestModules;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.metrics.AtlasMetrics;
 import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.runner.LocalSolrRunner;
 import org.apache.atlas.type.AtlasTypeRegistry;
+import org.apache.atlas.util.AtlasGremlin3QueryProvider;
 import org.apache.commons.configuration.Configuration;
-import org.mockito.invocation.InvocationOnMock;
-import org.mockito.stubbing.Answer;
 import org.testng.SkipException;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
-import org.testng.annotations.Guice;
 import org.testng.annotations.Test;
 
 import java.util.ArrayList;
@@ -45,34 +42,40 @@ import static org.mockito.Mockito.*;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
 
-@Guice(modules = TestModules.TestOnlyModule.class)
 public class MetricsServiceTest {
     private Configuration             mockConfig        = mock(Configuration.class);
     private AtlasTypeRegistry         mockTypeRegistry  = mock(AtlasTypeRegistry.class);
     private AtlasGraph                mockGraph         = mock(AtlasGraph.class);
     private MetricsService            metricsService;
 
-    private List<Map> mockMapList = new ArrayList<>();
-    private Number mockCount = 10;
+    private List<Map> activeEntityCountList = new ArrayList<>();
+    private List<Map> deletedEntityCountList = new ArrayList<>();
 
     @BeforeClass
     public void init() throws Exception {
         try {
-            Map<String, Object> mockMap = new HashMap<>();
-            mockMap.put("a", 1);
-            mockMap.put("b", 2);
-            mockMap.put("c", 3);
-            mockMapList.add(mockMap);
+            Map<String, Object> activeEntityCount = new HashMap<>();
+            activeEntityCount.put("a", 1);
+            activeEntityCount.put("b", 2);
+            activeEntityCount.put("d", 5);
+            activeEntityCount.put("e", 10);
+            activeEntityCount.put("f", 15);
+            activeEntityCountList.add(activeEntityCount);
+
+            Map<String, Object> deletedEntityCount = new HashMap<>();
+            deletedEntityCount.put("b", 5);
+            deletedEntityCountList.add(deletedEntityCount);
 
             when(mockConfig.getInt(anyString(), anyInt())).thenReturn(5);
             assertEquals(mockConfig.getInt("test", 1), 5);
-            when(mockConfig.getString(anyString(), anyString()))
-                    .thenReturn("count()", "count()", "count()", "count()", "count()", "toList()", "count()", "toList()");
-            when(mockTypeRegistry.getAllEntityDefNames()).thenReturn(Arrays.asList("a", "b", "c"));
+            when(mockConfig.getString(anyString(), anyString())).thenReturn("toList()", "toList()");
+            when(mockTypeRegistry.getAllTypeNames()).thenReturn(Arrays.asList("a", "b", "c", "d", "e", "f"));
             when(mockTypeRegistry.getAllEntityDefNames()).thenReturn(Arrays.asList("a", "b", "c"));
+            when(mockTypeRegistry.getAllClassificationDefNames()).thenReturn(Arrays.asList("d", "e", "f"));
+//            when(mockTypeRegistry.getAllEntityDefNames()).thenReturn(Arrays.asList("a", "b", "c"));
             setupMockGraph();
 
-            metricsService = new MetricsService(mockConfig, mockGraph);
+            metricsService = new MetricsService(mockConfig, mockGraph, mockTypeRegistry, new AtlasGremlin3QueryProvider());
         }
         catch(Exception e) {
             throw new SkipException("MetricsServicesTest: init failed!", e);
@@ -88,16 +91,9 @@ public class MetricsServiceTest {
 
     private void setupMockGraph() throws AtlasBaseException {
         if (mockGraph == null) mockGraph = mock(AtlasGraph.class);
-        when(mockGraph.executeGremlinScript(anyString(), eq(false))).thenAnswer(new Answer<Object>() {
-            @Override
-            public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
-                if (((String)invocationOnMock.getArguments()[0]).contains("count()")) {
-                    return mockCount;
-                } else {
-                    return mockMapList;
-                }
-            }
-        });
+        when(mockGraph.executeGremlinScript(anyString(), eq(false)))
+                .thenReturn(activeEntityCountList)
+                .thenReturn(deletedEntityCountList);
     }
 
     @Test
@@ -105,29 +101,28 @@ public class MetricsServiceTest {
         assertNotNull(metricsService);
         AtlasMetrics metrics = metricsService.getMetrics(false);
         assertNotNull(metrics);
-        Number aCount = metrics.getMetric("entity", "a");
-        assertNotNull(aCount);
-        assertEquals(aCount, 1);
+        Map activeMetrics = (Map) metrics.getMetric("entity", "entityActive");
+        assertNotNull(activeMetrics);
+        assertEquals(activeMetrics.get("a"), 1);
+        assertEquals(activeMetrics.get("b"), 2);
 
-        Number bCount = metrics.getMetric("entity", "b");
-        assertNotNull(bCount);
-        assertEquals(bCount, 2);
+        Map deletedMetrics = (Map) metrics.getMetric("entity", "entityDeleted");
+        assertEquals(deletedMetrics.get("b"), 5);
 
-        Number cCount = metrics.getMetric("entity", "c");
-        assertNotNull(cCount);
-        assertEquals(cCount, 3);
+        Number unusedTypeCount = metrics.getNumericMetric("general", "typeUnusedCount");
+        assertEquals(unusedTypeCount, 1);
 
-        Number aTags = metrics.getMetric("tag", "a");
-        assertNotNull(aTags);
-        assertEquals(aTags, 1);
+        Number cCount = metrics.getNumericMetric("general", "entityCount");
+        assertEquals(cCount, 8);
 
-        Number bTags = metrics.getMetric("tag", "b");
-        assertNotNull(bTags);
-        assertEquals(bTags, 2);
+        Number aTags = (Number) metrics.getMetric("general", "tagCount");
+        assertEquals(aTags, 3);
 
-        Number cTags = metrics.getMetric("tag", "c");
-        assertNotNull(cTags);
-        assertEquals(cTags, 3);
+        Map taggedEntityMetric = (Map) metrics.getMetric("tag", "tagEntities");
+        assertNotNull(taggedEntityMetric);
+        assertEquals(taggedEntityMetric.get("d"), 5);
+        assertEquals(taggedEntityMetric.get("e"), 10);
+        assertEquals(taggedEntityMetric.get("f"), 15);
 
         verify(mockGraph, atLeastOnce()).executeGremlinScript(anyString(), anyBoolean());
 

webapp/src/main/java/org/apache/atlas/web/resources/AdminResource.java

@@ -99,9 +99,9 @@ public class AdminResource {
 
     private Response version;
 
+    private static Configuration            atlasProperties;
     private final  ServiceState             serviceState;
     private final  MetricsService           metricsService;
-    private static Configuration           atlasProperties;
     private final  ExportService            exportService;
     private final  ImportService            importService;
     private final  SearchTracker            activeSearches;