modeldb-main/backend/server/src/main/java/ai/verta/modeldb/versioning/BlobDAORdbImpl.java
package ai.verta.modeldb.versioning;

import static ai.verta.modeldb.ModelDBConstants.DEFAULT_VERSIONING_BLOB_LOCATION;
import static java.util.stream.Collectors.toMap;

import ai.verta.common.ArtifactPart;
import ai.verta.common.KeyValue;
import ai.verta.modeldb.*;
import ai.verta.modeldb.authservice.MDBRoleService;
import ai.verta.modeldb.common.artifactStore.ArtifactStoreDAO;
import ai.verta.modeldb.common.authservice.UACApisUtil;
import ai.verta.modeldb.common.exceptions.ModelDBException;
import ai.verta.modeldb.cron_jobs.DeleteEntitiesCron;
import ai.verta.modeldb.entities.ArtifactPartEntity;
import ai.verta.modeldb.entities.AttributeEntity;
import ai.verta.modeldb.entities.UploadStatusEntity;
import ai.verta.modeldb.entities.versioning.CommitEntity;
import ai.verta.modeldb.entities.versioning.InternalFolderElementEntity;
import ai.verta.modeldb.entities.versioning.RepositoryEntity;
import ai.verta.modeldb.experimentRun.CommitMultipartFunction;
import ai.verta.modeldb.experimentRun.S3KeyFunction;
import ai.verta.modeldb.metadata.IDTypeEnum;
import ai.verta.modeldb.metadata.IdentificationType;
import ai.verta.modeldb.metadata.MetadataDAO;
import ai.verta.modeldb.utils.ModelDBHibernateUtil;
import ai.verta.modeldb.utils.ModelDBUtils;
import ai.verta.modeldb.versioning.DiffStatusEnum.DiffStatus;
import ai.verta.modeldb.versioning.FindRepositoriesBlobs.Response;
import ai.verta.modeldb.versioning.autogenerated._public.modeldb.versioning.model.*;
import ai.verta.modeldb.versioning.blob.container.BlobContainer;
import ai.verta.modeldb.versioning.blob.container.DatasetContainer;
import ai.verta.modeldb.versioning.blob.diff.ConflictGenerator;
import ai.verta.modeldb.versioning.blob.diff.DiffComputer;
import ai.verta.modeldb.versioning.blob.diff.DiffMerger;
import ai.verta.modeldb.versioning.blob.diff.TypeChecker;
import ai.verta.modeldb.versioning.blob.factory.BlobFactory;
import ai.verta.uac.Workspace;
import com.amazonaws.services.s3.model.PartETag;
import com.google.protobuf.ProtocolStringList;
import com.google.rpc.Code;
import java.security.NoSuchAlgorithmException;
import java.util.AbstractMap;
import java.util.AbstractMap.SimpleEntry;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.Date;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.hibernate.LockMode;
import org.hibernate.Session;
import org.hibernate.query.Query;

public class BlobDAORdbImpl implements BlobDAO {

  private static final Logger LOGGER = LogManager.getLogger(BlobDAORdbImpl.class);
  private static final ModelDBHibernateUtil modelDBHibernateUtil =
      ModelDBHibernateUtil.getInstance();

  public static final String TREE = "TREE";
  private static final String FOLDER_HASH_QUERY_PARAM = "folderHash";
  private static final String INTERNAL_PATH_QUERY_PARAM = "internalPath";
  private static final String COMPUTE_SHA_QUERY_PARAM = "computeSha";
  private final UACApisUtil uacApisUtil;
  private final MDBRoleService mdbRoleService;

  public BlobDAORdbImpl(UACApisUtil uacApisUtil, MDBRoleService mdbRoleService) {
    this.uacApisUtil = uacApisUtil;
    this.mdbRoleService = mdbRoleService;
  }

  /**
   * Goes through each BlobExpanded creating TREE/BLOB node top down and computing SHA bottom up
   * there is a rootSHA which holds one TREE node of each BlobExpanded
   *
   * @throws ModelDBException
   */
  @Override
  public String setBlobs(Session session, List<BlobContainer> blobContainers, FileHasher fileHasher)
      throws NoSuchAlgorithmException, ModelDBException {
    var rootTree = new TreeElem();
    Set<String> blobHashes = new HashSet<>();
    for (BlobContainer blobContainer : blobContainers) {
      // should save each blob during one session to avoid recurring entities ids
      blobContainer.process(session, rootTree, fileHasher, blobHashes);
    }
    Set<String> hashes = new HashSet<>();
    final var internalFolderElement = rootTree.saveFolders(session, fileHasher, hashes);
    return internalFolderElement.getElementSha();
  }

  @Override
  public void setBlobsAttributes(
      Session session,
      Long repoId,
      String commitHash,
      List<BlobContainer> blobContainers,
      boolean addAttribute)
      throws ModelDBException {
    for (BlobContainer blobContainer : blobContainers) {
      // should save attributes of each blob during one session to avoid recurring entities ids
      blobContainer.processAttribute(session, repoId, commitHash, addAttribute);
    }
  }

  @Override
  public DatasetVersion addUpdateDatasetVersionAttributes(
      RepositoryDAO repositoryDAO,
      CommitDAO commitDAO,
      MetadataDAO metadataDAO,
      Long repoId,
      String commitHash,
      List<KeyValue> attributes,
      boolean addAttribute)
      throws ModelDBException {
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {
      RepositoryEntity repositoryEntity;

      var repositoryIdentification = RepositoryIdentification.newBuilder();
      if (repoId == null) {
        var commitEntity = session.get(CommitEntity.class, commitHash);
        if (commitEntity == null) {
          throw new ModelDBException(
              ModelDBMessages.DATASET_VERSION_NOT_FOUND_ERROR, Code.NOT_FOUND);
        }

        if (commitEntity.getRepository() != null && commitEntity.getRepository().size() > 1) {
          throw new ModelDBException(
              String.format(
                  "DatasetVersion '%s' associated with multiple datasets",
                  commitEntity.getCommit_hash()),
              Code.INTERNAL);
        }
        assert commitEntity.getRepository() != null;
        Long newRepoId = new ArrayList<>(commitEntity.getRepository()).get(0).getId();
        repositoryIdentification.setRepoId(newRepoId);
      } else {
        repositoryIdentification.setRepoId(repoId);
      }
      repositoryEntity =
          repositoryDAO.getProtectedRepositoryById(repositoryIdentification.build(), true);
      addUpdateBlobAttributes(commitDAO, repositoryEntity, commitHash, attributes, addAttribute);
      return convertToDatasetVersion(
          repositoryDAO, metadataDAO, repositoryEntity, commitHash, true);
    } catch (Exception ex) {
      if (ModelDBUtils.needToRetry(ex)) {
        return addUpdateDatasetVersionAttributes(
            repositoryDAO, commitDAO, metadataDAO, repoId, commitHash, attributes, addAttribute);
      } else {
        throw ex;
      }
    }
  }

  @Override
  public void addUpdateBlobAttributes(
      CommitDAO commitDAO,
      RepositoryEntity repositoryEntity,
      String commitHash,
      List<KeyValue> attributes,
      boolean addAttribute)
      throws ModelDBException {
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {
      var blobBuilder = Blob.newBuilder();
      List<String> locations =
          Collections.singletonList(ModelDBConstants.DEFAULT_VERSIONING_BLOB_LOCATION);
      List<BlobContainer> blobList =
          Collections.singletonList(
              BlobContainer.create(
                  BlobExpanded.newBuilder()
                      .addAllLocation(locations)
                      .setBlob(blobBuilder.setDataset(DatasetBlob.newBuilder().build()).build())
                      .addAllAttributes(attributes)
                      .build()));

      session.beginTransaction();
      var commitEntity =
          commitDAO.getCommitEntity(session, commitHash, (session1 -> repositoryEntity));
      session.lock(commitEntity, LockMode.PESSIMISTIC_WRITE);
      setBlobsAttributes(
          session, repositoryEntity.getId(), commitEntity.getCommit_hash(), blobList, addAttribute);
      commitEntity.setDate_updated(new Date().getTime());
      commitEntity.increaseVersionNumber();
      session.update(commitEntity);
      session.getTransaction().commit();
    } catch (Exception ex) {
      if (ModelDBUtils.needToRetry(ex)) {
        addUpdateBlobAttributes(commitDAO, repositoryEntity, commitHash, attributes, addAttribute);
      } else {
        throw ex;
      }
    }
  }

  @Override
  public DatasetVersion deleteDatasetVersionAttributes(
      RepositoryDAO repositoryDAO,
      CommitDAO commitDAO,
      MetadataDAO metadataDAO,
      Long repoId,
      String commitHash,
      List<String> attributesKeys,
      List<String> location,
      boolean deleteAll)
      throws ModelDBException {
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {
      RepositoryEntity repositoryEntity;

      var repositoryIdentification = RepositoryIdentification.newBuilder();
      if (repoId == null) {
        var commitEntity = session.get(CommitEntity.class, commitHash);
        if (commitEntity == null) {
          throw new ModelDBException(
              ModelDBMessages.DATASET_VERSION_NOT_FOUND_ERROR, Code.NOT_FOUND);
        }

        if (commitEntity.getRepository() != null && commitEntity.getRepository().size() > 1) {
          throw new ModelDBException(
              String.format(
                  "DatasetVersion '%s' associated with multiple datasets",
                  commitEntity.getCommit_hash()),
              Code.INTERNAL);
        }
        assert commitEntity.getRepository() != null;
        Long newRepoId = new ArrayList<>(commitEntity.getRepository()).get(0).getId();
        repositoryIdentification.setRepoId(newRepoId);
      } else {
        repositoryIdentification.setRepoId(repoId);
      }
      repositoryEntity =
          repositoryDAO.getProtectedRepositoryById(repositoryIdentification.build(), true);

      deleteBlobAttributes(
          commitDAO, repositoryEntity, commitHash, attributesKeys, location, deleteAll);
      return convertToDatasetVersion(
          repositoryDAO, metadataDAO, repositoryEntity, commitHash, true);
    } catch (Exception ex) {
      if (ModelDBUtils.needToRetry(ex)) {
        return deleteDatasetVersionAttributes(
            repositoryDAO,
            commitDAO,
            metadataDAO,
            repoId,
            commitHash,
            attributesKeys,
            location,
            deleteAll);
      } else {
        throw ex;
      }
    }
  }

  @Override
  public void deleteBlobAttributes(
      CommitDAO commitDAO,
      RepositoryEntity repositoryEntity,
      String commitHash,
      List<String> attributesKeys,
      List<String> location,
      boolean deleteAll)
      throws ModelDBException {
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {
      session.beginTransaction();
      var commitEntity =
          commitDAO.getCommitEntity(session, commitHash, (session1 -> repositoryEntity));
      session.lock(commitEntity, LockMode.PESSIMISTIC_WRITE);
      if (deleteAll) {
        String entityHash =
            VersioningUtils.getVersioningCompositeId(
                repositoryEntity.getId(), commitEntity.getCommit_hash(), location);
        DeleteEntitiesCron.deleteAttribute(session, entityHash);
      } else {
        List<AttributeEntity> existingAttributes =
            VersioningUtils.getAttributeEntities(
                session, repositoryEntity.getId(), commitEntity.getCommit_hash(), location, null);
        for (String removeAttrKey : attributesKeys) {
          for (AttributeEntity existingAttribute : existingAttributes) {
            if (existingAttribute.getKey().equals(removeAttrKey)) {
              session.lock(existingAttribute, LockMode.PESSIMISTIC_WRITE);
              session.delete(existingAttribute);
            }
          }
        }
      }
      commitEntity.setDate_updated(new Date().getTime());
      commitEntity.increaseVersionNumber();
      session.update(commitEntity);
      session.getTransaction().commit();
    } catch (Exception ex) {
      if (ModelDBUtils.needToRetry(ex)) {
        deleteBlobAttributes(
            commitDAO, repositoryEntity, commitHash, attributesKeys, location, deleteAll);
      } else {
        throw ex;
      }
    }
  }

  @Override
  public List<KeyValue> getDatasetVersionAttributes(
      RepositoryDAO repositoryDAO,
      CommitDAO commitDAO,
      Long repoId,
      String commitHash,
      List<String> location,
      List<String> attributeKeysList)
      throws ModelDBException {
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {
      CommitEntity commitEntity = null;
      RepositoryEntity repositoryEntity;

      var repositoryIdentification = RepositoryIdentification.newBuilder();
      if (repoId == null) {
        commitEntity = session.get(CommitEntity.class, commitHash);
        if (commitEntity == null) {
          throw new ModelDBException(
              ModelDBMessages.DATASET_VERSION_NOT_FOUND_ERROR, Code.NOT_FOUND);
        }

        if (commitEntity.getRepository() != null && commitEntity.getRepository().size() > 1) {
          throw new ModelDBException(
              String.format(
                  "DatasetVersion '%s' associated with multiple datasets",
                  commitEntity.getCommit_hash()),
              Code.INTERNAL);
        }
        assert commitEntity.getRepository() != null;
        Long newRepoId = new ArrayList<>(commitEntity.getRepository()).get(0).getId();
        repositoryIdentification.setRepoId(newRepoId);
      } else {
        repositoryIdentification.setRepoId(repoId);
      }
      repositoryEntity =
          repositoryDAO.getProtectedRepositoryById(repositoryIdentification.build(), true);

      if (commitEntity == null) {
        commitEntity =
            commitDAO.getCommitEntity(session, commitHash, (session1 -> repositoryEntity));
      }

      return getBlobAttributes(
          repositoryEntity.getId(), commitEntity.getCommit_hash(), location, attributeKeysList);
    } catch (Exception ex) {
      if (ModelDBUtils.needToRetry(ex)) {
        return getDatasetVersionAttributes(
            repositoryDAO, commitDAO, repoId, commitHash, location, attributeKeysList);
      } else {
        throw ex;
      }
    }
  }

  @Override
  public List<KeyValue> getBlobAttributes(
      Long repositoryId, String commitHash, List<String> location, List<String> attributeKeysList)
      throws ModelDBException {
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {
      return VersioningUtils.getAttributes(
          session, repositoryId, commitHash, location, attributeKeysList);
    } catch (Exception ex) {
      if (ModelDBUtils.needToRetry(ex)) {
        return getBlobAttributes(repositoryId, commitHash, location, attributeKeysList);
      } else {
        throw ex;
      }
    }
  }

  private Blob getBlob(Session session, InternalFolderElementEntity folderElementEntity)
      throws ModelDBException {
    return BlobFactory.create(folderElementEntity).getBlob(session);
  }

  private Folder getFolder(Session session, String commitSha, String folderSha) {
    var query =
        session.createQuery("From InternalFolderElementEntity where folder_hash = :folder_hash");
    query.setParameter("folder_hash", folderSha);
    Optional result =
        query.list().stream()
            .map(
                d -> {
                  InternalFolderElementEntity entity = (InternalFolderElementEntity) d;
                  var folder = Folder.newBuilder();
                  FolderElement.Builder folderElement =
                      FolderElement.newBuilder().setElementName(entity.getElement_name());

                  if (entity.getElement_type().equals(TREE)) {
                    folder.addSubFolders(folderElement);
                  } else {
                    folder.addBlobs(folderElement);
                  }
                  return folder.build();
                })
            .reduce((a, b) -> ((Folder) a).toBuilder().mergeFrom((Folder) b).build());

    if (result.isPresent()) {
      return (Folder) result.get();
    } else {
      return null;
    }
  }

  // TODO : check if there is a way to optimize on the calls to data base.
  // We should fetch data  in a single query.
  @Override
  public GetCommitComponentRequest.Response getCommitComponent(
      RepositoryFunction repositoryFunction, String commitHash, ProtocolStringList locationList)
      throws ModelDBException {
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {
      RepositoryEntity repository = repositoryFunction.apply(session);
      CommitEntity commit = session.get(CommitEntity.class, commitHash);
      return getCommitComponent(session, repository.getId(), commit, locationList);
    } catch (Exception ex) {
      if (ModelDBUtils.needToRetry(ex)) {
        return getCommitComponent(repositoryFunction, commitHash, locationList);
      } else {
        throw ex;
      }
    }
  }

  private GetCommitComponentRequest.Response getCommitComponent(
      Session session, Long repoId, CommitEntity commit, List<String> locationList)
      throws ModelDBException {
    if (commit == null) {
      throw new ModelDBException("No such commit", Code.NOT_FOUND);
    }

    if (!VersioningUtils.commitRepositoryMappingExists(session, commit.getCommit_hash(), repoId)) {
      throw new ModelDBException("No such commit found in the repository", Code.NOT_FOUND);
    }

    String folderHash = commit.getRootSha();
    if (locationList.isEmpty()) { // getting root
      var folder = getFolder(session, commit.getCommit_hash(), folderHash);
      if (folder == null) { // root is empty
        return GetCommitComponentRequest.Response.newBuilder().build();
      }
      return GetCommitComponentRequest.Response.newBuilder().setFolder(folder).build();
    }
    for (var index = 0; index < locationList.size(); index++) {
      String folderLocation = locationList.get(index);
      var folderQueryHQL =
          "From InternalFolderElementEntity parentIfe WHERE parentIfe.element_name = :location AND parentIfe.folder_hash = :folderHash";
      Query<InternalFolderElementEntity> fetchTreeQuery = session.createQuery(folderQueryHQL);
      fetchTreeQuery.setParameter("location", folderLocation);
      fetchTreeQuery.setParameter(FOLDER_HASH_QUERY_PARAM, folderHash);
      InternalFolderElementEntity elementEntity = fetchTreeQuery.uniqueResult();

      if (elementEntity == null) {
        LOGGER.warn(
            "No such folder found : {}. Failed at index {} looking for {}",
            folderLocation,
            index,
            folderLocation);
        throw new ModelDBException(
            String.format("No such folder found : %s", folderLocation), Code.NOT_FOUND);
      }
      if (elementEntity.getElement_type().equals(TREE)) {
        folderHash = elementEntity.getElement_sha();
        if (index == locationList.size() - 1) {
          var folder = getFolder(session, commit.getCommit_hash(), folderHash);
          if (folder == null) { // folder is empty
            return GetCommitComponentRequest.Response.newBuilder().build();
          }
          return GetCommitComponentRequest.Response.newBuilder().setFolder(folder).build();
        }
      } else {
        if (index == locationList.size() - 1) {
          var blob = getBlob(session, elementEntity);
          List<KeyValue> attributeEntities =
              VersioningUtils.getAttributes(
                  session, repoId, commit.getCommit_hash(), locationList, null);
          return GetCommitComponentRequest.Response.newBuilder()
              .setBlob(blob)
              .addAllAttributes(attributeEntities)
              .build();
        } else {
          throw new ModelDBException(
              String.format("No such folder found : %s", locationList.get(index + 1)),
              Code.NOT_FOUND);
        }
      }
    }
    throw new ModelDBException("Unexpected logic issue found when fetching blobs", Code.UNKNOWN);
  }

  @Override
  public DatasetVersion convertToDatasetVersion(
      RepositoryDAO repositoryDAO,
      MetadataDAO metadataDAO,
      RepositoryEntity repositoryEntity,
      String commitHash,
      boolean checkWrite)
      throws ModelDBException {
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {
      if (repositoryEntity == null) {
        repositoryEntity =
            VersioningUtils.getDatasetRepositoryEntity(
                session, repositoryDAO, null, commitHash, checkWrite);
      }
      if (!repositoryEntity.isDataset()) {
        throw new ModelDBException(
            "Repository should be created from Dataset to add Dataset Version to it",
            Code.INVALID_ARGUMENT);
      }
      CommitEntity commit = session.get(CommitEntity.class, commitHash);
      List<String> locationList = Collections.singletonList(DEFAULT_VERSIONING_BLOB_LOCATION);
      var getComponentResponse =
          getCommitComponent(session, repositoryEntity.getId(), commit, locationList);
      if (getComponentResponse.hasBlob()) {
        var blob = getComponentResponse.getBlob();
        var datasetVersionBuilder = DatasetVersion.newBuilder();
        datasetVersionBuilder.setId(commit.getCommit_hash());
        if (commit.getParent_commits().size() != 0) {
          datasetVersionBuilder.setParentId(commit.getParent_commits().get(0).getCommit_hash());
        }
        datasetVersionBuilder.setDatasetId(String.valueOf(repositoryEntity.getId()));
        datasetVersionBuilder.setTimeLogged(commit.getDate_created());
        datasetVersionBuilder.setTimeUpdated(commit.getDate_updated());

        String compositeId =
            VersioningUtils.getVersioningCompositeId(
                repositoryEntity.getId(), commitHash, locationList);
        String blobDescription =
            metadataDAO.getProperty(
                session,
                IdentificationType.newBuilder()
                    .setIdType(IDTypeEnum.IDType.VERSIONING_REPO_COMMIT_BLOB)
                    .setStringId(compositeId)
                    .build(),
                "description");
        if (blobDescription != null) {
          datasetVersionBuilder.setDescription(blobDescription);
        }
        List<String> labels =
            metadataDAO.getLabels(
                session,
                IdentificationType.newBuilder()
                    .setIdType(IDTypeEnum.IDType.VERSIONING_REPO_COMMIT_BLOB)
                    .setStringId(compositeId)
                    .build());
        if (labels.size() > 0) {
          datasetVersionBuilder.addAllTags(labels);
        }
        String version =
            metadataDAO.getProperty(
                session,
                IdentificationType.newBuilder()
                    .setIdType(IDTypeEnum.IDType.VERSIONING_REPO_COMMIT_BLOB)
                    .setStringId(compositeId)
                    .build(),
                ModelDBConstants.VERSION);
        if (version != null) {
          datasetVersionBuilder.setVersion(Long.parseLong(version));
        }
        datasetVersionBuilder.addAllAttributes(getComponentResponse.getAttributesList());
        datasetVersionBuilder.setOwner(commit.getAuthor());
        datasetVersionBuilder.setVersionNumber(commit.getVersion_number());
        DatasetBlob dataset = blob.getDataset();
        var builderPathDatasetVersion = PathDatasetVersionInfo.newBuilder();
        List<DatasetPartInfo> components;
        if (dataset.hasPath()) {
          builderPathDatasetVersion.setLocationType(
              PathLocationTypeEnum.PathLocationType.LOCAL_FILE_SYSTEM);
          components =
              dataset.getPath().getComponentsList().stream()
                  .map(this::getPathInfo)
                  .collect(Collectors.toList());
        } else if (dataset.hasS3()) {
          builderPathDatasetVersion.setLocationType(
              PathLocationTypeEnum.PathLocationType.S3_FILE_SYSTEM);
          components =
              dataset.getS3().getComponentsList().stream()
                  .map(S3DatasetComponentBlob::getPath)
                  .map(this::getPathInfo)
                  .collect(Collectors.toList());
        } else if (dataset.hasQuery()) {
          components = Collections.emptyList();
          LOGGER.info("Found query dataset. Skipping populating datasetinfo");
        } else {
          var errorMessage = "Unknown blob type found while converting Blob to DatasetVersion";
          LOGGER.warn(errorMessage);
          throw new ModelDBException(errorMessage);
        }
        Optional<Long> sum = components.stream().map(DatasetPartInfo::getSize).reduce(Long::sum);
        sum.ifPresent(builderPathDatasetVersion::setSize);
        datasetVersionBuilder.setPathDatasetVersionInfo(
            builderPathDatasetVersion.addAllDatasetPartInfos(components));
        datasetVersionBuilder.setDatasetBlob(dataset);
        return datasetVersionBuilder.build();
      } else {
        throw new ModelDBException("No such blob found", Code.NOT_FOUND);
      }
    } catch (Exception ex) {
      if (ModelDBUtils.needToRetry(ex)) {
        return convertToDatasetVersion(
            repositoryDAO, metadataDAO, repositoryEntity, commitHash, checkWrite);
      } else {
        throw ex;
      }
    }
  }

  private DatasetPartInfo getPathInfo(PathDatasetComponentBlob pathDatasetComponentBlob) {
    return DatasetPartInfo.newBuilder()
        .setSize(pathDatasetComponentBlob.getSize())
        .setLastModifiedAtSource(pathDatasetComponentBlob.getLastModifiedAtSource())
        .setChecksum(pathDatasetComponentBlob.getMd5())
        .setPath(pathDatasetComponentBlob.getPath())
        .build();
  }

  /**
   * get the Folder Element pointed to by the parentFolderHash and elementName
   *
   * @param session
   * @param parentFolderHash : folder hash of the parent
   * @param elementName : element name of the element to be fetched
   * @return {@link List<InternalFolderElementEntity>}
   */
  private List<InternalFolderElementEntity> getFolderElement(
      Session session, String parentFolderHash, String elementName) {
    StringBuilder folderQueryHQLBuilder =
        new StringBuilder("From ")
            .append(InternalFolderElementEntity.class.getSimpleName())
            .append(" parentIfe WHERE parentIfe.folder_hash = :folderHash ");

    if (elementName != null && !elementName.isEmpty()) {
      folderQueryHQLBuilder.append("AND parentIfe.element_name = :elementName");
    }

    Query<InternalFolderElementEntity> fetchTreeQuery =
        session.createQuery(folderQueryHQLBuilder.toString());
    fetchTreeQuery.setParameter(FOLDER_HASH_QUERY_PARAM, parentFolderHash);
    if (elementName != null && !elementName.isEmpty()) {
      fetchTreeQuery.setParameter("elementName", elementName);
    }
    return fetchTreeQuery.list();
  }

  boolean childContains(Set<?> list, Set<?> sublist) {
    return Collections.indexOfSubList(new LinkedList<>(list), new LinkedList<>(sublist)) != -1;
  }

  private Map<String, Map.Entry<BlobExpanded, String>> getChildFolderBlobMap(
      Session session,
      List<String> requestedLocation,
      Set<String> parentLocation,
      String parentFolderHash,
      List<BlobType> blobTypeList)
      throws ModelDBException {
    var folderQueryHQL =
        "From InternalFolderElementEntity parentIfe WHERE parentIfe.folder_hash = :folderHash";
    Query<InternalFolderElementEntity> fetchTreeQuery = session.createQuery(folderQueryHQL);
    fetchTreeQuery.setParameter(FOLDER_HASH_QUERY_PARAM, parentFolderHash);
    List<InternalFolderElementEntity> childElementFolders = fetchTreeQuery.list();

    Map<String, Map.Entry<BlobExpanded, String>> childBlobExpandedMap = new LinkedHashMap<>();
    for (InternalFolderElementEntity childElementFolder : childElementFolders) {
      Set<String> childLocation = new LinkedHashSet<>(parentLocation);
      childLocation.add(childElementFolder.getElement_name());
      if (childContains(new LinkedHashSet<>(requestedLocation), childLocation)
          || childLocation.containsAll(requestedLocation)) {
        if (childElementFolder.getElement_type().equals(TREE)) {
          childBlobExpandedMap.putAll(
              getChildFolderBlobMap(
                  session,
                  requestedLocation,
                  childLocation,
                  childElementFolder.getElement_sha(),
                  blobTypeList));
        } else {
          if (parentLocation.containsAll(requestedLocation)
              || childLocation.containsAll(requestedLocation)) {
            var blob = getBlob(session, childElementFolder);
            if (blobTypeList != null && !blobTypeList.isEmpty()) {
              if (blobTypeExistsInList(blobTypeList, blob.getContentCase())) {
                setBlobInBlobExpandMap(
                    parentLocation, childBlobExpandedMap, childElementFolder, blob);
              }
            } else {
              setBlobInBlobExpandMap(
                  parentLocation, childBlobExpandedMap, childElementFolder, blob);
            }
          }
        }
      } else {
        if (parentLocation.containsAll(requestedLocation)) {
          var blob = getBlob(session, childElementFolder);
          if (blobTypeList != null && !blobTypeList.isEmpty()) {
            if (blobTypeExistsInList(blobTypeList, blob.getContentCase())) {
              setBlobInBlobExpandMap(
                  parentLocation, childBlobExpandedMap, childElementFolder, blob);
            }
          } else {
            setBlobInBlobExpandMap(parentLocation, childBlobExpandedMap, childElementFolder, blob);
          }
        }
      }
    }
    return childBlobExpandedMap;
  }

  private void setBlobInBlobExpandMap(
      Set<String> parentLocation,
      Map<String, Entry<BlobExpanded, String>> blobExpandedMap,
      InternalFolderElementEntity elementFolder,
      Blob blob) {
    BlobExpanded blobExpanded;
    if (parentLocation.size() == 1
        && new ArrayList<>(parentLocation).get(0).equals(elementFolder.getElement_name())) {
      blobExpanded =
          BlobExpanded.newBuilder()
              .addLocation(elementFolder.getElement_name())
              .setBlob(blob)
              .build();
    } else {
      blobExpanded =
          BlobExpanded.newBuilder()
              .addAllLocation(parentLocation)
              .addLocation(elementFolder.getElement_name())
              .setBlob(blob)
              .build();
    }
    blobExpandedMap.put(
        getStringFromLocationList(blobExpanded.getLocationList()),
        new SimpleEntry<>(blobExpanded, elementFolder.getElement_sha()));
  }

  /**
   * Given a folderHash and a location list, collects all the blobs along the location list and
   * returns them with their location as set
   *
   * @param session
   * @param folderHash : the base folder to start the search for location list
   * @param locationList : list of trees and psossibly terminating with blob
   * @return
   * @throws ModelDBException
   */
  @Override
  public Map<String, BlobExpanded> getCommitBlobMap(
      Session session, String folderHash, List<String> locationList) throws ModelDBException {
    return convertToLocationBlobMap(
        getCommitBlobMapWithHash(session, folderHash, locationList, Collections.emptyList()));
  }

  private Map<String, BlobExpanded> convertToLocationBlobMap(
      Map<String, Map.Entry<BlobExpanded, String>> commitBlobMapWithHash) {
    return commitBlobMapWithHash.entrySet().stream()
        .collect(
            Collectors.toMap(
                Entry::getKey, stringEntryEntry -> stringEntryEntry.getValue().getKey()));
  }

  @Override
  public Map<String, Map.Entry<BlobExpanded, String>> getCommitBlobMapWithHash(
      Session session, String folderHash, List<String> locationList, List<BlobType> blobTypeList)
      throws ModelDBException {

    String parentLocation = locationList.size() == 0 ? null : locationList.get(0);
    List<InternalFolderElementEntity> parentFolderElementList =
        getFolderElement(session, folderHash, parentLocation);
    if (parentFolderElementList == null || parentFolderElementList.isEmpty()) {
      if (parentLocation
          != null) { // = null mainly is supporting the call on init commit which is an empty commit
        throw new ModelDBException(
            String.format("No such folder found : %s", parentLocation), Code.NOT_FOUND);
      }
    }

    Map<String, Map.Entry<BlobExpanded, String>> finalLocationBlobMap = new LinkedHashMap<>();
    for (InternalFolderElementEntity parentFolderElement : parentFolderElementList) {
      if (!parentFolderElement.getElement_type().equals(TREE)) {
        var blob = getBlob(session, parentFolderElement);
        if (blobTypeList != null && !blobTypeList.isEmpty()) {
          if (blobTypeExistsInList(blobTypeList, blob.getContentCase())) {
            setBlobInBlobExpandMap(
                Collections.singleton(parentFolderElement.getElement_name()),
                finalLocationBlobMap,
                parentFolderElement,
                blob);
          }
        } else {
          setBlobInBlobExpandMap(
              Collections.singleton(parentFolderElement.getElement_name()),
              finalLocationBlobMap,
              parentFolderElement,
              blob);
        }
      } else {
        // if this is tree, search further
        Set<String> location = new LinkedHashSet<>();
        Map<String, Map.Entry<BlobExpanded, String>> locationBlobList =
            getChildFolderBlobMap(session, locationList, location, folderHash, blobTypeList);
        finalLocationBlobMap.putAll(locationBlobList);
      }
    }

    Comparator<Map.Entry<String, Map.Entry<BlobExpanded, String>>> locationComparator =
        Comparator.comparing(
            (Map.Entry<String, Map.Entry<BlobExpanded, String>> o) -> o.getKey().replace("#", ""));

    finalLocationBlobMap =
        finalLocationBlobMap.entrySet().stream()
            .sorted(locationComparator)
            .collect(
                toMap(Map.Entry::getKey, Map.Entry::getValue, (e1, e2) -> e2, LinkedHashMap::new));

    return finalLocationBlobMap;
  }

  @Override
  public ListCommitBlobsRequest.Response getCommitBlobsList(
      RepositoryFunction repositoryFunction, String commitHash, List<String> locationList)
      throws ModelDBException {
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {

      CommitEntity commit = session.get(CommitEntity.class, commitHash);
      if (commit == null) {
        throw new ModelDBException("No such commit", Code.NOT_FOUND);
      }

      RepositoryEntity repository = repositoryFunction.apply(session);
      if (!VersioningUtils.commitRepositoryMappingExists(session, commitHash, repository.getId())) {
        throw new ModelDBException("No such commit found in the repository", Code.NOT_FOUND);
      }
      Map<String, BlobExpanded> locationBlobMap =
          getCommitBlobMap(session, commit.getRootSha(), locationList);
      Set<BlobExpanded> blobExpandedSet = new HashSet<>();
      for (String location : locationBlobMap.keySet()) {
        List<KeyValue> attributes =
            VersioningUtils.getAttributes(
                session,
                repository.getId(),
                commit.getCommit_hash(),
                Collections.singletonList(location),
                null);
        var blobExpanded = locationBlobMap.get(location);
        blobExpanded = blobExpanded.toBuilder().addAllAttributes(attributes).build();
        blobExpandedSet.add(blobExpanded);
      }
      return ListCommitBlobsRequest.Response.newBuilder().addAllBlobs(blobExpandedSet).build();
    } catch (Exception ex) {
      if (ModelDBUtils.needToRetry(ex)) {
        return getCommitBlobsList(repositoryFunction, commitHash, locationList);
      } else {
        throw ex;
      }
    }
  }

  @Override
  public ComputeRepositoryDiffRequest.Response computeRepositoryDiff(
      RepositoryDAO repositoryDAO, ComputeRepositoryDiffRequest request) throws ModelDBException {
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {

      // validating request
      validateDiffMergeRequest(
          request.getBranchA(), request.getBranchB(), request.getCommitA(), request.getCommitB());

      var repositoryEntity = repositoryDAO.getRepositoryById(session, request.getRepositoryId());

      CommitEntity internalCommitA = null;
      CommitEntity internalCommitB = null;
      String errorNameA = null;
      String errorNameB = null;
      if (!request.getCommitA().isEmpty()) {
        errorNameA = request.getCommitA();
        internalCommitA =
            session.get(CommitEntity.class, request.getCommitA(), LockMode.PESSIMISTIC_WRITE);
      }

      if (!request.getCommitB().isEmpty()) {
        errorNameB = request.getCommitB();
        internalCommitB =
            session.get(CommitEntity.class, request.getCommitB(), LockMode.PESSIMISTIC_WRITE);
      }

      if (!request.getBranchA().isEmpty()) {
        errorNameA = request.getBranchA();
        var branchAEntity =
            repositoryDAO.getBranchEntity(session, repositoryEntity.getId(), request.getBranchA());
        internalCommitA =
            session.get(
                CommitEntity.class, branchAEntity.getCommit_hash(), LockMode.PESSIMISTIC_WRITE);
      }

      if (!request.getBranchB().isEmpty()) {
        errorNameB = request.getBranchB();
        var branchBEntity =
            repositoryDAO.getBranchEntity(session, repositoryEntity.getId(), request.getBranchB());
        internalCommitB =
            session.get(
                CommitEntity.class, branchBEntity.getCommit_hash(), LockMode.PESSIMISTIC_WRITE);
      }

      if (internalCommitA == null) {
        throw new ModelDBException(
            String.format("No such commit OR branch found : %s", errorNameA), Code.NOT_FOUND);
      }

      if (internalCommitB == null) {
        throw new ModelDBException(
            String.format("No such commit OR branch found : %s", errorNameB), Code.NOT_FOUND);
      }

      if (!VersioningUtils.commitRepositoryMappingExists(
          session, internalCommitA.getCommit_hash(), repositoryEntity.getId())) {
        throw new ModelDBException(
            String.format(
                "No such commit found in the repository : %s", internalCommitA.getCommit_hash()),
            Code.NOT_FOUND);
      }

      if (!VersioningUtils.commitRepositoryMappingExists(
          session, internalCommitB.getCommit_hash(), repositoryEntity.getId())) {
        throw new ModelDBException(
            String.format(
                "No such commit found in the repository : %s", internalCommitB.getCommit_hash()),
            Code.NOT_FOUND);
      }

      if (request.getReplaceAWithCommonAncestor()) {
        internalCommitA =
            getCommonParent(
                session, internalCommitA.getCommit_hash(), internalCommitB.getCommit_hash());
      }
      // get list of blob expanded in both commit and group them in a map based on location
      Map<String, Map.Entry<BlobExpanded, String>> locationBlobsMapCommitA =
          getCommitBlobMapWithHash(
              session, internalCommitA.getRootSha(), new ArrayList<>(), Collections.emptyList());

      Map<String, Map.Entry<BlobExpanded, String>> locationBlobsMapCommitB =
          getCommitBlobMapWithHash(
              session, internalCommitB.getRootSha(), new ArrayList<>(), Collections.emptyList());

      return computeDiffFromCommitMaps(locationBlobsMapCommitA, locationBlobsMapCommitB);
    } catch (Exception ex) {
      if (ModelDBUtils.needToRetry(ex)) {
        return computeRepositoryDiff(repositoryDAO, request);
      } else {
        throw ex;
      }
    }
  }

  private ComputeRepositoryDiffRequest.Response computeDiffFromCommitMaps(
      Map<String, Map.Entry<BlobExpanded, String>> locationBlobsMapCommitA,
      Map<String, Map.Entry<BlobExpanded, String>> locationBlobsMapCommitB) {
    // Added new blob location in the CommitB, locations in
    Set<String> addedLocations = new LinkedHashSet<>(locationBlobsMapCommitB.keySet());
    addedLocations.removeAll(locationBlobsMapCommitA.keySet());
    LOGGER.debug("Added location for Diff : {}", addedLocations);

    // deleted new blob location from the CommitA
    Set<String> deletedLocations = new LinkedHashSet<>(locationBlobsMapCommitA.keySet());
    deletedLocations.removeAll(locationBlobsMapCommitB.keySet());
    LOGGER.debug("Deleted location for Diff : {}", deletedLocations);

    // get B sha -> blobs
    Map<String, Set<BlobExpanded>> blobsB = getCollectToMap(locationBlobsMapCommitB);
    // get A sha -> blobs
    Map<String, Set<BlobExpanded>> blobsA = getCollectToMap(locationBlobsMapCommitA);
    // delete blobs same with A
    for (Map.Entry<String, Set<BlobExpanded>> entry : blobsA.entrySet()) {
      Set<BlobExpanded> ent = blobsB.get(entry.getKey());
      if (ent != null) {
        ent.removeAll(entry.getValue());
      }
    }
    // get modified location -> blob
    Map<String, BlobExpanded> locationBlobsModified =
        getLocationWiseBlobExpandedMapFromCollection(
            blobsB.values().stream().flatMap(Collection::stream).collect(Collectors.toList()));
    // remove added from modified
    locationBlobsModified.keySet().removeAll(addedLocations);
    Set<String> modifiedLocations = locationBlobsModified.keySet();
    LOGGER.debug("Modified location for Diff : {}", modifiedLocations);

    List<ai.verta.modeldb.versioning.BlobDiff> addedBlobDiffList =
        getAddedBlobDiff(addedLocations, convertToLocationBlobMap(locationBlobsMapCommitB));
    List<ai.verta.modeldb.versioning.BlobDiff> deletedBlobDiffList =
        getDeletedBlobDiff(deletedLocations, convertToLocationBlobMap(locationBlobsMapCommitA));
    List<ai.verta.modeldb.versioning.BlobDiff> modifiedBlobDiffList =
        getModifiedBlobDiff(
            modifiedLocations,
            convertToLocationBlobMap(locationBlobsMapCommitA),
            convertToLocationBlobMap(locationBlobsMapCommitB));

    return ComputeRepositoryDiffRequest.Response.newBuilder()
        .addAllDiffs(addedBlobDiffList)
        .addAllDiffs(deletedBlobDiffList)
        .addAllDiffs(modifiedBlobDiffList)
        .build();
  }

  public void validateDiffMergeRequest(
      String branchA, String branchB, String commitA, String commitB) throws ModelDBException {
    var invalidParam = false;
    if (!branchA.isEmpty() && !branchB.isEmpty() && !commitA.isEmpty() && !commitB.isEmpty()) {
      invalidParam = true;
    } else if (!branchA.isEmpty() && !commitA.isEmpty()) {
      invalidParam = true;
    } else if (branchA.isEmpty() && commitA.isEmpty()) {
      invalidParam = true;
    } else if (!branchB.isEmpty() && !commitB.isEmpty()) {
      invalidParam = true;
    } else if (branchB.isEmpty() && commitB.isEmpty()) {
      invalidParam = true;
    }

    if (invalidParam) {
      throw new ModelDBException(
          "Branches and Commits both are not allowed in the request", Code.INVALID_ARGUMENT);
    }
  }

  @Override
  public MergeRepositoryCommitsRequest.Response mergeCommit(
      RepositoryDAO repositoryDAO, MergeRepositoryCommitsRequest request)
      throws ModelDBException, NoSuchAlgorithmException {
    // validating request
    LOGGER.debug("Validating MergeRepositoryCommitsRequest");
    validateDiffMergeRequest(
        request.getBranchA(),
        request.getBranchB(),
        request.getCommitShaA(),
        request.getCommitShaB());
    LOGGER.debug("MergeRepositoryCommitsRequest validated");

    Map<String, Map.Entry<BlobExpanded, String>> locationBlobsMapCommitA =
        new HashMap<String, Map.Entry<BlobExpanded, String>>();
    Map<String, Map.Entry<BlobExpanded, String>> locationBlobsMapCommitB =
        new HashMap<String, Map.Entry<BlobExpanded, String>>();
    Map<String, Map.Entry<BlobExpanded, String>> locationBlobsMapParentCommit =
        new HashMap<String, Map.Entry<BlobExpanded, String>>();
    Map<String, BlobExpanded> locationBlobsMapCommitASimple = new HashMap<String, BlobExpanded>();

    CommitEntity internalCommitA = null;
    CommitEntity internalCommitB = null;
    CommitEntity parentCommit;
    Commit parentCommitProto;
    String branchOrCommitA = null;
    String branchOrCommitB = null;
    RepositoryEntity repositoryEntity = null;
    try (var readSession = modelDBHibernateUtil.getSessionFactory().openSession()) {
      repositoryEntity = repositoryDAO.getRepositoryById(readSession, request.getRepositoryId());

      if (!request.getBranchA().isEmpty()) {
        LOGGER.debug("Branch A found in request");
        branchOrCommitA = request.getBranchA();
        var branchAEntity =
            repositoryDAO.getBranchEntity(
                readSession, repositoryEntity.getId(), request.getBranchA());
        internalCommitA =
            readSession.get(
                CommitEntity.class, branchAEntity.getCommit_hash(), LockMode.PESSIMISTIC_WRITE);
      }

      if (!request.getBranchB().isEmpty()) {
        LOGGER.debug("Branch B found in request");
        branchOrCommitB = request.getBranchB();
        var branchBEntity =
            repositoryDAO.getBranchEntity(
                readSession, repositoryEntity.getId(), request.getBranchB());
        internalCommitB =
            readSession.get(
                CommitEntity.class, branchBEntity.getCommit_hash(), LockMode.PESSIMISTIC_WRITE);
      }

      if (!request.getCommitShaA().isEmpty()) {
        LOGGER.debug("Commit A found in request");
        branchOrCommitA = request.getCommitShaA().substring(0, 7);
        internalCommitA =
            readSession.get(
                CommitEntity.class, request.getCommitShaA(), LockMode.PESSIMISTIC_WRITE);
      }

      if (!request.getCommitShaB().isEmpty()) {
        LOGGER.debug("Commit B found in request");
        branchOrCommitB = request.getCommitShaB().substring(0, 7);
        internalCommitB =
            readSession.get(
                CommitEntity.class, request.getCommitShaB(), LockMode.PESSIMISTIC_WRITE);
      }

      if (internalCommitA == null) {
        throw new ModelDBException(
            String.format("No such commit OR branch found : %s", branchOrCommitA), Code.NOT_FOUND);
      }

      if (internalCommitB == null) {
        throw new ModelDBException(
            String.format("No such commit OR branch found : %s", branchOrCommitB), Code.NOT_FOUND);
      }

      if (!VersioningUtils.commitRepositoryMappingExists(
          readSession, internalCommitA.getCommit_hash(), repositoryEntity.getId())) {
        throw new ModelDBException(
            String.format(
                "No such commit found in the repository : %s", internalCommitA.getCommit_hash()),
            Code.NOT_FOUND);
      }

      if (!VersioningUtils.commitRepositoryMappingExists(
          readSession, internalCommitB.getCommit_hash(), repositoryEntity.getId())) {
        throw new ModelDBException(
            String.format(
                "No such commit found in the repository : %s", internalCommitB.getCommit_hash()),
            Code.NOT_FOUND);
      }

      parentCommit =
          getCommonParent(
              readSession, internalCommitA.getCommit_hash(), internalCommitB.getCommit_hash());
      parentCommitProto = parentCommit.toCommitProto();
      locationBlobsMapCommitA =
          getCommitBlobMapWithHash(
              readSession,
              internalCommitA.getRootSha(),
              new ArrayList<>(),
              Collections.emptyList());

      locationBlobsMapCommitASimple = convertToLocationBlobMap(locationBlobsMapCommitA);

      locationBlobsMapCommitB =
          getCommitBlobMapWithHash(
              readSession,
              internalCommitB.getRootSha(),
              new ArrayList<>(),
              Collections.emptyList());

      locationBlobsMapParentCommit =
          getCommitBlobMapWithHash(
              readSession, parentCommit.getRootSha(), new ArrayList<>(), Collections.emptyList());
    }
    try (var writeSession = modelDBHibernateUtil.getSessionFactory().openSession()) {
      List<ai.verta.modeldb.versioning.BlobDiff> diffB =
          computeDiffFromCommitMaps(locationBlobsMapParentCommit, locationBlobsMapCommitB)
              .getDiffsList();

      HashMap<String, List<String>> conflictLocationMap = new HashMap<String, List<String>>();

      List<BlobContainer> blobContainerList =
          getBlobContainers(
              diffB.stream().map(AutogenBlobDiff::fromProto).collect(Collectors.toList()),
              locationBlobsMapCommitASimple,
              conflictLocationMap);

      if (conflictLocationMap.isEmpty()) {
        if (request.getIsDryRun()) {
          return MergeRepositoryCommitsRequest.Response.getDefaultInstance();
        }
        String mergeMessage = request.getContent().getMessage();
        List<String> parentSHAs =
            Arrays.asList(internalCommitB.getCommit_hash(), internalCommitA.getCommit_hash());
        Map<Integer, CommitEntity> parentCommits = new HashMap<>();
        parentCommits.put(0, internalCommitB);
        parentCommits.put(1, internalCommitA);
        if (mergeMessage.isEmpty()) {
          mergeMessage = "Merge " + branchOrCommitA + " into " + branchOrCommitB;
        }

        writeSession.beginTransaction();
        var commitEntity =
            getCommitEntityObj(
                writeSession,
                repositoryEntity,
                parentSHAs,
                parentCommits,
                blobContainerList,
                mergeMessage);
        writeSession.saveOrUpdate(commitEntity);
        writeSession.getTransaction().commit();
        return MergeRepositoryCommitsRequest.Response.newBuilder()
            .setCommit(commitEntity.toCommitProto())
            .build();
      } else {

        Map<String, BlobExpanded> locationBlobsMapParentCommitSimple =
            convertToLocationBlobMap(locationBlobsMapParentCommit);

        List<ai.verta.modeldb.versioning.BlobDiff> diffA =
            computeDiffFromCommitMaps(locationBlobsMapParentCommit, locationBlobsMapCommitA)
                .getDiffsList();
        List<BlobDiff> blobDiffList =
            getConflictDiff(diffA, diffB, conflictLocationMap, locationBlobsMapParentCommitSimple);
        LOGGER.debug("conflict found {}", conflictLocationMap);
        return MergeRepositoryCommitsRequest.Response.newBuilder()
            .setCommonBase(parentCommitProto)
            .addAllConflicts(blobDiffList)
            .build();
      }
    }
  }

  private CommitEntity getCommitEntityObj(
      Session writeSession,
      RepositoryEntity repositoryEntity,
      List<String> parentSHAs,
      Map<Integer, CommitEntity> parentCommits,
      List<BlobContainer> blobContainerList,
      String commitMessage)
      throws NoSuchAlgorithmException, ModelDBException {
    String rootSha;
    if (blobContainerList != null && !blobContainerList.isEmpty()) {
      rootSha = setBlobs(writeSession, blobContainerList, new FileHasher());
    } else {
      rootSha = FileHasher.getSha(new String());
    }
    long timeCreated = new Date().getTime();

    try {
      var currentLoginUserInfo = uacApisUtil.getCurrentLoginUserInfo().blockAndGet();
      String author = uacApisUtil.getVertaIdFromUserInfo(currentLoginUserInfo);
      final String commitSha =
          VersioningUtils.generateCommitSHA(
              parentSHAs, commitMessage, timeCreated, author, rootSha);

      var internalCommit =
          Commit.newBuilder()
              .setDateCreated(timeCreated)
              .setDateUpdated(timeCreated)
              .setAuthor(author)
              .setMessage(commitMessage)
              .setCommitSha(commitSha)
              .build();
      return new CommitEntity(repositoryEntity, parentCommits, internalCommit, rootSha);
    } catch (Exception e) {
      throw new RuntimeException(e);
    }
  }

  @Override
  public RevertRepositoryCommitsRequest.Response revertCommit(
      RepositoryDAO repositoryDAO, RevertRepositoryCommitsRequest request)
      throws ModelDBException, NoSuchAlgorithmException {
    // validating request
    LOGGER.debug("Validating RevertRepositoryCommitsRequest");
    if (request.getCommitToRevertSha().isEmpty()) {
      throw new ModelDBException(
          "Revert Commit SHA not found in the request", Code.INVALID_ARGUMENT);
    }
    if (request.getBaseCommitSha().isEmpty()) {
      throw new ModelDBException(
          "Revert Base Commit SHA not found in the request", Code.INVALID_ARGUMENT);
    }
    LOGGER.debug("RevertRepositoryCommitsRequest validated");

    List<BlobContainer> blobContainerList;
    RepositoryEntity repositoryEntity;
    CommitEntity baseCommitEntity;
    CommitEntity commitToRevertEntity;
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {
      repositoryEntity = repositoryDAO.getRepositoryById(session, request.getRepositoryId());

      if (!VersioningUtils.commitRepositoryMappingExists(
          session, request.getCommitToRevertSha(), repositoryEntity.getId())) {
        throw new ModelDBException(
            "No such revert commit found in the repository : " + request.getCommitToRevertSha(),
            Code.NOT_FOUND);
      }

      if (!VersioningUtils.commitRepositoryMappingExists(
          session, request.getBaseCommitSha(), repositoryEntity.getId())) {
        throw new ModelDBException(
            "No such base commit found in the repository : " + request.getCommitToRevertSha(),
            Code.NOT_FOUND);
      }

      commitToRevertEntity =
          session.get(
              CommitEntity.class, request.getCommitToRevertSha(), LockMode.PESSIMISTIC_WRITE);
      baseCommitEntity =
          session.get(CommitEntity.class, request.getBaseCommitSha(), LockMode.PESSIMISTIC_WRITE);

      if (commitToRevertEntity.getParent_commits() == null
          || commitToRevertEntity.getParent_commits().isEmpty()) {
        throw new ModelDBException(
            "No parent found for commit : " + request.getCommitToRevertSha(), Code.INTERNAL);
      }
      CommitEntity firstParentOfCommitToRevert = commitToRevertEntity.getParent_commits().get(0);

      Map<String, Map.Entry<BlobExpanded, String>> locationBlobsMapFirstParentCommit =
          getCommitBlobMapWithHash(
              session,
              firstParentOfCommitToRevert.getRootSha(),
              new ArrayList<>(),
              Collections.emptyList());

      Map<String, Map.Entry<BlobExpanded, String>> locationBlobsMapBaseCommit =
          getCommitBlobMapWithHash(
              session, baseCommitEntity.getRootSha(), new ArrayList<>(), Collections.emptyList());

      Map<String, Map.Entry<BlobExpanded, String>> locationBlobsMapCommitToRevert =
          getCommitBlobMapWithHash(
              session,
              commitToRevertEntity.getRootSha(),
              new ArrayList<>(),
              Collections.emptyList());

      List<ai.verta.modeldb.versioning.BlobDiff> commitToRevertShaBlobDiff =
          computeDiffFromCommitMaps(
                  locationBlobsMapCommitToRevert, locationBlobsMapFirstParentCommit)
              .getDiffsList();
      Map<String, BlobExpanded> locationBlobsMapBaseCommitSimple =
          convertToLocationBlobMap(locationBlobsMapBaseCommit);

      blobContainerList =
          getBlobContainers(
              commitToRevertShaBlobDiff.stream()
                  .map(AutogenBlobDiff::fromProto)
                  .collect(Collectors.toList()),
              locationBlobsMapBaseCommitSimple,
              new LinkedHashMap<>());
    } catch (Exception ex) {
      if (ModelDBUtils.needToRetry(ex)) {
        return revertCommit(repositoryDAO, request);
      } else {
        throw ex;
      }
    }

    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {
      List<String> parentSHAs = Collections.singletonList(request.getBaseCommitSha());
      Map<Integer, CommitEntity> parentCommits = new HashMap<>();
      parentCommits.put(0, baseCommitEntity);
      String revertMessage = request.getContent().getMessage();
      if (revertMessage.isEmpty()) {
        revertMessage = VersioningUtils.revertCommitMessage(commitToRevertEntity.toCommitProto());
      }

      var transaction = session.beginTransaction();
      var commitEntity =
          getCommitEntityObj(
              session,
              repositoryEntity,
              parentSHAs,
              parentCommits,
              blobContainerList,
              revertMessage);
      session.saveOrUpdate(commitEntity);
      transaction.commit();
      return RevertRepositoryCommitsRequest.Response.newBuilder()
          .setCommit(commitEntity.toCommitProto())
          .build();
    } catch (Exception ex) {
      if (ModelDBUtils.needToRetry(ex)) {
        return revertCommit(repositoryDAO, request);
      } else {
        throw ex;
      }
    }
  }

  private List<BlobDiff> getConflictDiff(
      List<BlobDiff> diffListA,
      List<BlobDiff> diffListB,
      HashMap<String, List<String>> conflictLocationMap,
      Map<String, BlobExpanded> locationBlobsMapParentCommitSimple)
      throws ModelDBException {
    List<BlobDiff> blobDiffList = new LinkedList<BlobDiff>(); // TODO sort?
    HashMap<String, List<BlobDiff>> diffMapA =
        getLocationMapDiff(diffListA, conflictLocationMap.keySet());
    HashMap<String, List<BlobDiff>> diffMapB =
        getLocationMapDiff(diffListB, conflictLocationMap.keySet());

    for (Entry<String, List<String>> entry : conflictLocationMap.entrySet()) {
      LOGGER.debug(entry.getKey());
      LOGGER.debug(entry.getValue());
      List<BlobDiff> locSpecificBlobDiffA =
          diffMapA.getOrDefault(entry.getKey(), Collections.emptyList());
      List<BlobDiff> locSpecificBlobDiffB =
          diffMapB.getOrDefault(entry.getKey(), Collections.emptyList());
      var parentBlobExpanded =
          locationBlobsMapParentCommitSimple.getOrDefault(
              entry.getKey(), BlobExpanded.newBuilder().build());
      var diffBuilder = BlobDiff.newBuilder().setStatus(DiffStatus.CONFLICTED);
      if (!locSpecificBlobDiffA.isEmpty()) {
        diffBuilder.addAllLocation(locSpecificBlobDiffA.get(0).getLocationList());
      } else {
        diffBuilder.addAllLocation(locSpecificBlobDiffB.get(0).getLocationList());
      }
      blobDiffList.addAll(
          ConflictGenerator.setConflictBlobsInDiff(
              diffBuilder.build(),
              locSpecificBlobDiffA,
              locSpecificBlobDiffB,
              parentBlobExpanded.getBlob()));
    }
    return blobDiffList;
  }

  private HashMap<String, List<BlobDiff>> getLocationMapDiff(
      List<BlobDiff> diffList, Set<String> interestSet) {
    HashMap<String, List<BlobDiff>> diffMap = new HashMap<String, List<BlobDiff>>();
    for (BlobDiff diff : diffList) {
      var locKey = getStringFromLocationList(diff.getLocationList());
      if (interestSet.contains(locKey)) {
        if (!diffMap.containsKey(locKey)) {
          diffMap.put(locKey, new LinkedList<BlobDiff>());
        }
        diffMap.get(locKey).add(diff);
      }
    }
    return diffMap;
  }

  private CommitEntity getCommonParent(Session session, String commitA, String commitB)
      throws ModelDBException {
    LOGGER.debug("Branch B found in request");
    List<CommitEntity> parentCommitA = VersioningUtils.getParentCommits(session, commitA);
    List<CommitEntity> parentCommitB = VersioningUtils.getParentCommits(session, commitB);

    CommitEntity commonParent = null;
    var itrA = 0;
    var itrB = 0;
    // TODO : this algorithm does not require all the commits to be available before starting.
    while (itrA < parentCommitA.size() && itrB < parentCommitB.size()) {
      CommitEntity candidateA = parentCommitA.get(itrA);
      CommitEntity candidateB = parentCommitB.get(itrB);
      if (candidateA.getCommit_hash().equals(candidateB.getCommit_hash())) {
        return candidateA;
      } else if (candidateA.getDate_created() > candidateB.getDate_created()) {
        itrA++;
      } else {
        itrB++;
      }
    }
    // Should never happen, since we have the initial commit
    throw new ModelDBException("Could not find base commit for merge", Code.INTERNAL);
  }

  private Map<String, Set<BlobExpanded>> getCollectToMap(
      Map<String, Entry<BlobExpanded, String>> locationBlobsMapCommit) {
    return locationBlobsMapCommit.values().stream()
        .collect(
            Collectors.toMap(
                Entry::getValue,
                entry -> new LinkedHashSet<>(Collections.singletonList(entry.getKey())),
                (m1, m2) -> {
                  LinkedHashSet<BlobExpanded> newHash = new LinkedHashSet<>(m1);
                  newHash.addAll(m2);
                  return newHash;
                },
                LinkedHashMap::new));
  }

  List<ai.verta.modeldb.versioning.BlobDiff> getAddedBlobDiff(
      Set<String> addedLocations, Map<String, BlobExpanded> locationBlobsMapCommitB) {
    return addedLocations.stream()
        .map(
            location -> {
              var blobExpanded = locationBlobsMapCommitB.get(location);
              AutogenBlobDiff diff =
                  DiffComputer.computeBlobDiff(null, fromBlobProto(blobExpanded));
              diff.setStatus(AutogenDiffStatusEnumDiffStatus.fromProto(DiffStatus.ADDED));
              diff.setLocation(blobExpanded.getLocationList());
              return diff.toProto().build();
            })
        .collect(Collectors.toList());
  }

  private AutogenBlob fromBlobProto(BlobExpanded blobExpanded) {
    return AutogenBlob.fromProto(blobExpanded.getBlob());
  }

  List<ai.verta.modeldb.versioning.BlobDiff> getDeletedBlobDiff(
      Set<String> deletedLocations, Map<String, BlobExpanded> locationBlobsMapCommitA) {
    return deletedLocations.stream()
        .map(
            location -> {
              var blobExpanded = locationBlobsMapCommitA.get(location);

              AutogenBlobDiff diff =
                  DiffComputer.computeBlobDiff(fromBlobProto(blobExpanded), null);
              diff.setStatus(AutogenDiffStatusEnumDiffStatus.fromProto(DiffStatus.DELETED));
              diff.setLocation(blobExpanded.getLocationList());
              return diff.toProto().build();
            })
        .collect(Collectors.toList());
  }

  List<ai.verta.modeldb.versioning.BlobDiff> getModifiedBlobDiff(
      Set<String> modifiedLocations,
      Map<String, BlobExpanded> locationBlobsMapCommitA,
      Map<String, BlobExpanded> locationBlobsMapCommitB) {
    return modifiedLocations.stream()
        .flatMap(
            location -> {
              var blobExpandedCommitA = locationBlobsMapCommitA.get(location);
              var blobExpandedCommitB = locationBlobsMapCommitB.get(location);
              final AutogenBlob a = fromBlobProto(blobExpandedCommitA);
              final AutogenBlob b = fromBlobProto(blobExpandedCommitB);
              if (TypeChecker.sameType(a, b)) {
                AutogenBlobDiff blobDiff = DiffComputer.computeBlobDiff(a, b);
                // diff can be null because the old hash computation could detect two same entities
                // evaluate to different sha because it evaluated the list it contains in random
                // order
                if (blobDiff != null) {
                  return Stream.of(
                      blobDiff
                          .setLocation(blobExpandedCommitA.getLocationList())
                          .setStatus(AutogenDiffStatusEnumDiffStatus.fromProto(DiffStatus.MODIFIED))
                          .toProto()
                          .build());
                }
                return null;
              } else {
                return Stream.of(
                    DiffComputer.computeBlobDiff(a, null)
                        .setLocation(blobExpandedCommitA.getLocationList())
                        .setStatus(AutogenDiffStatusEnumDiffStatus.fromProto(DiffStatus.DELETED))
                        .toProto()
                        .build(),
                    DiffComputer.computeBlobDiff(null, b)
                        .setLocation(blobExpandedCommitB.getLocationList())
                        .setStatus(AutogenDiffStatusEnumDiffStatus.fromProto(DiffStatus.ADDED))
                        .toProto()
                        .build());
              }
            })
        .collect(Collectors.toList());
  }

  private Map<String, BlobExpanded> getLocationWiseBlobExpandedMapFromCollection(
      Collection<BlobExpanded> blobExpandeds) {
    return blobExpandeds.stream()
        .collect(
            Collectors.toMap(
                // TODO: Here used the `#` for joining the locations but if folder locations contain
                // TODO: - the `#` then this functionality will break.
                blobExpanded -> getStringFromLocationList(blobExpanded.getLocationList()),
                blobExpanded -> blobExpanded));
  }

  private String getStringFromLocationList(List<String> locationList) {
    return String.join("#", locationList);
  }

  @Override
  public List<BlobContainer> convertBlobDiffsToBlobs(
      List<AutogenBlobDiff> diffs,
      RepositoryFunction repositoryFunction,
      CommitFunction commitFunction)
      throws ModelDBException {
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {
      var repositoryEntity = repositoryFunction.apply(session);
      var commitEntity = commitFunction.apply(session, session1 -> repositoryEntity);
      Map<String, BlobExpanded> locationBlobsMap =
          getCommitBlobMap(session, commitEntity.getRootSha(), new ArrayList<>());

      return getBlobContainers(diffs, locationBlobsMap, new HashMap<String, List<String>>());
    } catch (Exception ex) {
      if (ModelDBUtils.needToRetry(ex)) {
        return convertBlobDiffsToBlobs(diffs, repositoryFunction, commitFunction);
      } else {
        throw ex;
      }
    }
  }

  private List<BlobContainer> getBlobContainers(
      List<AutogenBlobDiff> diffs,
      Map<String, BlobExpanded> locationBlobsMap,
      HashMap<String, List<String>> conflictLocationMap)
      throws ModelDBException {
    Map<String, BlobExpanded> locationBlobsMapNew = new LinkedHashMap<>();
    for (AutogenBlobDiff blobDiff : diffs) {
      final List<String> locationList = blobDiff.getLocation();
      if (locationList == null || locationList.isEmpty()) {
        throw new ModelDBException("Location in BlobDiff should not be empty", Code.INTERNAL);
      }
      var blobExpanded = locationBlobsMap.get(getStringFromLocationList(locationList));

      HashSet<String> conflictKeys = new HashSet<>();
      AutogenBlob blob =
          DiffMerger.mergeBlob(
              blobExpanded == null ? null : AutogenBlob.fromProto(blobExpanded.getBlob()),
              blobDiff,
              conflictKeys);
      if (!conflictKeys.isEmpty()) {
        if (!conflictLocationMap.containsKey(getStringFromLocationList(locationList))) {
          conflictLocationMap.put(
              getStringFromLocationList(locationList), new LinkedList<String>());
        }
        conflictLocationMap.get(getStringFromLocationList(locationList)).addAll(conflictKeys);
      }
      locationBlobsMapNew.put(
          getStringFromLocationList(locationList),
          blob == null
              ? null
              : BlobExpanded.newBuilder()
                  .addAllLocation(locationList)
                  .setBlob(blob.toProto())
                  .build());
    }

    locationBlobsMap.putAll(locationBlobsMapNew);
    List<BlobContainer> blobContainerList = new LinkedList<>();
    for (Map.Entry<String, BlobExpanded> blobExpandedEntry : locationBlobsMap.entrySet()) {
      if (blobExpandedEntry.getValue() != null) {
        blobContainerList.add(BlobContainer.create(blobExpandedEntry.getValue()));
      }
    }
    return blobContainerList;
  }

  private Boolean blobTypeExistsInList(List<BlobType> blobTypeList, Blob.ContentCase contentCase)
      throws ModelDBException {
    switch (contentCase) {
      case DATASET:
        return blobTypeList.contains(BlobType.DATASET_BLOB);
      case CONFIG:
        return blobTypeList.contains(BlobType.CONFIG_BLOB);
      case CODE:
        return blobTypeList.contains(BlobType.CODE_BLOB);
      case ENVIRONMENT:
        return blobTypeList.contains(BlobType.ENVIRONMENT_BLOB);
      default:
        throw new ModelDBException(
            "Invalid blob type found in DB Blob : " + contentCase.name(), Code.INTERNAL);
    }
  }

  @Override
  public Response findRepositoriesBlobs(
      CommitDAO commitDAO, FindRepositoriesBlobs request, List<Repository> repositories)
      throws ModelDBException {
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {

      var currentLoginUserInfo = uacApisUtil.getCurrentLoginUserInfo().blockAndGet();
      var commitPaginationDTO =
          commitDAO.findCommits(
              session,
              request,
              currentLoginUserInfo,
              false,
              true,
              false,
              ModelDBConstants.DATE_UPDATED,
              false);

      Set<String> rootShaList =
          commitPaginationDTO.getCommitEntities().stream()
              .map(CommitEntity::getRootSha)
              .collect(Collectors.toSet());
      Long totalRecords = commitPaginationDTO.getTotalRecords();

      Set<BlobExpanded> blobExpandedSet = new LinkedHashSet<>();
      for (String rootSha : rootShaList) {
        Map<String, BlobExpanded> blobExpandedMap =
            convertToLocationBlobMap(
                getCommitBlobMapWithHash(
                    session, rootSha, request.getLocationPrefixList(), request.getBlobTypeList()));
        blobExpandedSet.addAll(blobExpandedMap.values());
      }

      Map<Long, Workspace> cacheWorkspaceMap = new HashMap<>();
      Function<RepositoryEntity, Repository> toProto =
          (repositoryEntity) ->
              repositoryEntity.toProto(
                  mdbRoleService, uacApisUtil, cacheWorkspaceMap, new HashMap<>());
      repositories.addAll(
          commitPaginationDTO.getCommitEntities().stream()
              .flatMap(commitEntity -> commitEntity.getRepository().stream())
              .map(toProto)
              .collect(Collectors.toList()));

      return FindRepositoriesBlobs.Response.newBuilder()
          .addAllBlobs(blobExpandedSet)
          .setTotalRecords(totalRecords)
          .build();
    } catch (Exception ex) {
      if (ModelDBUtils.needToRetry(ex)) {
        return findRepositoriesBlobs(commitDAO, request, repositories);
      } else {
        throw new ModelDBException(ex);
      }
    }
  }

  @Override
  public GetUrlForDatasetBlobVersioned.Response getUrlForVersionedDatasetBlob(
      ArtifactStoreDAO artifactStoreDAO,
      RepositoryDAO repositoryDAO,
      String datasetId,
      CommitFunction commitFunction,
      GetUrlForDatasetBlobVersioned request)
      throws ModelDBException {
    var getUrlForBlobVersionedRequest =
        GetUrlForBlobVersioned.newBuilder()
            .setRepositoryId(
                RepositoryIdentification.newBuilder()
                    .setRepoId(Long.parseLong(request.getDatasetId()))
                    .build())
            .setCommitSha(request.getDatasetVersionId())
            .setMethod(request.getMethod())
            .setPartNumber(request.getPartNumber())
            .addLocation(DEFAULT_VERSIONING_BLOB_LOCATION)
            .setPathDatasetComponentBlobPath(request.getPathDatasetComponentBlobPath())
            .build();
    var response =
        getUrlForVersionedBlob(
            artifactStoreDAO,
            (session) ->
                VersioningUtils.getDatasetRepositoryEntity(
                    session,
                    repositoryDAO,
                    datasetId,
                    request.getDatasetVersionId(),
                    request.getMethod().equalsIgnoreCase("put")),
            commitFunction,
            getUrlForBlobVersionedRequest);
    return GetUrlForDatasetBlobVersioned.Response.newBuilder()
        .setUrl(response.getUrl())
        .setMultipartUploadOk(response.getMultipartUploadOk())
        .build();
  }

  @Override
  public GetUrlForBlobVersioned.Response getUrlForVersionedBlob(
      ArtifactStoreDAO artifactStoreDAO,
      RepositoryFunction repositoryFunction,
      CommitFunction commitFunction,
      GetUrlForBlobVersioned request)
      throws ModelDBException {
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {
      var repositoryEntity = repositoryFunction.apply(session);
      var commitEntity = commitFunction.apply(session, session1 -> repositoryEntity);

      Map<String, Map.Entry<BlobExpanded, String>> locationBlobWithHashMap =
          getCommitBlobMapWithHash(
              session,
              commitEntity.getRootSha(),
              request.getLocationList(),
              Collections.singletonList(BlobType.DATASET_BLOB));

      var locationKey = String.join("#", request.getLocationList());
      if (!locationBlobWithHashMap.containsKey(locationKey)) {
        throw new ModelDBException(
            String.format(
                "Blob Location '%s' not found in commit blobs", request.getLocationList()),
            Code.INVALID_ARGUMENT);
      }
      Map.Entry<BlobExpanded, String> blobExpandedMap = locationBlobWithHashMap.get(locationKey);
      var blob = blobExpandedMap.getKey().getBlob();

      var presignedUrlResponse = getPresignedUrl(artifactStoreDAO, session, request, blob);
      var responseBuilder = GetUrlForBlobVersioned.Response.newBuilder();
      if (presignedUrlResponse.getUrl() != null && !presignedUrlResponse.getUrl().isEmpty()) {
        responseBuilder.setUrl(presignedUrlResponse.getUrl());
        responseBuilder.setMultipartUploadOk(presignedUrlResponse.getMultipartUploadOk());
      }
      return responseBuilder.build();
    }
  }

  private GetUrlForArtifact.Response getPresignedUrl(
      ArtifactStoreDAO artifactStoreDAO, Session session, GetUrlForBlobVersioned request, Blob blob)
      throws ModelDBException {
    Map<String, String> componentWithSHAMap =
        getDatasetComponentBlob(blob, request.getPathDatasetComponentBlobPath());

    if (blob.getContentCase().equals(Blob.ContentCase.DATASET)) {
      String internalPath = componentWithSHAMap.get(INTERNAL_PATH_QUERY_PARAM);
      String computeSha = componentWithSHAMap.get(COMPUTE_SHA_QUERY_PARAM);
      Map.Entry<String, String> s3KeyUploadId =
          getS3PathAndMultipartUploadId(
              session,
              computeSha,
              internalPath,
              blob.getDataset().getContentCase(),
              request.getPartNumber() > 0,
              key -> artifactStoreDAO.initializeMultipart(internalPath));
      var errorMessage = "S3Key not found";
      String s3Key = s3KeyUploadId.getKey();
      String uploadId = s3KeyUploadId.getValue();
      if (s3Key == null || s3Key.isEmpty()) {
        LOGGER.warn(errorMessage);
        throw new ModelDBException(errorMessage, Code.NOT_FOUND);
      }
      return artifactStoreDAO.getUrlForArtifactMultipart(
          s3Key, request.getMethod(), request.getPartNumber(), uploadId);
    } else {
      throw new ModelDBException("Invalid Blob type found", Code.INVALID_ARGUMENT);
    }
  }

  /**
   * select Component matching the path passed in pathDatasetComponentBlobPath. returns a map
   * containing with keys 'internalPath' and 'computeSha' and values populated from the details of
   * the Component
   *
   * @param blob : commit blob
   * @param pathDatasetComponentBlobPath : pathDatasetComponentBlobPath
   * @return
   * @throws ModelDBException
   */
  private Map<String, String> getDatasetComponentBlob(
      Blob blob, String pathDatasetComponentBlobPath) throws ModelDBException {
    if (blob.getContentCase().equals(Blob.ContentCase.DATASET)) {
      var datasetBlob = blob.getDataset();
      switch (datasetBlob.getContentCase()) {
        case PATH:
          Optional<PathDatasetComponentBlob> pathComponentBlob =
              datasetBlob.getPath().getComponentsList().stream()
                  .filter(
                      componentBlob ->
                          componentBlob.getPath().equals(pathDatasetComponentBlobPath)
                              && componentBlob.getInternalVersionedPath() != null
                              && !componentBlob.getInternalVersionedPath().isEmpty())
                  .findFirst();
          if (pathComponentBlob.isPresent()
              && !pathComponentBlob.get().getInternalVersionedPath().isEmpty()) {
            var pathDatasetComponentBlob = pathComponentBlob.get();
            var autogenPathDatasetComponentBlob =
                AutogenPathDatasetComponentBlob.fromProto(pathDatasetComponentBlob);
            String computeSha = null;
            try {
              computeSha = DatasetContainer.computeSHA(autogenPathDatasetComponentBlob);
            } catch (NoSuchAlgorithmException e) {
              throw new ModelDBException(e);
            }
            Map<String, String> componentWithSHAMap = new HashMap<>();
            componentWithSHAMap.put(
                INTERNAL_PATH_QUERY_PARAM, pathDatasetComponentBlob.getInternalVersionedPath());
            componentWithSHAMap.put(COMPUTE_SHA_QUERY_PARAM, computeSha);
            return componentWithSHAMap;
          }
          break;
        case S3:
          Optional<S3DatasetComponentBlob> s3PathComponentBlob =
              datasetBlob.getS3().getComponentsList().stream()
                  .filter(
                      componentBlob ->
                          componentBlob.getPath().getPath().equals(pathDatasetComponentBlobPath)
                              && componentBlob.getPath().getInternalVersionedPath() != null
                              && !componentBlob.getPath().getInternalVersionedPath().isEmpty())
                  .findFirst();
          if (s3PathComponentBlob.isPresent()
              && !s3PathComponentBlob.get().getPath().getInternalVersionedPath().isEmpty()) {
            var componentBlob = AutogenS3DatasetComponentBlob.fromProto(s3PathComponentBlob.get());
            String computeS3Sha = null;
            try {
              computeS3Sha = DatasetContainer.computeSHA(componentBlob);
            } catch (NoSuchAlgorithmException e) {
              throw new ModelDBException(e);
            }
            Map<String, String> componentWithSHAMap = new HashMap<>();
            componentWithSHAMap.put(
                INTERNAL_PATH_QUERY_PARAM,
                s3PathComponentBlob.get().getPath().getInternalVersionedPath());
            componentWithSHAMap.put(COMPUTE_SHA_QUERY_PARAM, computeS3Sha);
            return componentWithSHAMap;
          }
          break;
        default:
          // Do nothing
          break;
      }
    } else {
      throw new ModelDBException("Invalid Blob type found", Code.INVALID_ARGUMENT);
    }
    throw new ModelDBException(
        "Dataset component blob not found for the path: '" + pathDatasetComponentBlobPath + "'",
        Code.NOT_FOUND);
  }

  private Map.Entry<String, String> getS3PathAndMultipartUploadId(
      Session session,
      String computeSha,
      String internalVersionedPath,
      DatasetBlob.ContentCase contentCase,
      boolean partNumberSpecified,
      S3KeyFunction initializeMultipart)
      throws ModelDBException {
    List<UploadStatusEntity> uploadStatusEntities =
        getUploadStatusEntity(session, computeSha, contentCase);
    UploadStatusEntity uploadStatusEntity = null;
    if (uploadStatusEntities != null && !uploadStatusEntities.isEmpty()) {
      if (uploadStatusEntities.size() > 1) {
        LOGGER.warn(
            "Multiple upload status found for datasetComponentPathId : "
                + computeSha
                + ", go ahead with first upload status");
      }
      uploadStatusEntity = uploadStatusEntities.get(0);
    }

    String uploadId;
    if (partNumberSpecified) {
      uploadId =
          uploadStatusEntity == null
                  || uploadStatusEntity.getUploadId() == null
                  || uploadStatusEntity.getUploadId().isEmpty()
              ? null
              : uploadStatusEntity.getUploadId();

      String message = null;
      if (uploadId == null) {
        if (initializeMultipart == null) {
          message = "Multipart wasn't initialized";
        } else {
          uploadId = initializeMultipart.apply(internalVersionedPath).orElse(null);
          if (uploadStatusEntity == null) {
            uploadStatusEntity = new UploadStatusEntity();
          }
          uploadStatusEntity.setDataset_component_blob_id(computeSha);
          if (contentCase.equals(DatasetBlob.ContentCase.PATH)) {
            uploadStatusEntity.setComponent_blob_type(
                UploadStatusEntity.PATH_DATASET_COMPONENT_BLOB);
          } else {
            uploadStatusEntity.setComponent_blob_type(UploadStatusEntity.S3_DATASET_COMPONENT_BLOB);
          }
        }
      }
      if (message != null) {
        LOGGER.info(message);
        throw new ModelDBException(message, Code.FAILED_PRECONDITION);
      }

      if (!Objects.equals(uploadId, uploadStatusEntity.getUploadId())
          || uploadStatusEntity.isUploadCompleted()) {
        session.beginTransaction();
        uploadStatusEntity.setUploadId(uploadId);
        uploadStatusEntity.setUploadCompleted(false);
        session.saveOrUpdate(uploadStatusEntity);
        session.getTransaction().commit();
      }
    } else {
      uploadId = null;
    }
    return new AbstractMap.SimpleEntry<>(internalVersionedPath, uploadId);
  }

  private List<UploadStatusEntity> getUploadStatusEntity(
      Session session, String datasetComponentPathId, DatasetBlob.ContentCase contentCase)
      throws ModelDBException {
    var getUploadStatusQuery =
        new StringBuilder("From " + UploadStatusEntity.class.getSimpleName() + " us WHERE ");
    getUploadStatusQuery.append(" us.dataset_component_blob_id = :pathId ");
    if (contentCase.equals(DatasetBlob.ContentCase.PATH)) {
      getUploadStatusQuery.append(
          " AND us.component_blob_type = " + UploadStatusEntity.PATH_DATASET_COMPONENT_BLOB);
    } else if (contentCase.equals(DatasetBlob.ContentCase.S3)) {
      getUploadStatusQuery.append(
          " AND us.component_blob_type = " + UploadStatusEntity.S3_DATASET_COMPONENT_BLOB);
    } else {
      throw new ModelDBException("Invalid content case found in DatasetBlob", Code.INTERNAL);
    }
    var query = session.createQuery(getUploadStatusQuery.toString());
    query.setParameter("pathId", datasetComponentPathId);
    return query.list();
  }

  @Override
  public void commitVersionedDatasetBlobArtifactPart(
      RepositoryDAO repositoryDAO,
      String datasetId,
      CommitFunction commitFunction,
      CommitVersionedDatasetBlobArtifactPart request)
      throws ModelDBException {
    commitVersionedBlobArtifactPart(
        (session) ->
            VersioningUtils.getDatasetRepositoryEntity(
                session, repositoryDAO, datasetId, request.getDatasetVersionId(), true),
        commitFunction,
        Collections.singletonList(DEFAULT_VERSIONING_BLOB_LOCATION),
        request.getPathDatasetComponentBlobPath(),
        request.getArtifactPart());
  }

  @Override
  public CommitVersionedBlobArtifactPart.Response commitVersionedBlobArtifactPart(
      RepositoryFunction repositoryFunction,
      CommitFunction commitFunction,
      List<String> location,
      String pathDatasetComponentBlobPath,
      ArtifactPart artifactPart)
      throws ModelDBException {
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {
      RepositoryEntity repository = repositoryFunction.apply(session);
      var commitEntity = commitFunction.apply(session, session1 -> repository);
      session.lock(commitEntity, LockMode.PESSIMISTIC_WRITE);

      Map<String, Map.Entry<BlobExpanded, String>> locationBlobWithHashMap =
          getCommitBlobMapWithHash(
              session,
              commitEntity.getRootSha(),
              location,
              Collections.singletonList(BlobType.DATASET_BLOB));

      var locationKey = String.join("#", location);
      if (!locationBlobWithHashMap.containsKey(locationKey)) {
        throw new ModelDBException(
            String.format("Blob Location '%s' not found in commit blobs", location),
            Code.INVALID_ARGUMENT);
      }
      Map.Entry<BlobExpanded, String> blobExpandedMap = locationBlobWithHashMap.get(locationKey);

      Map<String, String> componentWithSHAMap =
          getDatasetComponentBlob(blobExpandedMap.getKey().getBlob(), pathDatasetComponentBlobPath);
      String computeSha = componentWithSHAMap.get(COMPUTE_SHA_QUERY_PARAM);

      VersioningUtils.saveArtifactPartEntity(
          artifactPart, session, computeSha, ArtifactPartEntity.VERSION_BLOB_ARTIFACT);
      return CommitVersionedBlobArtifactPart.Response.newBuilder().build();
    } catch (Exception e) {
      throw new ModelDBException(e);
    }
  }

  @Override
  public GetCommittedVersionedDatasetBlobArtifactParts.Response
      getCommittedVersionedDatasetBlobArtifactParts(
          RepositoryDAO repositoryDAO,
          String datasetId,
          CommitFunction commitFunction,
          GetCommittedVersionedDatasetBlobArtifactParts request)
          throws ModelDBException {
    var blobArtifactPartsResponse =
        getCommittedVersionedBlobArtifactParts(
            (session) ->
                VersioningUtils.getDatasetRepositoryEntity(
                    session, repositoryDAO, datasetId, request.getDatasetVersionId(), true),
            commitFunction,
            Collections.singletonList(DEFAULT_VERSIONING_BLOB_LOCATION),
            request.getPathDatasetComponentBlobPath());
    return GetCommittedVersionedDatasetBlobArtifactParts.Response.newBuilder()
        .addAllArtifactParts(blobArtifactPartsResponse.getArtifactPartsList())
        .build();
  }

  @Override
  public GetCommittedVersionedBlobArtifactParts.Response getCommittedVersionedBlobArtifactParts(
      RepositoryFunction repositoryFunction,
      CommitFunction commitFunction,
      List<String> location,
      String pathDatasetComponentBlobPath)
      throws ModelDBException {
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {
      RepositoryEntity repository = repositoryFunction.apply(session);
      var commitEntity = commitFunction.apply(session, session1 -> repository);

      Map<String, Map.Entry<BlobExpanded, String>> locationBlobWithHashMap =
          getCommitBlobMapWithHash(
              session,
              commitEntity.getRootSha(),
              location,
              Collections.singletonList(BlobType.DATASET_BLOB));

      var locationKey = String.join("#", location);
      if (!locationBlobWithHashMap.containsKey(locationKey)) {
        throw new ModelDBException(
            String.format("Blob Location '%s' not found in commit blobs", location),
            Code.INVALID_ARGUMENT);
      }
      Map.Entry<BlobExpanded, String> blobExpandedMap = locationBlobWithHashMap.get(locationKey);
      Map<String, String> componentWithSHAMap =
          getDatasetComponentBlob(blobExpandedMap.getKey().getBlob(), pathDatasetComponentBlobPath);
      String computeSha = componentWithSHAMap.get(COMPUTE_SHA_QUERY_PARAM);

      Set<ArtifactPartEntity> artifactPartEntities =
          VersioningUtils.getArtifactPartEntities(
              session, computeSha, ArtifactPartEntity.VERSION_BLOB_ARTIFACT);
      var response = GetCommittedVersionedBlobArtifactParts.Response.newBuilder();
      artifactPartEntities.forEach(
          artifactPartEntity -> response.addArtifactParts(artifactPartEntity.toProto()));
      return response.build();
    } catch (Exception e) {
      throw new ModelDBException(e);
    }
  }

  @Override
  public void commitMultipartVersionedDatasetBlobArtifact(
      RepositoryDAO repositoryDAO,
      String datasetId,
      CommitFunction commitFunction,
      CommitMultipartVersionedDatasetBlobArtifact request,
      CommitMultipartFunction commitMultipartFunction)
      throws ModelDBException {
    commitMultipartVersionedBlobArtifact(
        (session) ->
            VersioningUtils.getDatasetRepositoryEntity(
                session, repositoryDAO, datasetId, request.getDatasetVersionId(), true),
        commitFunction,
        Collections.singletonList(DEFAULT_VERSIONING_BLOB_LOCATION),
        request.getPathDatasetComponentBlobPath(),
        commitMultipartFunction);
  }

  @Override
  public CommitMultipartVersionedBlobArtifact.Response commitMultipartVersionedBlobArtifact(
      RepositoryFunction repositoryFunction,
      CommitFunction commitFunction,
      List<String> location,
      String pathDatasetComponentBlobPath,
      CommitMultipartFunction commitMultipartFunction)
      throws ModelDBException {
    List<PartETag> partETags;
    try (var session = modelDBHibernateUtil.getSessionFactory().openSession()) {
      RepositoryEntity repository = repositoryFunction.apply(session);
      var commitEntity = commitFunction.apply(session, session1 -> repository);
      session.lock(commitEntity, LockMode.PESSIMISTIC_WRITE);

      Map<String, Map.Entry<BlobExpanded, String>> locationBlobWithHashMap =
          getCommitBlobMapWithHash(
              session,
              commitEntity.getRootSha(),
              location,
              Collections.singletonList(BlobType.DATASET_BLOB));

      var locationKey = String.join("#", location);
      if (!locationBlobWithHashMap.containsKey(locationKey)) {
        throw new ModelDBException(
            String.format("Blob Location '%s' not found in commit blobs", location),
            Code.INVALID_ARGUMENT);
      }
      Map.Entry<BlobExpanded, String> blobExpandedMap = locationBlobWithHashMap.get(locationKey);
      var blob = blobExpandedMap.getKey().getBlob();
      Map<String, String> componentWithSHAMap =
          getDatasetComponentBlob(blob, pathDatasetComponentBlobPath);
      String internalPath = componentWithSHAMap.get(INTERNAL_PATH_QUERY_PARAM);
      String computeSha = componentWithSHAMap.get(COMPUTE_SHA_QUERY_PARAM);

      if (!blob.getContentCase().equals(Blob.ContentCase.DATASET)) {
        throw new ModelDBException(
            "Invalid Blob type found for given location", Code.INVALID_ARGUMENT);
      }

      List<UploadStatusEntity> uploadStatusEntities =
          getUploadStatusEntity(session, computeSha, blob.getDataset().getContentCase());
      UploadStatusEntity uploadStatusEntity;
      if (uploadStatusEntities == null || uploadStatusEntities.size() == 0) {
        throw new ModelDBException(
            "Multipart wasn't initialized OR Artifact upload status not found in DB datasetComponentPathId : "
                + computeSha,
            Code.FAILED_PRECONDITION);
      } else if (uploadStatusEntities.size() > 1) {
        throw new ModelDBException(
            "Multiple upload status found for datasetComponentPathId : " + computeSha,
            Code.FAILED_PRECONDITION);
      } else {
        uploadStatusEntity = uploadStatusEntities.get(0);
      }

      Set<ArtifactPartEntity> artifactPartEntities =
          VersioningUtils.getArtifactPartEntities(
              session, computeSha, ArtifactPartEntity.VERSION_BLOB_ARTIFACT);
      partETags =
          artifactPartEntities.stream()
              .map(ArtifactPartEntity::toPartETag)
              .collect(Collectors.toList());
      commitMultipartFunction.apply(internalPath, uploadStatusEntity.getUploadId(), partETags);
      session.beginTransaction();
      uploadStatusEntity.setUploadCompleted(true);
      uploadStatusEntity.setUploadId(null);
      artifactPartEntities.forEach(session::delete);
      artifactPartEntities.clear();
      session.saveOrUpdate(uploadStatusEntity);
      session.getTransaction().commit();
    } catch (Exception e) {
      throw new ModelDBException(e);
    }
    return CommitMultipartVersionedBlobArtifact.Response.newBuilder().build();
  }
}
