private int countOutgoingEdges(AtlasVertex typeVertex, String edgeLabel) {

    Iterator<AtlasEdge> outGoingEdgesByLabel =
        GraphHelper.getInstance().getOutGoingEdgesByLabel(typeVertex, edgeLabel);
    int edgeCount = 0;
    for (Iterator<AtlasEdge> iterator = outGoingEdgesByLabel; iterator.hasNext(); ) {
      iterator.next();
      edgeCount++;
    }
    return edgeCount;
  }
public abstract class DeleteHandler {
  public static final Logger LOG = LoggerFactory.getLogger(DeleteHandler.class);

  protected static final GraphHelper graphHelper = GraphHelper.getInstance();

  protected TypeSystem typeSystem;
  private boolean shouldUpdateReverseAttribute;
  private boolean softDelete;

  public DeleteHandler(
      TypeSystem typeSystem, boolean shouldUpdateReverseAttribute, boolean softDelete) {
    this.typeSystem = typeSystem;
    this.shouldUpdateReverseAttribute = shouldUpdateReverseAttribute;
    this.softDelete = softDelete;
  }

  /**
   * Deletes the entity vertex - deletes the traits and all the references
   *
   * @param instanceVertex
   * @throws AtlasException
   */
  public void deleteEntity(Vertex instanceVertex) throws AtlasException {
    RequestContext requestContext = RequestContext.get();
    String guid = GraphHelper.getIdFromVertex(instanceVertex);
    Id.EntityState state = GraphHelper.getState(instanceVertex);
    if (requestContext.getDeletedEntityIds().contains(guid) || state == Id.EntityState.DELETED) {
      LOG.debug("Skipping deleting {} as its already deleted", guid);
      return;
    }
    String typeName = GraphHelper.getTypeName(instanceVertex);
    requestContext.recordEntityDelete(guid, typeName);

    deleteAllTraits(instanceVertex);

    deleteTypeVertex(instanceVertex, false);
  }

  protected abstract void deleteEdge(Edge edge, boolean force) throws AtlasException;

  /**
   * Deletes a type vertex - can be entity(class type) or just vertex(struct/trait type)
   *
   * @param instanceVertex
   * @param typeCategory
   * @throws AtlasException
   */
  protected void deleteTypeVertex(
      Vertex instanceVertex, DataTypes.TypeCategory typeCategory, boolean force)
      throws AtlasException {
    switch (typeCategory) {
      case STRUCT:
      case TRAIT:
        deleteTypeVertex(instanceVertex, force);
        break;

      case CLASS:
        deleteEntity(instanceVertex);
        break;

      default:
        throw new IllegalStateException("Type category " + typeCategory + " not handled");
    }
  }

  /**
   * Deleting any type vertex. Goes over the complex attributes and removes the references
   *
   * @param instanceVertex
   * @throws AtlasException
   */
  protected void deleteTypeVertex(Vertex instanceVertex, boolean force) throws AtlasException {
    LOG.debug("Deleting {}", string(instanceVertex));
    String typeName = GraphHelper.getTypeName(instanceVertex);
    IDataType type = typeSystem.getDataType(IDataType.class, typeName);
    FieldMapping fieldMapping = getFieldMapping(type);

    for (AttributeInfo attributeInfo : fieldMapping.fields.values()) {
      LOG.debug("Deleting attribute {} for {}", attributeInfo.name, string(instanceVertex));
      String edgeLabel = GraphHelper.getEdgeLabel(type, attributeInfo);

      switch (attributeInfo.dataType().getTypeCategory()) {
        case CLASS:
          // If its class attribute, delete the reference
          deleteEdgeReference(
              instanceVertex, edgeLabel, DataTypes.TypeCategory.CLASS, attributeInfo.isComposite);
          break;

        case STRUCT:
          // If its struct attribute, delete the reference
          deleteEdgeReference(instanceVertex, edgeLabel, DataTypes.TypeCategory.STRUCT, false);
          break;

        case ARRAY:
          // For array attribute, if the element is struct/class, delete all the references
          IDataType elementType = ((DataTypes.ArrayType) attributeInfo.dataType()).getElemType();
          DataTypes.TypeCategory elementTypeCategory = elementType.getTypeCategory();
          if (elementTypeCategory == DataTypes.TypeCategory.STRUCT
              || elementTypeCategory == DataTypes.TypeCategory.CLASS) {
            Iterator<Edge> edges = GraphHelper.getOutGoingEdgesByLabel(instanceVertex, edgeLabel);
            if (edges != null) {
              while (edges.hasNext()) {
                Edge edge = edges.next();
                deleteEdgeReference(
                    edge, elementType.getTypeCategory(), attributeInfo.isComposite, false);
              }
            }
          }
          break;

        case MAP:
          // For map attribute, if the value type is struct/class, delete all the references
          DataTypes.MapType mapType = (DataTypes.MapType) attributeInfo.dataType();
          DataTypes.TypeCategory valueTypeCategory = mapType.getValueType().getTypeCategory();
          String propertyName = GraphHelper.getQualifiedFieldName(type, attributeInfo.name);

          if (valueTypeCategory == DataTypes.TypeCategory.STRUCT
              || valueTypeCategory == DataTypes.TypeCategory.CLASS) {
            List<String> keys = instanceVertex.getProperty(propertyName);
            if (keys != null) {
              for (String key : keys) {
                String mapEdgeLabel = GraphHelper.getQualifiedNameForMapKey(edgeLabel, key);
                deleteEdgeReference(
                    instanceVertex, mapEdgeLabel, valueTypeCategory, attributeInfo.isComposite);
              }
            }
          }
      }
    }

    deleteVertex(instanceVertex, force);
  }

  /**
   * Force delete is used to remove struct/trait in case of entity updates
   *
   * @param edge
   * @param typeCategory
   * @param isComposite
   * @param forceDeleteStructTrait
   * @return returns true if the edge reference is hard deleted
   * @throws AtlasException
   */
  public boolean deleteEdgeReference(
      Edge edge,
      DataTypes.TypeCategory typeCategory,
      boolean isComposite,
      boolean forceDeleteStructTrait)
      throws AtlasException {
    LOG.debug("Deleting {}", string(edge));
    boolean forceDelete =
        (typeCategory == DataTypes.TypeCategory.STRUCT
                || typeCategory == DataTypes.TypeCategory.TRAIT)
            ? forceDeleteStructTrait
            : false;
    if (typeCategory == DataTypes.TypeCategory.STRUCT
        || typeCategory == DataTypes.TypeCategory.TRAIT
        || (typeCategory == DataTypes.TypeCategory.CLASS && isComposite)) {
      // If the vertex is of type struct/trait, delete the edge and then the reference vertex as the
      // vertex is not shared by any other entities.
      // If the vertex is of type class, and its composite attribute, this reference vertex'
      // lifecycle is controlled
      // through this delete, hence delete the edge and the reference vertex.
      Vertex vertexForDelete = edge.getVertex(Direction.IN);

      // If deleting the edge and then the in vertex, reverse attribute shouldn't be updated
      deleteEdge(edge, false, forceDelete);
      deleteTypeVertex(vertexForDelete, typeCategory, forceDelete);
    } else {
      // If the vertex is of type class, and its not a composite attributes, the reference vertex'
      // lifecycle is not controlled
      // through this delete. Hence just remove the reference edge. Leave the reference vertex as is

      // If deleting just the edge, reverse attribute should be updated for any references
      // For example, for the department type system, if the person's manager edge is deleted,
      // subordinates of manager should be updated
      deleteEdge(edge, true, false);
    }
    return !softDelete || forceDelete;
  }

  public void deleteEdgeReference(
      Vertex outVertex, String edgeLabel, DataTypes.TypeCategory typeCategory, boolean isComposite)
      throws AtlasException {
    Edge edge = GraphHelper.getEdgeForLabel(outVertex, edgeLabel);
    if (edge != null) {
      deleteEdgeReference(edge, typeCategory, isComposite, false);
    }
  }

  protected void deleteEdge(Edge edge, boolean updateReverseAttribute, boolean force)
      throws AtlasException {
    // update reverse attribute
    if (updateReverseAttribute) {
      AttributeInfo attributeInfo = getAttributeForEdge(edge.getLabel());
      if (attributeInfo.reverseAttributeName != null) {
        deleteEdgeBetweenVertices(
            edge.getVertex(Direction.IN),
            edge.getVertex(Direction.OUT),
            attributeInfo.reverseAttributeName);
      }
    }

    deleteEdge(edge, force);
  }

  protected void deleteVertex(Vertex instanceVertex, boolean force) throws AtlasException {
    // Update external references(incoming edges) to this vertex
    LOG.debug(
        "Setting the external references to {} to null(removing edges)", string(instanceVertex));
    Iterator<Edge> edges = instanceVertex.getEdges(Direction.IN).iterator();

    while (edges.hasNext()) {
      Edge edge = edges.next();
      Id.EntityState edgeState = GraphHelper.getState(edge);
      if (edgeState == Id.EntityState.ACTIVE) {
        // Delete only the active edge references
        AttributeInfo attribute = getAttributeForEdge(edge.getLabel());
        // TODO use delete edge instead??
        deleteEdgeBetweenVertices(
            edge.getVertex(Direction.OUT), edge.getVertex(Direction.IN), attribute.name);
      }
    }
    _deleteVertex(instanceVertex, force);
  }

  protected abstract void _deleteVertex(Vertex instanceVertex, boolean force);

  /**
   * Deletes the edge between outvertex and inVertex. The edge is for attribute attributeName of
   * outVertex
   *
   * @param outVertex
   * @param inVertex
   * @param attributeName
   * @throws AtlasException
   */
  protected void deleteEdgeBetweenVertices(Vertex outVertex, Vertex inVertex, String attributeName)
      throws AtlasException {
    LOG.debug(
        "Removing edge from {} to {} with attribute name {}",
        string(outVertex),
        string(inVertex),
        attributeName);
    String typeName = GraphHelper.getTypeName(outVertex);
    String outId = GraphHelper.getIdFromVertex(outVertex);
    Id.EntityState state = GraphHelper.getState(outVertex);
    if ((outId != null && RequestContext.get().isDeletedEntity(outId))
        || state == Id.EntityState.DELETED) {
      // If the reference vertex is marked for deletion, skip updating the reference
      return;
    }

    IDataType type = typeSystem.getDataType(IDataType.class, typeName);
    AttributeInfo attributeInfo = getFieldMapping(type).fields.get(attributeName);
    String propertyName = GraphHelper.getQualifiedFieldName(type, attributeName);
    String edgeLabel = EDGE_LABEL_PREFIX + propertyName;
    Edge edge = null;

    switch (attributeInfo.dataType().getTypeCategory()) {
      case CLASS:
        // If its class attribute, its the only edge between two vertices
        if (attributeInfo.multiplicity.nullAllowed()) {
          edge = GraphHelper.getEdgeForLabel(outVertex, edgeLabel);
          if (shouldUpdateReverseAttribute) {
            GraphHelper.setProperty(outVertex, propertyName, null);
          }
        } else {
          // Cannot unset a required attribute.
          throw new NullRequiredAttributeException(
              "Cannot unset required attribute "
                  + GraphHelper.getQualifiedFieldName(type, attributeName)
                  + " on "
                  + string(outVertex)
                  + " edge = "
                  + edgeLabel);
        }
        break;

      case ARRAY:
        // If its array attribute, find the right edge between the two vertices and update array
        // property
        List<String> elements = outVertex.getProperty(propertyName);
        if (elements != null) {
          elements =
              new ArrayList<>(
                  elements); // Make a copy, else list.remove reflects on titan.getProperty()
          for (String elementEdgeId : elements) {
            Edge elementEdge = graphHelper.getEdgeByEdgeId(outVertex, edgeLabel, elementEdgeId);
            if (elementEdge == null) {
              continue;
            }

            Vertex elementVertex = elementEdge.getVertex(Direction.IN);
            if (elementVertex.getId().toString().equals(inVertex.getId().toString())) {
              edge = elementEdge;

              // TODO element.size includes deleted items as well. should exclude
              if (!attributeInfo.multiplicity.nullAllowed()
                  && elements.size() <= attributeInfo.multiplicity.lower) {
                // Deleting this edge would violate the attribute's lower bound.
                throw new NullRequiredAttributeException(
                    "Cannot remove array element from required attribute "
                        + GraphHelper.getQualifiedFieldName(type, attributeName)
                        + " on "
                        + string(outVertex)
                        + " "
                        + string(elementEdge));
              }

              if (shouldUpdateReverseAttribute) {
                // if composite attribute, remove the reference as well. else, just remove the edge
                // for example, when table is deleted, process still references the table
                // but when column is deleted, table will not reference the deleted column
                LOG.debug(
                    "Removing edge {} from the array attribute {}",
                    string(elementEdge),
                    attributeName);
                elements.remove(elementEdge.getId().toString());
                GraphHelper.setProperty(outVertex, propertyName, elements);
                break;
              }
            }
          }
        }
        break;

      case MAP:
        // If its map attribute, find the right edge between two vertices and update map property
        List<String> keys = outVertex.getProperty(propertyName);
        if (keys != null) {
          keys =
              new ArrayList<>(
                  keys); // Make a copy, else list.remove reflects on titan.getProperty()
          for (String key : keys) {
            String keyPropertyName = GraphHelper.getQualifiedNameForMapKey(propertyName, key);
            String mapEdgeId = outVertex.getProperty(keyPropertyName);
            Edge mapEdge = graphHelper.getEdgeByEdgeId(outVertex, keyPropertyName, mapEdgeId);
            Vertex mapVertex = mapEdge.getVertex(Direction.IN);
            if (mapVertex.getId().toString().equals(inVertex.getId().toString())) {
              // TODO keys.size includes deleted items as well. should exclude
              if (attributeInfo.multiplicity.nullAllowed()
                  || keys.size() > attributeInfo.multiplicity.lower) {
                edge = mapEdge;
              } else {
                // Deleting this entry would violate the attribute's lower bound.
                throw new NullRequiredAttributeException(
                    "Cannot remove map entry "
                        + keyPropertyName
                        + " from required attribute "
                        + GraphHelper.getQualifiedFieldName(type, attributeName)
                        + " on "
                        + string(outVertex)
                        + " "
                        + string(mapEdge));
              }

              if (shouldUpdateReverseAttribute) {
                // remove this key
                LOG.debug(
                    "Removing edge {}, key {} from the map attribute {}",
                    string(mapEdge),
                    key,
                    attributeName);
                keys.remove(key);
                GraphHelper.setProperty(outVertex, propertyName, keys);
                GraphHelper.setProperty(outVertex, keyPropertyName, null);
              }
              break;
            }
          }
        }
        break;

      case STRUCT:
      case TRAIT:
        break;

      default:
        throw new IllegalStateException(
            "There can't be an edge from "
                + string(outVertex)
                + " to "
                + string(inVertex)
                + " with attribute name "
                + attributeName
                + " which is not class/array/map attribute");
    }

    if (edge != null) {
      deleteEdge(edge, false);
      RequestContext requestContext = RequestContext.get();
      GraphHelper.setProperty(
          outVertex,
          Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY,
          requestContext.getRequestTime());
      requestContext.recordEntityUpdate(outId);
    }
  }

  protected AttributeInfo getAttributeForEdge(String edgLabel) throws AtlasException {
    AtlasEdgeLabel atlasEdgeLabel = new AtlasEdgeLabel(edgLabel);
    IDataType referenceType = typeSystem.getDataType(IDataType.class, atlasEdgeLabel.getTypeName());
    return getFieldMapping(referenceType).fields.get(atlasEdgeLabel.getAttributeName());
  }

  protected FieldMapping getFieldMapping(IDataType type) {
    switch (type.getTypeCategory()) {
      case CLASS:
      case TRAIT:
        return ((HierarchicalType) type).fieldMapping();

      case STRUCT:
        return ((StructType) type).fieldMapping();

      default:
        throw new IllegalStateException("Type " + type + " doesn't have any fields!");
    }
  }

  /**
   * Delete all traits from the specified vertex.
   *
   * @param instanceVertex
   * @throws AtlasException
   */
  private void deleteAllTraits(Vertex instanceVertex) throws AtlasException {
    List<String> traitNames = GraphHelper.getTraitNames(instanceVertex);
    LOG.debug("Deleting traits {} for {}", traitNames, string(instanceVertex));
    String typeName = GraphHelper.getTypeName(instanceVertex);

    for (String traitNameToBeDeleted : traitNames) {
      String relationshipLabel = GraphHelper.getTraitLabel(typeName, traitNameToBeDeleted);
      deleteEdgeReference(instanceVertex, relationshipLabel, DataTypes.TypeCategory.TRAIT, false);
    }
  }
}
/** An implementation backed by a Graph database provided as a Graph Service. */
@Singleton
public class GraphBackedMetadataRepository implements MetadataRepository {

  private static final Logger LOG = LoggerFactory.getLogger(GraphBackedMetadataRepository.class);

  private static TypeSystem typeSystem = TypeSystem.getInstance();

  private static final GraphHelper graphHelper = GraphHelper.getInstance();

  private final TitanGraph titanGraph;

  private DeleteHandler deleteHandler;

  private GraphToTypedInstanceMapper graphToInstanceMapper;

  @Inject
  public GraphBackedMetadataRepository(
      GraphProvider<TitanGraph> graphProvider, DeleteHandler deleteHandler) {
    this.titanGraph = graphProvider.get();
    graphToInstanceMapper = new GraphToTypedInstanceMapper(titanGraph);
    this.deleteHandler = deleteHandler;
  }

  public GraphToTypedInstanceMapper getGraphToInstanceMapper() {
    return graphToInstanceMapper;
  }

  @Override
  public String getTypeAttributeName() {
    return Constants.ENTITY_TYPE_PROPERTY_KEY;
  }

  @Override
  public String getStateAttributeName() {
    return Constants.STATE_PROPERTY_KEY;
  }

  /**
   * Returns the property key used to store super type names.
   *
   * @return property key used to store super type names.
   */
  @Override
  public String getSuperTypeAttributeName() {
    return Constants.SUPER_TYPES_PROPERTY_KEY;
  }

  public String getIdAttributeName() {
    return Constants.GUID_PROPERTY_KEY;
  }

  @Override
  public String getTraitLabel(IDataType<?> dataType, String traitName) {
    return GraphHelper.getTraitLabel(dataType.getName(), traitName);
  }

  @Override
  public String getFieldNameInVertex(IDataType<?> dataType, AttributeInfo aInfo)
      throws AtlasException {
    if (aInfo.name.startsWith(Constants.INTERNAL_PROPERTY_KEY_PREFIX)) {
      return aInfo.name;
    }
    return GraphHelper.getQualifiedFieldName(dataType, aInfo.name);
  }

  public String getFieldNameInVertex(IDataType<?> dataType, String attrName) throws AtlasException {
    return GraphHelper.getQualifiedFieldName(dataType, attrName);
  }

  @Override
  public String getEdgeLabel(IDataType<?> dataType, AttributeInfo aInfo) throws AtlasException {
    return GraphHelper.getEdgeLabel(dataType, aInfo);
  }

  @Override
  @GraphTransaction
  public List<String> createEntities(ITypedReferenceableInstance... entities)
      throws RepositoryException, EntityExistsException {
    LOG.debug("adding entities={}", entities);
    try {
      TypedInstanceToGraphMapper instanceToGraphMapper =
          new TypedInstanceToGraphMapper(graphToInstanceMapper, deleteHandler);
      instanceToGraphMapper.mapTypedInstanceToGraph(
          TypedInstanceToGraphMapper.Operation.CREATE, entities);
      return RequestContext.get().getCreatedEntityIds();
    } catch (EntityExistsException e) {
      throw e;
    } catch (AtlasException e) {
      throw new RepositoryException(e);
    }
  }

  @Override
  @GraphTransaction
  public ITypedReferenceableInstance getEntityDefinition(String guid)
      throws RepositoryException, EntityNotFoundException {
    LOG.debug("Retrieving entity with guid={}", guid);

    Vertex instanceVertex = graphHelper.getVertexForGUID(guid);

    try {
      return graphToInstanceMapper.mapGraphToTypedInstance(guid, instanceVertex);
    } catch (AtlasException e) {
      throw new RepositoryException(e);
    }
  }

  @Override
  @GraphTransaction
  public ITypedReferenceableInstance getEntityDefinition(
      String entityType, String attribute, Object value) throws AtlasException {
    LOG.debug("Retrieving entity with type={} and {}={}", entityType, attribute, value);
    IDataType type = typeSystem.getDataType(IDataType.class, entityType);
    String propertyKey = getFieldNameInVertex(type, attribute);
    Vertex instanceVertex =
        graphHelper.findVertex(
            propertyKey,
            value,
            Constants.ENTITY_TYPE_PROPERTY_KEY,
            entityType,
            Constants.STATE_PROPERTY_KEY,
            Id.EntityState.ACTIVE.name());

    String guid = instanceVertex.getProperty(Constants.GUID_PROPERTY_KEY);
    return graphToInstanceMapper.mapGraphToTypedInstance(guid, instanceVertex);
  }

  @Override
  @GraphTransaction
  public List<String> getEntityList(String entityType) throws RepositoryException {
    LOG.debug("Retrieving entity list for type={}", entityType);
    GraphQuery query = titanGraph.query().has(Constants.ENTITY_TYPE_PROPERTY_KEY, entityType);
    Iterator<Vertex> results = query.vertices().iterator();
    if (!results.hasNext()) {
      return Collections.emptyList();
    }

    ArrayList<String> entityList = new ArrayList<>();
    while (results.hasNext()) {
      Vertex vertex = results.next();
      entityList.add(vertex.<String>getProperty(Constants.GUID_PROPERTY_KEY));
    }

    return entityList;
  }

  /**
   * Gets the list of trait names for a given entity represented by a guid.
   *
   * @param guid globally unique identifier for the entity
   * @return a list of trait names for the given entity guid
   * @throws RepositoryException
   */
  @Override
  @GraphTransaction
  public List<String> getTraitNames(String guid) throws AtlasException {
    LOG.debug("Retrieving trait names for entity={}", guid);
    Vertex instanceVertex = graphHelper.getVertexForGUID(guid);
    return GraphHelper.getTraitNames(instanceVertex);
  }

  /**
   * Adds a new trait to an existing entity represented by a guid.
   *
   * @param guid globally unique identifier for the entity
   * @param traitInstance trait instance that needs to be added to entity
   * @throws RepositoryException
   */
  @Override
  @GraphTransaction
  public void addTrait(String guid, ITypedStruct traitInstance) throws RepositoryException {
    Preconditions.checkNotNull(traitInstance, "Trait instance cannot be null");
    final String traitName = traitInstance.getTypeName();
    LOG.debug("Adding a new trait={} for entity={}", traitName, guid);

    try {
      Vertex instanceVertex = graphHelper.getVertexForGUID(guid);

      // add the trait instance as a new vertex
      final String typeName = GraphHelper.getTypeName(instanceVertex);

      TypedInstanceToGraphMapper instanceToGraphMapper =
          new TypedInstanceToGraphMapper(graphToInstanceMapper, deleteHandler);
      instanceToGraphMapper.mapTraitInstanceToVertex(
          traitInstance, typeSystem.getDataType(ClassType.class, typeName), instanceVertex);

      // update the traits in entity once adding trait instance is successful
      GraphHelper.addProperty(instanceVertex, Constants.TRAIT_NAMES_PROPERTY_KEY, traitName);
      GraphHelper.setProperty(
          instanceVertex,
          Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY,
          RequestContext.get().getRequestTime());

    } catch (RepositoryException e) {
      throw e;
    } catch (Exception e) {
      throw new RepositoryException(e);
    }
  }

  /**
   * Deletes a given trait from an existing entity represented by a guid.
   *
   * @param guid globally unique identifier for the entity
   * @param traitNameToBeDeleted name of the trait
   * @throws RepositoryException
   */
  @Override
  @GraphTransaction
  public void deleteTrait(String guid, String traitNameToBeDeleted)
      throws TraitNotFoundException, EntityNotFoundException, RepositoryException {
    LOG.debug("Deleting trait={} from entity={}", traitNameToBeDeleted, guid);

    Vertex instanceVertex = graphHelper.getVertexForGUID(guid);

    List<String> traitNames = GraphHelper.getTraitNames(instanceVertex);
    if (!traitNames.contains(traitNameToBeDeleted)) {
      throw new TraitNotFoundException(
          "Could not find trait="
              + traitNameToBeDeleted
              + " in the repository for entity: "
              + guid);
    }

    try {
      final String entityTypeName = GraphHelper.getTypeName(instanceVertex);
      String relationshipLabel = GraphHelper.getTraitLabel(entityTypeName, traitNameToBeDeleted);
      Edge edge = GraphHelper.getEdgeForLabel(instanceVertex, relationshipLabel);
      deleteHandler.deleteEdgeReference(edge, DataTypes.TypeCategory.TRAIT, false, true);

      // update the traits in entity once trait removal is successful
      traitNames.remove(traitNameToBeDeleted);
      updateTraits(instanceVertex, traitNames);
    } catch (Exception e) {
      throw new RepositoryException(e);
    }
  }

  private void updateTraits(Vertex instanceVertex, List<String> traitNames) {
    // remove the key
    instanceVertex.removeProperty(Constants.TRAIT_NAMES_PROPERTY_KEY);

    // add it back again
    for (String traitName : traitNames) {
      GraphHelper.addProperty(instanceVertex, Constants.TRAIT_NAMES_PROPERTY_KEY, traitName);
    }
    GraphHelper.setProperty(
        instanceVertex,
        Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY,
        RequestContext.get().getRequestTime());
  }

  @Override
  @GraphTransaction
  public AtlasClient.EntityResult updateEntities(ITypedReferenceableInstance... entitiesUpdated)
      throws RepositoryException {
    LOG.debug("updating entity {}", entitiesUpdated);
    try {
      TypedInstanceToGraphMapper instanceToGraphMapper =
          new TypedInstanceToGraphMapper(graphToInstanceMapper, deleteHandler);
      instanceToGraphMapper.mapTypedInstanceToGraph(
          TypedInstanceToGraphMapper.Operation.UPDATE_FULL, entitiesUpdated);
      RequestContext requestContext = RequestContext.get();
      return new AtlasClient.EntityResult(
          requestContext.getCreatedEntityIds(),
          requestContext.getUpdatedEntityIds(),
          requestContext.getDeletedEntityIds());
    } catch (AtlasException e) {
      throw new RepositoryException(e);
    }
  }

  @Override
  @GraphTransaction
  public AtlasClient.EntityResult updatePartial(ITypedReferenceableInstance entity)
      throws RepositoryException {
    LOG.debug("updating entity {}", entity);
    try {
      TypedInstanceToGraphMapper instanceToGraphMapper =
          new TypedInstanceToGraphMapper(graphToInstanceMapper, deleteHandler);
      instanceToGraphMapper.mapTypedInstanceToGraph(
          TypedInstanceToGraphMapper.Operation.UPDATE_PARTIAL, entity);
      RequestContext requestContext = RequestContext.get();
      return new AtlasClient.EntityResult(
          requestContext.getCreatedEntityIds(),
          requestContext.getUpdatedEntityIds(),
          requestContext.getDeletedEntityIds());
    } catch (AtlasException e) {
      throw new RepositoryException(e);
    }
  }

  @Override
  @GraphTransaction
  public AtlasClient.EntityResult deleteEntities(List<String> guids) throws RepositoryException {

    if (guids == null || guids.size() == 0) {
      throw new IllegalArgumentException("guids must be non-null and non-empty");
    }

    for (String guid : guids) {
      if (guid == null) {
        LOG.warn("deleteEntities: Ignoring null guid");
        continue;
      }
      try {
        Vertex instanceVertex = graphHelper.getVertexForGUID(guid);
        deleteHandler.deleteEntity(instanceVertex);
      } catch (EntityNotFoundException e) {
        // Entity does not exist - treat as non-error, since the caller
        // wanted to delete the entity and it's already gone.
        LOG.info("Deletion request ignored for non-existent entity with guid " + guid);
        continue;
      } catch (AtlasException e) {
        throw new RepositoryException(e);
      }
    }
    RequestContext requestContext = RequestContext.get();
    return new AtlasClient.EntityResult(
        requestContext.getCreatedEntityIds(),
        requestContext.getUpdatedEntityIds(),
        requestContext.getDeletedEntityIds());
  }
}