add guard rails

atlanhq · Nov 29, 2024 · 7dd9642 · 7dd9642
1 parent e060f61
commit 7dd9642
Show file tree

Hide file tree

Showing 6 changed files with 154 additions and 2 deletions.
diff --git a/common/src/main/java/org/apache/atlas/utils/AtlasPerfMetrics.java b/common/src/main/java/org/apache/atlas/utils/AtlasPerfMetrics.java
@@ -50,6 +50,24 @@ public void recordMetric(MetricRecorder recorder) {
         }
     }
 
+    public void recordMetricWithInvocations(MetricRecorder recorder, long count) {
+        if (recorder != null) {
+            final String name = recorder.name;
+            final long timeTaken = recorder.getElapsedTime();
+
+            Metric metric = metrics.get(name);
+
+            if (metric == null) {
+                metric = new Metric(name);
+
+                metrics.put(name, metric);
+            }
+
+            metric.invocations += count;
+            metric.totalTimeMSecs += timeTaken;
+        }
+    }
+
     public void clear() {
         metrics.clear();
     }
@@ -129,7 +147,6 @@ public String getName() {
         public long getInvocations() {
             return invocations;
         }
-
         public void setTotalTimeMSecs(long totalTimeMSecs) {
             this.totalTimeMSecs = totalTimeMSecs;
         }

diff --git a/intg/src/main/java/org/apache/atlas/AtlasConfiguration.java b/intg/src/main/java/org/apache/atlas/AtlasConfiguration.java
@@ -117,8 +117,9 @@ public enum AtlasConfiguration {
     ATLAS_INDEXSEARCH_QUERY_SIZE_MAX_LIMIT("atlas.indexsearch.query.size.max.limit", 100000),
     ATLAS_INDEXSEARCH_LIMIT_UTM_TAGS("atlas.indexsearch.limit.ignore.utm.tags", ""),
     ATLAS_INDEXSEARCH_ENABLE_API_LIMIT("atlas.indexsearch.enable.api.limit", false),
+    ATLAS_INDEXSEARCH_ENABLE_JANUS_OPTIMISATION("atlas.indexsearch.enable.janus.optimization", false),
     ATLAS_INDEXSEARCH_ENABLE_FETCHING_NON_PRIMITIVE_ATTRIBUTES("atlas.indexsearch.enable.fetching.non.primitive.attributes", false),
-
+    ATLAS_INDEXSEARCH_ATTRIBUTES_MIN_LIMIT("atlas.indexsearch.attributes.min.limit", 8),
     ATLAS_MAINTENANCE_MODE("atlas.maintenance.mode", false),
 
     ATLAS_UD_RELATIONSHIPS_MAX_COUNT("atlas.ud.relationship.max.count", 100);

diff --git a/...sitory/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java b/...sitory/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
@@ -49,6 +49,7 @@
 import org.apache.atlas.repository.graphdb.AtlasElement;
 import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
+import org.apache.atlas.repository.util.AccessControlUtils;
 import org.apache.atlas.type.AtlasArrayType;
 import org.apache.atlas.type.AtlasBuiltInTypes.AtlasObjectIdType;
 import org.apache.atlas.type.AtlasEntityType;
@@ -1027,9 +1028,128 @@ private Map<String, Object> preloadProperties(AtlasVertex entityVertex) {
         return propertiesMap;
     }
 
+
+    private boolean isPolicyAttribute(Set<String> attributes) {
+        Set<String> exclusionSet = new HashSet<>(Arrays.asList(AccessControlUtils.ATTR_POLICY_TYPE,
+                AccessControlUtils.ATTR_POLICY_USERS,
+                AccessControlUtils.ATTR_POLICY_GROUPS,
+                AccessControlUtils.ATTR_POLICY_ROLES,
+                AccessControlUtils.ATTR_POLICY_ACTIONS,
+                AccessControlUtils.ATTR_POLICY_CATEGORY,
+                AccessControlUtils.ATTR_POLICY_SUB_CATEGORY,
+                AccessControlUtils.ATTR_POLICY_RESOURCES,
+                AccessControlUtils.ATTR_POLICY_IS_ENABLED,
+                AccessControlUtils.ATTR_POLICY_RESOURCES_CATEGORY,
+                AccessControlUtils.ATTR_POLICY_SERVICE_NAME,
+                AccessControlUtils.ATTR_POLICY_PRIORITY,
+                AccessControlUtils.REL_ATTR_POLICIES,
+                AccessControlUtils.ATTR_SERVICE_SERVICE_TYPE,
+                AccessControlUtils.ATTR_SERVICE_TAG_SERVICE,
+                AccessControlUtils.ATTR_SERVICE_IS_ENABLED,
+                AccessControlUtils.ATTR_SERVICE_LAST_SYNC));
+
+        return exclusionSet.stream().anyMatch(attributes::contains);
+    }
+
     private AtlasEntityHeader mapVertexToAtlasEntityHeader(AtlasVertex entityVertex, Set<String> attributes) throws AtlasBaseException {
+        boolean shouldPrefetch = attributes.size() > AtlasConfiguration.ATLAS_INDEXSEARCH_ATTRIBUTES_MIN_LIMIT.getInt()
+                && !isPolicyAttribute(attributes)
+                && AtlasConfiguration.ATLAS_INDEXSEARCH_ENABLE_JANUS_OPTIMISATION.getBoolean();
+
+        if (shouldPrefetch) {
+            return mapVertexToAtlasEntityHeaderWithPrefetch(entityVertex, attributes);
+        } else {
+            return mapVertexToAtlasEntityHeaderWithoutPrefetch(entityVertex, attributes);
+        }
+    }
+
+    private AtlasEntityHeader mapVertexToAtlasEntityHeaderWithoutPrefetch(AtlasVertex entityVertex, Set<String> attributes) throws AtlasBaseException {
         AtlasPerfMetrics.MetricRecorder metricRecorder = RequestContext.get().startMetricRecord("mapVertexToAtlasEntityHeader");
         AtlasEntityHeader ret = new AtlasEntityHeader();
+        try {
+            String  typeName     = entityVertex.getProperty(Constants.TYPE_NAME_PROPERTY_KEY, String.class);
+            String  guid         = entityVertex.getProperty(Constants.GUID_PROPERTY_KEY, String.class);
+            Boolean isIncomplete = isEntityIncomplete(entityVertex);
+
+            ret.setTypeName(typeName);
+            ret.setGuid(guid);
+            ret.setStatus(GraphHelper.getStatus(entityVertex));
+            RequestContext context = RequestContext.get();
+            boolean includeClassifications = context.includeClassifications();
+            boolean includeClassificationNames = context.isIncludeClassificationNames();
+            if(includeClassifications){
+                ret.setClassificationNames(getAllTraitNamesFromAttribute(entityVertex));
+            } else if (!includeClassifications && includeClassificationNames) {
+                ret.setClassificationNames(getAllTraitNamesFromAttribute(entityVertex));
+            }
+            ret.setIsIncomplete(isIncomplete);
+            ret.setLabels(getLabels(entityVertex));
+
+            ret.setCreatedBy(GraphHelper.getCreatedByAsString(entityVertex));
+            ret.setUpdatedBy(GraphHelper.getModifiedByAsString(entityVertex));
+            ret.setCreateTime(new Date(GraphHelper.getCreatedTime(entityVertex)));
+            ret.setUpdateTime(new Date(GraphHelper.getModifiedTime(entityVertex)));
+
+            if(RequestContext.get().includeMeanings()) {
+                List<AtlasTermAssignmentHeader> termAssignmentHeaders = mapAssignedTerms(entityVertex);
+                ret.setMeanings(termAssignmentHeaders);
+                ret.setMeaningNames(
+                        termAssignmentHeaders.stream().map(AtlasTermAssignmentHeader::getDisplayText)
+                                .collect(Collectors.toList()));
+            }
+            AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName);
+
+            if (entityType != null) {
+                for (AtlasAttribute headerAttribute : entityType.getHeaderAttributes().values()) {
+                    Object attrValue = getVertexAttribute(entityVertex, headerAttribute);
+
+                    if (attrValue != null) {
+                        ret.setAttribute(headerAttribute.getName(), attrValue);
+                    }
+                }
+
+                Object displayText = getDisplayText(entityVertex, entityType);
+
+                if (displayText != null) {
+                    ret.setDisplayText(displayText.toString());
+                }
+
+                if (CollectionUtils.isNotEmpty(attributes)) {
+                    for (String attrName : attributes) {
+                        AtlasAttribute attribute = entityType.getAttribute(attrName);
+
+                        if (attribute == null) {
+                            attrName = toNonQualifiedName(attrName);
+
+                            if (ret.hasAttribute(attrName)) {
+                                continue;
+                            }
+
+                            attribute = entityType.getAttribute(attrName);
+
+                            if (attribute == null) {
+                                attribute = entityType.getRelationshipAttribute(attrName, null);
+                            }
+                        }
+
+                        Object attrValue = getVertexAttribute(entityVertex, attribute);
+
+                        if (attrValue != null) {
+                            ret.setAttribute(attrName, attrValue);
+                        }
+                    }
+                }
+            }
+        }
+        finally {
+            RequestContext.get().endMetricRecord(metricRecorder);
+        }
+        return ret;
+    }
+
+    private AtlasEntityHeader mapVertexToAtlasEntityHeaderWithPrefetch(AtlasVertex entityVertex, Set<String> attributes) throws AtlasBaseException {
+        AtlasPerfMetrics.MetricRecorder metricRecorder = RequestContext.get().startMetricRecord("mapVertexToAtlasEntityHeaderWithPrefetch");
+        AtlasEntityHeader ret = new AtlasEntityHeader();
         try {
             //pre-fetching the properties
             Map<String, Object> properties = preloadProperties(entityVertex);

diff --git a/repository/src/main/java/org/apache/atlas/repository/util/AccessControlUtils.java b/repository/src/main/java/org/apache/atlas/repository/util/AccessControlUtils.java
@@ -59,6 +59,11 @@ public final class AccessControlUtils {
     public static final String ATTR_PERSONA_USERS   = "personaUsers";
     public static final String ATTR_PERSONA_GROUPS  = "personaGroups";
 
+    public static final String ATTR_SERVICE_SERVICE_TYPE = "authServiceType";
+    public static final String ATTR_SERVICE_TAG_SERVICE  = "tagService";
+    public static final String ATTR_SERVICE_IS_ENABLED   = "authServiceIsEnabled";
+    public static final String ATTR_SERVICE_LAST_SYNC    = "authServicePolicyLastSync";
+
     public static final String ATTR_PURPOSE_CLASSIFICATIONS  = "purposeClassifications";
 
     public static final String ATTR_POLICY_TYPE  = "policyType";

diff --git a/server-api/src/main/java/org/apache/atlas/RequestContext.java b/server-api/src/main/java/org/apache/atlas/RequestContext.java
@@ -657,6 +657,13 @@ public void endMetricRecord(MetricRecorder recorder) {
         }
     }
 
+    public void endMetricRecordWithInvocations(MetricRecorder recorder, long invocationCount) {
+        if (metrics != null && recorder != null) {
+            metrics.recordMetricWithInvocations(recorder, invocationCount);
+        }
+    }
+
+
     public void recordEntityGuidUpdate(AtlasEntity entity, String guidInRequest) {
         recordEntityGuidUpdate(new EntityGuidPair(entity, guidInRequest));
     }

diff --git a/webapp/src/main/java/org/apache/atlas/web/rest/DiscoveryREST.java b/webapp/src/main/java/org/apache/atlas/web/rest/DiscoveryREST.java
@@ -401,6 +401,8 @@ public AtlasSearchResult indexSearch(@Context HttpServletRequest servletRequest,
                 perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "DiscoveryREST.indexSearch(" + parameters + ")");
             }
 
+            RequestContext.get().endMetricRecordWithInvocations(RequestContext.get().startMetricRecord("dslSize"), parameters.getQuerySize());
+            RequestContext.get().endMetricRecordWithInvocations(RequestContext.get().startMetricRecord("attributeSize"), parameters.getAttributes().size());
             if (AtlasConfiguration.ATLAS_INDEXSEARCH_ENABLE_API_LIMIT.getBoolean() && parameters.getQuerySize() > AtlasConfiguration.ATLAS_INDEXSEARCH_QUERY_SIZE_MAX_LIMIT.getLong()) {
                 if(CollectionUtils.isEmpty(parameters.getUtmTags())) {
                     throw new AtlasBaseException(AtlasErrorCode.INVALID_DSL_QUERY_SIZE, String.valueOf(AtlasConfiguration.ATLAS_INDEXSEARCH_QUERY_SIZE_MAX_LIMIT.getLong()));