Merge pull request #2370 from atlanhq/rawquerymigration

[stag] PLT-1886 Migrate Query.rawQuery to Query.rawQueryText
atlanhq · Sep 15, 2023 · 6e8474b · 6e8474b
2 parents 736597e + a428b48
commit 6e8474b
Showing 1 changed file with 21 additions and 9 deletions.
diff --git a/webapp/src/main/java/org/apache/atlas/web/rest/EntityREST.java b/webapp/src/main/java/org/apache/atlas/web/rest/EntityREST.java
@@ -69,8 +69,6 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.util.*;
-import java.util.function.Predicate;
-import java.util.stream.Collectors;
 
 import static org.apache.atlas.AtlasErrorCode.BAD_REQUEST;
 import static org.apache.atlas.AtlasErrorCode.DEPRECATED_API;
@@ -93,6 +91,10 @@ public class EntityREST {
     public static final String PREFIX_ATTR_ = "attr_";
     public static final String QUALIFIED_NAME  = "qualifiedName";
     private static final int HUNDRED_THOUSAND = 100000;
+    private static final int TWO_MILLION = HUNDRED_THOUSAND * 10 * 2;
+    private static final Set<String> ATTRS_WITH_TWO_MILLION_LIMIT = new HashSet<String>() {{
+        add("rawQueryText");
+    }};
 
 
     private final AtlasTypeRegistry      typeRegistry;
@@ -899,15 +901,25 @@ public EntityMutationResponse createOrUpdate(AtlasEntitiesWithExtInfo entities,
     }
 
     public static void validateAttributeLength(final List<AtlasEntity> entities) throws AtlasBaseException {
-        //Predicate to check attribute value exceeding length
-        Predicate<Map.Entry<String, Object>> predicateOfAttributeLengthExceedingLimit = attribute ->
-                attribute.getValue() instanceof String && ((String) attribute.getValue()).length() > HUNDRED_THOUSAND;
+        List<String> errorMessages = new ArrayList<>();
 
         for (final AtlasEntity atlasEntity : entities) {
-            Set<String> attributeKeys = org.apache.commons.collections4.MapUtils.emptyIfNull(atlasEntity.getAttributes())
-                    .entrySet().stream().filter(predicateOfAttributeLengthExceedingLimit).map(Map.Entry::getKey).collect(Collectors.toSet());
-            if (!attributeKeys.isEmpty()) {
-                throw new AtlasBaseException("Attribute(s) " + String.join(",", attributeKeys) + " exceeds limit of "+HUNDRED_THOUSAND+" characters");
+            for (Map.Entry<String, Object> attribute : atlasEntity.getAttributes().entrySet()) {
+
+                if (attribute.getValue() instanceof String && ((String) attribute.getValue()).length() > HUNDRED_THOUSAND) {
+
+                    if (ATTRS_WITH_TWO_MILLION_LIMIT.contains(attribute.getKey())) {
+                        if (((String) attribute.getValue()).length() > TWO_MILLION) {
+                            errorMessages.add("Attribute " + attribute.getKey() + " exceeds limit of " + TWO_MILLION + " characters");
+                        }
+                    } else {
+                        errorMessages.add("Attribute " + attribute.getKey() + " exceeds limit of " + HUNDRED_THOUSAND + " characters");
+                    }
+                }
+            }
+
+            if (errorMessages.size() > 0) {
+                throw new AtlasBaseException(AtlasType.toJson(errorMessages));
             }
         }
     }