[VUMM-344] Exception handling for spring async endpoint (#3033)

* Spring boot async endpoint

* https://jenkins.dev.verta.ai/job/build/job/autoformat/job/modeldb-backend/1022/

* Implemented Async endpoint using CompletableFuture<T>

* https://jenkins.dev.verta.ai/job/build/job/autoformat/job/modeldb-backend/1029/

* Exception handling for spring async endpoint

* https://jenkins.dev.verta.ai/job/build/job/autoformat/job/modeldb-backend/1036/

* Fix compilation for exception code changes for MDB

* Fix error code as per https://github.com/grpc-ecosystem/grpc-gateway/blob/master/runtime/errors.go instead of manual handling

* https://jenkins.dev.verta.ai/job/build/job/autoformat/job/modeldb-backend/1041/

* Fix compilation for new changes

* fix error message formate

* Removed error message from response while getting 500

* fix method name with lower case

* fix message for internal server error

* Fix import

* Extracted to future utils

* https://jenkins.dev.verta.ai/job/build/job/autoformat/job/modeldb-backend/1055/

* Created final concrete class instead of abstraction

* https://jenkins.dev.verta.ai/job/build/job/autoformat/job/modeldb-backend/1065/

* Fix review comment on method naming convention

* Fix compilation due to merge conflicts

* Fix review comment

* https://jenkins.dev.verta.ai/job/build/job/autoformat/job/modeldb-backend/86/

* Added unit test for CommonUtils.observeErrorWithResponse

* https://jenkins.dev.verta.ai/job/build/job/autoformat/job/modeldb-backend/98/

Co-authored-by: Jenkins <jenkins@verta.ai>

backend/common/src/main/java/ai/verta/modeldb/common/CommonUtils.java

@@ -23,10 +23,14 @@
 import java.net.SocketException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
+import java.util.concurrent.ExecutionException;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.jdbi.v3.core.statement.Query;
+import org.springframework.http.HttpStatus;
+import org.springframework.web.server.ResponseStatusException;
 
 public class CommonUtils {
   private static final Logger LOGGER = LogManager.getLogger(CommonUtils.class);
@@ -174,6 +178,64 @@ public static <T extends GeneratedMessageV3> void observeError(
     responseObserver.onError(logError(e));
   }
 
+  public static <T> void observeError(CompletableFuture<T> completableFuture, Throwable e) {
+    completableFuture.completeExceptionally(observeErrorWithResponse(e, e.getMessage(), LOGGER));
+  }
+
+  public static ResponseStatusException observeErrorWithResponse(
+      Throwable e, String message, Logger logger) {
+    if (e instanceof ModelDBException
+        && e.getCause() instanceof ExecutionException
+        && e.getCause().getCause() != null) {
+      return observeErrorWithResponse(e.getCause().getCause(), message, logger);
+    }
+
+    if (e instanceof CompletionException) {
+      return observeErrorWithResponse(e.getCause(), message, logger);
+    }
+    final HttpStatus code;
+    if (e instanceof StatusRuntimeException) {
+      StatusRuntimeException ex = (StatusRuntimeException) e;
+      code = ModelDBException.httpStatusFromCode(ex.getStatus().getCode());
+      if (ex.getStatus().getCode() == io.grpc.Status.Code.NOT_FOUND) {
+        logger.debug(e.getMessage());
+      } else {
+        logger.error(e.getMessage());
+      }
+      message += " Details: " + e.getMessage();
+      printStackTrace(logger, e);
+    } else if (e instanceof IllegalArgumentException) {
+      code = HttpStatus.BAD_REQUEST;
+      String logMessage = "Common Service Exception occurred: {}";
+      message = e.getMessage();
+      logger.debug(logMessage, message);
+    } else if (e instanceof ModelDBException) {
+      ModelDBException uacServiceException = (ModelDBException) e;
+
+      code = uacServiceException.getHttpCode();
+      message += " Details: " + e.getMessage();
+      String logMessage = "Common Service Exception occurred: {}";
+      if (uacServiceException.getCodeValue() == Code.INTERNAL_VALUE) {
+        // If getting 500 then we will not expose about error to the user.
+        message = "Internal Server Error";
+        logger.error(logMessage, e.getMessage());
+        printStackTrace(logger, e);
+      } else {
+        if (uacServiceException.getCodeValue() == Code.RESOURCE_EXHAUSTED_VALUE) {
+          message = e.getMessage();
+        }
+        logger.debug(logMessage, e.getMessage());
+      }
+    } else {
+      code = HttpStatus.INTERNAL_SERVER_ERROR;
+      // If getting 500 then we will not expose about error to the user.
+      message = "Internal Server Error";
+      logger.error(e.getMessage());
+      printStackTrace(logger, e);
+    }
+    return new ResponseStatusException(code, message, e);
+  }
+
   public static <T extends GeneratedMessageV3> void observeError(
       StreamObserver<T> responseObserver, Exception e, T defaultInstance) {
     responseObserver.onError(logError(e, defaultInstance));

backend/common/src/main/java/ai/verta/modeldb/common/exceptions/ModelDBException.java

@@ -1,9 +1,13 @@
 package ai.verta.modeldb.common.exceptions;
 
 import io.grpc.Status.Code;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.springframework.http.HttpStatus;
 
 public class ModelDBException extends RuntimeException {
 
+  private static final Logger LOGGER = LogManager.getLogger(ModelDBException.class);
   private static final long serialVersionUID = 1L;
 
   private final Code code;
@@ -47,9 +51,56 @@ public Code getCode() {
     return code;
   }
 
+  public int getCodeValue() {
+    return getCode().value();
+  }
+
   protected ModelDBException(
       String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) {
     super(message, cause, enableSuppression, writableStackTrace);
     code = Code.INTERNAL;
   }
+
+  public HttpStatus getHttpCode() {
+    return httpStatusFromCode(this.code);
+  }
+
+  public static HttpStatus httpStatusFromCode(Code code) {
+    switch (code) {
+      case OK:
+        return HttpStatus.OK;
+      case CANCELLED:
+        return HttpStatus.REQUEST_TIMEOUT;
+      case UNKNOWN:
+      case INTERNAL:
+      case DATA_LOSS:
+        return HttpStatus.INTERNAL_SERVER_ERROR;
+      case INVALID_ARGUMENT:
+      case FAILED_PRECONDITION:
+      case OUT_OF_RANGE:
+        // Note, this deliberately doesn't translate to the similarly named '412 Precondition
+        // Failed' HTTP response status.
+        return HttpStatus.BAD_REQUEST;
+      case DEADLINE_EXCEEDED:
+        return HttpStatus.GATEWAY_TIMEOUT;
+      case NOT_FOUND:
+        return HttpStatus.NOT_FOUND;
+      case ALREADY_EXISTS:
+      case ABORTED:
+        return HttpStatus.CONFLICT;
+      case PERMISSION_DENIED:
+        return HttpStatus.FORBIDDEN;
+      case UNAUTHENTICATED:
+        return HttpStatus.UNAUTHORIZED;
+      case RESOURCE_EXHAUSTED:
+        return HttpStatus.TOO_MANY_REQUESTS;
+      case UNIMPLEMENTED:
+        return HttpStatus.NOT_IMPLEMENTED;
+      case UNAVAILABLE:
+        return HttpStatus.SERVICE_UNAVAILABLE;
+      default:
+        LOGGER.info("Unknown gRPC error code: {}", code);
+        return HttpStatus.INTERNAL_SERVER_ERROR;
+    }
+  }
 }

backend/common/src/main/java/ai/verta/modeldb/common/futures/FutureRest.java

@@ -1,5 +1,6 @@
 package ai.verta.modeldb.common.futures;
 
+import ai.verta.modeldb.common.CommonUtils;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 
@@ -15,10 +16,10 @@ public static <T> CompletableFuture<T> serverResponse(InternalFuture<T> future,
             try {
               promise.complete(v);
             } catch (Throwable e) {
-              promise.completeExceptionally(e);
+              CommonUtils.observeError(promise, t);
             }
           } else {
-            promise.completeExceptionally(t);
+            CommonUtils.observeError(promise, t);
           }
         },
         ex);

backend/src/main/java/ai/verta/modeldb/datasetVersion/DatasetVersionServiceImpl.java

@@ -174,7 +174,7 @@ private DatasetVersion getDatasetVersionFromRequest(
 
     if (!request.hasPathDatasetVersionInfo() && !request.hasDatasetBlob()) {
       LOGGER.info("Request {}", request);
-      throw new ModelDBException("Not supported", io.grpc.Status.Code.UNIMPLEMENTED);
+      throw new ModelDBException("Not supported", Code.UNIMPLEMENTED);
     }
     datasetVersionBuilder.setPathDatasetVersionInfo(request.getPathDatasetVersionInfo());
 

backend/src/main/java/ai/verta/modeldb/entities/config/ConfigBlobEntity.java

@@ -1,7 +1,7 @@
 package ai.verta.modeldb.entities.config;
 
 import ai.verta.modeldb.common.exceptions.ModelDBException;
-import io.grpc.Status;
+import com.google.rpc.Code;
 import java.io.Serializable;
 import java.util.Objects;
 import javax.persistence.CascadeType;
@@ -34,7 +34,7 @@ public ConfigBlobEntity(String blobHash, Integer configSeqNumber, Object blobEnt
           (HyperparameterElementConfigBlobEntity) blobEntity;
       this.hyperparameter_type = HYPERPARAMETER;
     } else {
-      throw new ModelDBException("Invalid blob object found", Status.Code.INVALID_ARGUMENT);
+      throw new ModelDBException("Invalid blob object found", Code.INVALID_ARGUMENT);
     }
   }
 

backend/src/main/java/ai/verta/modeldb/entities/config/HyperparameterElementConfigBlobEntity.java

@@ -2,7 +2,7 @@
 
 import ai.verta.modeldb.common.exceptions.ModelDBException;
 import ai.verta.modeldb.versioning.HyperparameterValuesConfigBlob;
-import io.grpc.Status;
+import com.google.rpc.Code;
 import java.io.Serializable;
 import javax.persistence.Column;
 import javax.persistence.Entity;
@@ -37,7 +37,7 @@ public HyperparameterElementConfigBlobEntity(
       case VALUE_NOT_SET:
       default:
         throw new ModelDBException(
-            "Invalid value found in HyperparameterValuesConfigBlob", Status.Code.INVALID_ARGUMENT);
+            "Invalid value found in HyperparameterValuesConfigBlob", Code.INVALID_ARGUMENT);
     }
   }
 

backend/src/main/java/ai/verta/modeldb/entities/config/HyperparameterElementMappingEntity.java

@@ -4,7 +4,7 @@
 import ai.verta.modeldb.common.exceptions.ModelDBException;
 import ai.verta.modeldb.entities.ExperimentRunEntity;
 import ai.verta.modeldb.versioning.HyperparameterValuesConfigBlob;
-import io.grpc.Status;
+import com.google.rpc.Code;
 import java.io.Serializable;
 import java.util.Objects;
 import javax.persistence.*;
@@ -36,7 +36,7 @@ public HyperparameterElementMappingEntity(
       case VALUE_NOT_SET:
       default:
         throw new ModelDBException(
-            "Invalid value found in HyperparameterValuesConfigBlob", Status.Code.INVALID_ARGUMENT);
+            "Invalid value found in HyperparameterValuesConfigBlob", Code.INVALID_ARGUMENT);
     }
 
     this.entity_type = entity.getClass().getSimpleName();

backend/src/main/java/ai/verta/modeldb/experimentRun/ExperimentRunDAORdbImpl.java

@@ -210,7 +210,7 @@ private Map<String, Map.Entry<BlobExpanded, String>> validateVersioningEntity(
     }
 
     if (errorMessage != null) {
-      throw new ModelDBException(errorMessage, io.grpc.Status.Code.INVALID_ARGUMENT);
+      throw new ModelDBException(errorMessage, Code.INVALID_ARGUMENT);
     }
     var repositoryIdentification =
         RepositoryIdentification.newBuilder().setRepoId(versioningEntry.getRepositoryId()).build();
@@ -234,7 +234,7 @@ private Map<String, Map.Entry<BlobExpanded, String>> validateVersioningEntity(
                   + "' for key '"
                   + locationBlobKeyMap.getKey()
                   + "' not found in commit blobs",
-              io.grpc.Status.Code.INVALID_ARGUMENT);
+              Code.INVALID_ARGUMENT);
         }
         requestedLocationBlobWithHashMap.put(locationKey, locationBlobWithHashMap.get(locationKey));
       }
@@ -2083,8 +2083,7 @@ public void logVersionedInput(LogVersionedInput request) throws ModelDBException
                   .equals(versioningModeldbFirstEntityMapping.getCommit())) {
             if (!OVERWRITE_VERSION_MAP) {
               throw new ModelDBException(
-                  ModelDBConstants.DIFFERENT_REPOSITORY_OR_COMMIT_MESSAGE,
-                  io.grpc.Status.Code.ALREADY_EXISTS);
+                  ModelDBConstants.DIFFERENT_REPOSITORY_OR_COMMIT_MESSAGE, Code.ALREADY_EXISTS);
             }
             var cb = session.getCriteriaBuilder();
             CriteriaDelete<VersioningModeldbEntityMapping> delete =
@@ -2437,7 +2436,7 @@ private SimpleEntry<String, String> getS3PathAndMultipartUploadId(
       }
       if (message != null) {
         LOGGER.info(message);
-        throw new ModelDBException(message, io.grpc.Status.Code.FAILED_PRECONDITION);
+        throw new ModelDBException(message, Code.FAILED_PRECONDITION);
       }
       if (!Objects.equals(uploadId, artifactEntity.getUploadId())
           || artifactEntity.isUploadCompleted()) {

backend/src/main/java/ai/verta/modeldb/experimentRun/subtypes/ArtifactHandler.java

@@ -16,7 +16,7 @@
 import ai.verta.modeldb.experimentRun.S3KeyFunction;
 import ai.verta.modeldb.utils.ModelDBHibernateUtil;
 import ai.verta.modeldb.versioning.VersioningUtils;
-import io.grpc.Status;
+import io.grpc.Status.Code;
 import java.util.*;
 import java.util.concurrent.Executor;
 import java.util.stream.Collectors;
@@ -61,7 +61,7 @@ public ArtifactHandler(
     } else if (entityName.equals("ExperimentRunEntity")) {
       this.artifactEntityType = ArtifactPartEntity.EXP_RUN_ARTIFACT;
     } else {
-      throw new ModelDBException("Invalid entity type for ArtifactPart", Status.Code.INTERNAL);
+      throw new ModelDBException("Invalid entity type for ArtifactPart", Code.INTERNAL);
     }
   }
 
@@ -167,7 +167,7 @@ private AbstractMap.SimpleEntry<String, String> getS3PathAndMultipartUploadId(
       }
       if (message != null) {
         LOGGER.info(message);
-        throw new ModelDBException(message, io.grpc.Status.Code.FAILED_PRECONDITION);
+        throw new ModelDBException(message, Code.FAILED_PRECONDITION);
       }
       if (!Objects.equals(uploadId, artifactEntity.getUploadId())
           || artifactEntity.isUploadCompleted()) {

backend/src/main/java/ai/verta/modeldb/metadata/MetadataServiceImpl.java

@@ -70,7 +70,7 @@ public void updateLabels(
       if (request.getId() == null
           || (request.getId().getIntId() == 0 && request.getId().getStringId().isEmpty())) {
         throw new ModelDBException(
-            "Invalid parameter set in AddLabelsRequest.Id", io.grpc.Status.Code.INVALID_ARGUMENT);
+            "Invalid parameter set in AddLabelsRequest.Id", Code.INVALID_ARGUMENT);
       }
 
       boolean status = metadataDAO.updateLabels(request.getId(), request.getLabelsList());