Merge pull request #2412 from atlanhq/batch-error

Add special error code for Batch size limit error
atlanhq · Oct 20, 2023 · 470a353 · 470a353
2 parents 3935ae6 + 0ee1385
commit 470a353
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/intg/src/main/java/org/apache/atlas/AtlasErrorCode.java b/intg/src/main/java/org/apache/atlas/AtlasErrorCode.java
@@ -271,7 +271,8 @@ public enum AtlasErrorCode {
     FAILED_TO_REFRESH_TYPE_DEF_CACHE(500, "ATLAS-500-00-20", "Failed to refresh type-def cache"),
     CINV_UNHEALTHY(500, "ATLAS-500-00-21", "Unable to process type-definition operations"),
     RUNTIME_EXCEPTION(500, "ATLAS-500-00-020", "Runtime exception {0}"),
-    KEYCLOAK_INIT_FAILED(500, "ATLAS-500-00-021", "Failed to initialize keycloak client: {0}"),
+    KEYCLOAK_INIT_FAILED(500, "ATLAS-500-00-022", "Failed to initialize keycloak client: {0}"),
+    BATCH_SIZE_TOO_LARGE(500, "ATLAS-500-00-023", "Batch size is too large, please use a smaller batch size"),
 
 
     CLASSIFICATION_CURRENTLY_BEING_PROPAGATED(400, "ATLAS-400-00-105", "Classification {0} is currently being propagated."),

diff --git a/repository/src/main/java/org/apache/atlas/GraphTransactionInterceptor.java b/repository/src/main/java/org/apache/atlas/GraphTransactionInterceptor.java
@@ -128,7 +128,11 @@ public Object invoke(MethodInvocation invocation) throws Throwable {
                 } else {
                     doRollback(logRollback, t);
                 }
-                throw t;
+                if (checkForBatchTooLargeError(t)) {
+                    throw new AtlasBaseException(AtlasErrorCode.BATCH_SIZE_TOO_LARGE, t);
+                } else {
+                    throw t;
+                }
             }
         } finally {
             RequestContext.get().endMetricRecord(metric);
@@ -167,6 +171,22 @@ public Object invoke(MethodInvocation invocation) throws Throwable {
         }
     }
 
+    public boolean checkForBatchTooLargeError(Throwable t) {
+        Throwable currentCause = t;
+        while (currentCause != null) {
+            String message = currentCause.getMessage();
+            if (message != null &&
+                    message.contains("Batch too large") &&
+                    currentCause.getClass().getSimpleName().equals("InvalidQueryException")) {
+                return true;
+            }
+            currentCause = currentCause.getCause();
+        }
+        return false;
+    }
+
+
+
     private void doCommitOrRollback(final String invokingClass, final String invokedMethodName) {
         if (innerFailure.get()) {
             if (LOG.isDebugEnabled()) {