ATLAS-1436: Metrics collection using gremlin

client/src/main/java/org/apache/atlas/AtlasAdminClient.java

@@ -18,6 +18,8 @@
 
 package org.apache.atlas;
 
+import org.apache.atlas.model.metrics.AtlasMetrics;
+import org.apache.atlas.type.AtlasType;
 import org.apache.atlas.utils.AuthenticationUtil;
 import org.apache.commons.cli.CommandLine;
 import org.apache.commons.cli.CommandLineParser;
@@ -44,6 +46,7 @@ import java.util.Arrays;
 public class AtlasAdminClient {
 
     private static final Option STATUS = new Option("status", false, "Get the status of an atlas instance");
+    private static final Option STATS = new Option("stats", false, "Get the metrics of an atlas instance");
     private static final Options OPTIONS = new Options();
 
     private static final int INVALID_OPTIONS_STATUS = 1;
@@ -51,6 +54,7 @@ public class AtlasAdminClient {
 
     static {
         OPTIONS.addOption(STATUS);
+        OPTIONS.addOption(STATS);
     }
 
     public static void main(String[] args) throws AtlasException, ParseException {
@@ -88,6 +92,16 @@ public class AtlasAdminClient {
                 System.err.println("Could not retrieve status of the server at " + Arrays.toString(atlasServerUri));
                 printStandardHttpErrorDetails(e);
             }
+        } else if (commandLine.hasOption(STATS.getOpt())) {
+            try {
+                AtlasMetrics atlasMetrics = atlasClient.getAtlasMetrics();
+                String json = AtlasType.toJson(atlasMetrics);
+                System.out.println(json);
+                cmdStatus = 0;
+            } catch (AtlasServiceException e) {
+                System.err.println("Could not retrieve metrics of the server at " + Arrays.toString(atlasServerUri));
+                printStandardHttpErrorDetails(e);
+            }
         } else {
             System.err.println("Unsupported option. Refer to usage for valid options.");
             printUsage(INVALID_OPTIONS_STATUS);

client/src/main/java/org/apache/atlas/AtlasBaseClient.java

@@ -26,6 +26,7 @@ import com.sun.jersey.api.client.config.DefaultClientConfig;
 import com.sun.jersey.api.client.filter.HTTPBasicAuthFilter;
 import com.sun.jersey.api.json.JSONConfiguration;
 import com.sun.jersey.client.urlconnection.URLConnectionClientHandler;
+import org.apache.atlas.model.metrics.AtlasMetrics;
 import org.apache.atlas.security.SecureClientUtils;
 import org.apache.atlas.type.AtlasType;
 import org.apache.atlas.utils.AuthenticationUtil;
@@ -54,10 +55,12 @@ public abstract class AtlasBaseClient {
     public static final String TYPES = "types";
     public static final String ADMIN_VERSION = "admin/version";
     public static final String ADMIN_STATUS = "admin/status";
+    public static final String ADMIN_METRICS = "admin/metrics";
     public static final String HTTP_AUTHENTICATION_ENABLED = "atlas.http.authentication.enabled";
     //Admin operations
     public static final APIInfo VERSION = new APIInfo(BASE_URI + ADMIN_VERSION, HttpMethod.GET, Response.Status.OK);
     public static final APIInfo STATUS = new APIInfo(BASE_URI + ADMIN_STATUS, HttpMethod.GET, Response.Status.OK);
+    public static final APIInfo METRICS = new APIInfo(BASE_URI + ADMIN_METRICS, HttpMethod.GET, Response.Status.OK);
     static final String JSON_MEDIA_TYPE = MediaType.APPLICATION_JSON + "; charset=UTF-8";
     static final String UNKNOWN_STATUS = "Unknown status";
     static final String ATLAS_CLIENT_HA_RETRIES_KEY = "atlas.client.ha.retries";
@@ -366,6 +369,14 @@ public abstract class AtlasBaseClient {
         return result;
     }
 
+    /**
+     * @return Return metrics of the service instance the client is pointing to
+     * @throws AtlasServiceException
+     */
+    public AtlasMetrics getAtlasMetrics () throws AtlasServiceException {
+        return callAPI(METRICS, AtlasMetrics.class, null);
+    }
+
     boolean isRetryableException(ClientHandlerException che) {
         return che.getCause().getClass().equals(IOException.class)
                 || che.getCause().getClass().equals(ConnectException.class);

distro/src/conf/atlas-application.properties

@@ -205,5 +205,15 @@ atlas.feature.taxonomy.enable=true
 #atlas.sso.knox.providerurl=https://<knox gateway ip>:8443/gateway/knoxsso/api/v1/websso
 #atlas.sso.knox.publicKey=
 
-
-
+############ Atlas Metric/Stats configs ################
+# Format: atlas.metric.query.<key>.<name>
+#atlas.metric.query.general.typeCount=
+#atlas.metric.query.general.typeUnusedCount=
+#atlas.metric.query.general.entityCount=
+#atlas.metric.query.general.tagCount=
+#atlas.metric.query.general.entityDeleted=
+#
+#atlas.metric.query.entity.typeEntities=
+#atlas.metric.query.entity.entityTagged=
+#
+#atlas.metric.query.tags.entityTags=

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

+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License 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 org.apache.atlas.model.metrics;
+
+
+import org.codehaus.jackson.annotate.JsonAutoDetect;
+import org.codehaus.jackson.annotate.JsonIgnore;
+import org.codehaus.jackson.annotate.JsonIgnoreProperties;
+import org.codehaus.jackson.map.annotate.JsonSerialize;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.NONE;
+import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.PUBLIC_ONLY;
+
+/**
+ * Atlas metrics
+ */
+@JsonAutoDetect(getterVisibility=PUBLIC_ONLY, setterVisibility=PUBLIC_ONLY, fieldVisibility=NONE)
+@JsonSerialize(include=JsonSerialize.Inclusion.NON_NULL)
+@JsonIgnoreProperties(ignoreUnknown=true)
+public class AtlasMetrics {
+    private Map<String, Map<String, Number>> data;
+
+    public AtlasMetrics() {
+        setData(null);
+    }
+
+    public AtlasMetrics(Map<String, Map<String, Number>> data) {
+        setData(data);
+    }
+
+    public AtlasMetrics(AtlasMetrics other) {
+        if (other != null) {
+            setData(other.getData());
+        }
+    }
+
+    public Map<String, Map<String, Number>> getData() {
+        return data;
+    }
+
+    public void setData(Map<String, Map<String, Number>> data) {
+        this.data = data;
+    }
+
+    @JsonIgnore
+    public void addData(String groupKey, String key, Integer value) {
+        Map<String, Map<String, Number>> 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);
+        }
+        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);
+            }
+        }
+    }
+}

repository/src/main/java/org/apache/atlas/services/MetricsService.java

+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License 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 org.apache.atlas.services;
+
+import com.google.inject.Singleton;
+import org.apache.atlas.ApplicationProperties;
+import org.apache.atlas.AtlasException;
+import org.apache.atlas.model.metrics.AtlasMetrics;
+import org.apache.atlas.repository.graph.AtlasGraphProvider;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
+import org.apache.commons.configuration.Configuration;
+import org.apache.commons.lang.StringUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.script.ScriptException;
+import java.util.List;
+import java.util.Map;
+
+@Singleton
+public class MetricsService {
+    private static final Logger LOG = LoggerFactory.getLogger(MetricsService.class);
+
+    public static final String METRIC_QUERY_PREFIX = "atlas.metric.query.";
+
+    public static final String TYPE = "type";
+    public static final String ENTITY = "entity";
+    public static final String TAG = "tag";
+    public static final String GENERAL = "general";
+
+    public static final String METRIC_TYPE_COUNT = TYPE + "Count";
+    public static final String METRIC_TYPE_UNUSED_COUNT = TYPE + "UnusedCount";
+    public static final String METRIC_TYPE_ENTITIES = TYPE + "Entities";
+
+    public static final String METRIC_ENTITY_COUNT = ENTITY + "Count";
+    public static final String METRIC_ENTITY_DELETED = ENTITY + "Deleted";
+    public static final String METRIC_TAGGED_ENTITIES = ENTITY + "Tagged";
+    public static final String METRIC_TAGS_PER_ENTITY = ENTITY + "Tags";
+
+    public static final String METRIC_TAG_COUNT = TAG + "Count";
+    public static final String METRIC_ENTITIES_PER_TAG = TAG + "Entities";
+
+    private static AtlasGraph atlasGraph;
+    private static Configuration configuration;
+
+    public MetricsService() throws AtlasException {
+        atlasGraph = AtlasGraphProvider.getGraphInstance();
+        configuration = ApplicationProperties.get();
+    }
+
+    @SuppressWarnings("unchecked")
+    public AtlasMetrics getMetrics() {
+        AtlasMetrics metrics = new AtlasMetrics();
+        for (MetricQuery metricQuery : MetricQuery.values()) {
+            try {
+                Object result = atlasGraph.executeGremlinScript(metricQuery.query, false);
+                if (result instanceof Number) {
+                    metrics.addData(metricQuery.type, metricQuery.name, ((Number) result).intValue());
+                } else if (result instanceof List) {
+                    for (Map resultMap : (List<Map>) result) {
+                        metrics.addData(metricQuery.type, (String) resultMap.get("key"), ((Number) resultMap.get("value")).intValue());
+                    }
+                } else {
+                    LOG.warn("Unhandled return type {} for {}. Ignoring", result.getClass().getSimpleName(), metricQuery);
+                }
+            } catch (ScriptException e) {
+                LOG.error("Gremlin execution failed for metric {}", metricQuery, e);
+            }
+        }
+
+        return metrics;
+    }
+
+    /**
+     * MetricQuery enum has the capability of reading the queries from the externalized config.
+     *
+     * The default behavior is to read from the properties and override the statically type query if the configured
+     * query is not blank/empty.
+     */
+    enum MetricQuery {
+        TYPE_COUNT(GENERAL, METRIC_TYPE_COUNT, "g.V().has('__type', 'typeSystem').filter({it.'__type.category'.name() != 'TRAIT'}).count()"),
+        UNUSED_TYPE_COUNT(GENERAL, METRIC_TYPE_UNUSED_COUNT, "g.V('__type', 'typeSystem').filter({ it.'__type.category'.name() != 'TRAIT' && it.inE.count() == 0}).count()"),
+        ENTITY_COUNT(GENERAL, METRIC_ENTITY_COUNT, "g.V().has('__superTypeNames', T.in, ['Referenceable']).count()"),
+        TAGS_COUNT(GENERAL, METRIC_TAG_COUNT, "g.V().has('__type', 'typeSystem').filter({it.'__type.category'.name() == 'TRAIT'}).count()"),
+        DELETED_ENTITY_COUNT(GENERAL, METRIC_ENTITY_DELETED, "g.V().has('__superTypeNames', T.in, ['Referenceable']).has('__status', 'DELETED').count()"),
+
+        ENTITIES_PER_TYPE(ENTITY, METRIC_TYPE_ENTITIES, "g.V().has('__type', 'typeSystem').has('__type.name').filter({it.'__type.category'.name() != 'TRAIT'}).transform{[key: it.'__type.name', value: it.inE.count()]}.dedup().toList()"),
+        TAGGED_ENTITIES(ENTITY, METRIC_TAGGED_ENTITIES, "g.V().has('__superTypeNames', T.in, ['Referenceable']).has('__traitNames').count()"),
+
+        TAGS_PER_ENTITY(TAG, METRIC_TAGS_PER_ENTITY, "g.V().has('__superTypeNames', T.in, ['Referenceable']).has('__traitNames').transform{[ key: it.'Referenceable.qualifiedName', value: it.'__traitNames'.size()]}.dedup().toList()"),
+        ;
+        private String type;
+        private String name;
+        private String query;
+
+        private static String getQuery(String type, String name) {
+            String metricQueryKey = METRIC_QUERY_PREFIX + type + "." + name;
+            if (LOG.isDebugEnabled()) {
+                LOG.debug("Looking for configured query {}", metricQueryKey);
+            }
+            return configuration.getString(metricQueryKey, "");
+        }
+
+        MetricQuery(String type, String name, String query) {
+            this.type = type;
+            this.name = name;
+            String configuredQuery = getQuery(type, name);
+            this.query = StringUtils.isNotEmpty(configuredQuery) ? configuredQuery : query;
+        }
+
+        @Override
+        public String toString() {
+            return "MetricQuery{" + "type='" + type + '\'' +
+                    ", name='" + name + '\'' +
+                    ", query='" + query + '\'' +
+                    '}';
+        }
+    }
+}

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

@@ -20,6 +20,8 @@ package org.apache.atlas.web.resources;
 
 import com.google.inject.Inject;
 import org.apache.atlas.AtlasClient;
+import org.apache.atlas.model.metrics.AtlasMetrics;
+import org.apache.atlas.services.MetricsService;
 import org.apache.atlas.authorize.AtlasActionTypes;
 import org.apache.atlas.authorize.AtlasResourceTypes;
 import org.apache.atlas.authorize.simple.AtlasAuthorizationUtils;
@@ -65,10 +67,12 @@ public class AdminResource {
     private static final String isEntityCreateAllowed = "atlas.entity.create.allowed";
     private Response version;
     private ServiceState serviceState;
+    private MetricsService metricsService;
 
     @Inject
-    public AdminResource(ServiceState serviceState) {
+    public AdminResource(ServiceState serviceState, MetricsService metricsService) {
         this.serviceState = serviceState;
+        this.metricsService = metricsService;
     }
 
     /**
@@ -223,4 +227,21 @@ public class AdminResource {
 
         return response;
     }
+
+    @GET
+    @Path("metrics")
+    @Produces(Servlets.JSON_MEDIA_TYPE)
+    public AtlasMetrics getMetrics() {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> AdminResource.getMetrics()");
+        }
+
+        AtlasMetrics metrics = metricsService.getMetrics();
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("<== AdminResource.getMetrics()");
+        }
+
+        return metrics;
+    }
 }

webapp/src/main/resources/spring-security.xml

@@ -33,6 +33,7 @@
     <security:http pattern="/js/**" security="none" />
     <security:http pattern="/ieerror.html" security="none" />
     <security:http pattern="/api/atlas/admin/status" security="none" />
+    <security:http pattern="/api/atlas/admin/metrics" security="none" />
 
     <security:http disable-url-rewriting="true"
                    use-expressions="true" create-session="always"

webapp/src/test/java/org/apache/atlas/web/resources/AdminResourceTest.java

@@ -48,7 +48,7 @@ public class AdminResourceTest {
 
         when(serviceState.getState()).thenReturn(ServiceState.ServiceStateValue.ACTIVE);
 
-        AdminResource adminResource = new AdminResource(serviceState);
+        AdminResource adminResource = new AdminResource(serviceState, null);
         Response response = adminResource.getStatus();
         assertEquals(response.getStatus(), HttpServletResponse.SC_OK);
         JSONObject entity = (JSONObject) response.getEntity();
@@ -59,7 +59,7 @@ public class AdminResourceTest {
     public void testResourceGetsValueFromServiceState() throws JSONException {
         when(serviceState.getState()).thenReturn(ServiceState.ServiceStateValue.PASSIVE);
 
-        AdminResource adminResource = new AdminResource(serviceState);
+        AdminResource adminResource = new AdminResource(serviceState, null);
         Response response = adminResource.getStatus();
 
         verify(serviceState).getState();