diff --git a/repository/src/main/java/org/apache/atlas/GraphTransactionInterceptor.java b/repository/src/main/java/org/apache/atlas/GraphTransactionInterceptor.java
index 90e937e3f7..e34655a79b 100644
--- a/repository/src/main/java/org/apache/atlas/GraphTransactionInterceptor.java
+++ b/repository/src/main/java/org/apache/atlas/GraphTransactionInterceptor.java
@@ -17,7 +17,6 @@
 
 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;
@@ -174,7 +173,7 @@ public boolean checkForBatchTooLargeError(Throwable t) {
             String message = currentCause.getMessage();
             if (message != null &&
                     message.contains("Batch too large") &&
-                    currentCause.getClass().equals(InvalidQueryException.class)) {
+                    currentCause.getClass().getSimpleName().equals("InvalidQueryException")) {
                 return true;
             }
             currentCause = currentCause.getCause();