From a4a1c61b6d25659ed3861a7cb8a1e6c4ed5559bb Mon Sep 17 00:00:00 2001
From: Suman Das <59254445+sumandas0@users.noreply.github.com>
Date: Wed, 11 Oct 2023 15:51:02 +0530
Subject: [PATCH] feat: Add exception message for error batch too large

---
 .../java/org/apache/atlas/AtlasErrorCode.java |  3 +-
 .../atlas/GraphTransactionInterceptor.java    | 29 +++++++++++++++----
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/intg/src/main/java/org/apache/atlas/AtlasErrorCode.java b/intg/src/main/java/org/apache/atlas/AtlasErrorCode.java
index 05a7b72b16..0abe2a7521 100644
--- a/intg/src/main/java/org/apache/atlas/AtlasErrorCode.java
+++ b/intg/src/main/java/org/apache/atlas/AtlasErrorCode.java
@@ -269,7 +269,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
index 5ce23cf5a4..90e937e3f7 100644
--- a/repository/src/main/java/org/apache/atlas/GraphTransactionInterceptor.java
+++ b/repository/src/main/java/org/apache/atlas/GraphTransactionInterceptor.java
@@ -17,6 +17,7 @@
 
 package org.apache.atlas;
 
+import com.datastax.oss.driver.api.core.servererrors.InvalidQueryException;
 import com.google.common.annotations.VisibleForTesting;
 import org.aopalliance.intercept.MethodInterceptor;
 import org.aopalliance.intercept.MethodInvocation;
@@ -34,11 +35,7 @@
 import javax.inject.Inject;
 import javax.ws.rs.core.Response;
 import java.lang.reflect.Method;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.locks.ReentrantLock;
 
@@ -128,7 +125,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 +168,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().equals(InvalidQueryException.class)) {
+                return true;
+            }
+            currentCause = currentCause.getCause();
+        }
+        return false;
+    }
+
+
+
     private void doCommitOrRollback(final String invokingClass, final String invokedMethodName) {
         if (innerFailure.get()) {
             if (LOG.isDebugEnabled()) {