@BeforeClass
 public void setUp() throws Exception {
   ts = TypeSystem.getInstance();
   ts.reset();
   TestUtils.defineDeptEmployeeTypes(ts);
 }
Esempio n. 2
0
/** Utility class for graph operations. */
public final class GraphHelper {

  private static final Logger LOG = LoggerFactory.getLogger(GraphHelper.class);
  public static final String EDGE_LABEL_PREFIX = "__";

  private static final TypeSystem typeSystem = TypeSystem.getInstance();

  private static final GraphHelper INSTANCE =
      new GraphHelper(TitanGraphProvider.getGraphInstance());

  private TitanGraph titanGraph;

  private GraphHelper(TitanGraph titanGraph) {
    this.titanGraph = titanGraph;
  }

  public static GraphHelper getInstance() {
    return INSTANCE;
  }

  public Vertex createVertexWithIdentity(
      ITypedReferenceableInstance typedInstance, Set<String> superTypeNames) {
    final Vertex vertexWithIdentity =
        createVertexWithoutIdentity(
            typedInstance.getTypeName(), typedInstance.getId(), superTypeNames);

    // add identity
    final String guid = UUID.randomUUID().toString();
    setProperty(vertexWithIdentity, Constants.GUID_PROPERTY_KEY, guid);

    // add version information
    setProperty(vertexWithIdentity, Constants.VERSION_PROPERTY_KEY, typedInstance.getId().version);

    return vertexWithIdentity;
  }

  public Vertex createVertexWithoutIdentity(
      String typeName, Id typedInstanceId, Set<String> superTypeNames) {
    LOG.debug(
        "Creating vertex for type {} id {}",
        typeName,
        typedInstanceId != null ? typedInstanceId._getId() : null);
    final Vertex vertexWithoutIdentity = titanGraph.addVertex(null);

    // add type information
    setProperty(vertexWithoutIdentity, Constants.ENTITY_TYPE_PROPERTY_KEY, typeName);

    // add super types
    for (String superTypeName : superTypeNames) {
      addProperty(vertexWithoutIdentity, Constants.SUPER_TYPES_PROPERTY_KEY, superTypeName);
    }

    // add timestamp information
    setProperty(
        vertexWithoutIdentity, Constants.TIMESTAMP_PROPERTY_KEY, System.currentTimeMillis());

    return vertexWithoutIdentity;
  }

  public Edge addEdge(Vertex fromVertex, Vertex toVertex, String edgeLabel) {
    LOG.debug("Adding edge for {} -> label {} -> {}", fromVertex, edgeLabel, toVertex);

    return titanGraph.addEdge(null, fromVertex, toVertex, edgeLabel);
  }

  public Vertex findVertex(String propertyKey, Object value) {
    LOG.debug("Finding vertex for {}={}", propertyKey, value);

    GraphQuery query = titanGraph.query().has(propertyKey, value);
    Iterator<Vertex> results = query.vertices().iterator();
    // returning one since entityType, qualifiedName should be unique
    return results.hasNext() ? results.next() : null;
  }

  public static Iterable<Edge> getOutGoingEdgesByLabel(Vertex instanceVertex, String edgeLabel) {
    if (instanceVertex != null && edgeLabel != null) {
      return instanceVertex.getEdges(Direction.OUT, edgeLabel);
    }
    return null;
  }

  public Edge getOutGoingEdgeById(String edgeId) {
    if (edgeId != null) {
      return titanGraph.getEdge(edgeId);
    }
    return null;
  }

  public static String vertexString(final Vertex vertex) {
    StringBuilder properties = new StringBuilder();
    for (String propertyKey : vertex.getPropertyKeys()) {
      properties
          .append(propertyKey)
          .append("=")
          .append(vertex.getProperty(propertyKey).toString())
          .append(", ");
    }

    return "v[" + vertex.getId() + "], Properties[" + properties + "]";
  }

  public static String edgeString(final Edge edge) {
    return "e["
        + edge.getLabel()
        + "], ["
        + edge.getVertex(Direction.OUT)
        + " -> "
        + edge.getLabel()
        + " -> "
        + edge.getVertex(Direction.IN)
        + "]";
  }

  public static void setProperty(Vertex vertex, String propertyName, Object value) {
    LOG.debug("Setting property {} = \"{}\" to vertex {}", propertyName, value, vertex);
    Object existValue = vertex.getProperty(propertyName);
    if (value == null || (value instanceof Collection && ((Collection) value).isEmpty())) {
      if (existValue != null) {
        LOG.info("Removing property - {} value from vertex {}", propertyName, vertex);
        vertex.removeProperty(propertyName);
      }
    } else {
      if (!value.equals(existValue)) {
        vertex.setProperty(propertyName, value);
      }
    }
  }

  public static void addProperty(Vertex vertex, String propertyName, Object value) {
    LOG.debug("Setting property {} = \"{}\" to vertex {}", propertyName, value, vertex);
    ((TitanVertex) vertex).addProperty(propertyName, value);
  }

  public Edge removeRelation(String edgeId, boolean cascade) {
    LOG.debug("Removing edge with id {}", edgeId);
    final Edge edge = titanGraph.getEdge(edgeId);
    titanGraph.removeEdge(edge);
    LOG.info("Removed edge {}", edge);
    if (cascade) {
      Vertex referredVertex = edge.getVertex(Direction.IN);
      titanGraph.removeVertex(referredVertex);
      LOG.info("Removed vertex {}", referredVertex);
    }
    return edge;
  }

  public Vertex getVertexForGUID(String guid) throws EntityNotFoundException {
    return getVertexForProperty(Constants.GUID_PROPERTY_KEY, guid);
  }

  public Vertex getVertexForProperty(String propertyKey, Object value)
      throws EntityNotFoundException {
    Vertex instanceVertex = findVertex(propertyKey, value);
    if (instanceVertex == null) {
      LOG.debug("Could not find a vertex with {}={}", propertyKey, value);
      throw new EntityNotFoundException(
          "Could not find an entity in the repository with " + propertyKey + "=" + value);
    } else {
      LOG.debug("Found a vertex {} with {}={}", instanceVertex, propertyKey, value);
    }

    return instanceVertex;
  }

  public static String getQualifiedFieldName(
      ITypedInstance typedInstance, AttributeInfo attributeInfo) throws AtlasException {
    IDataType dataType = typeSystem.getDataType(IDataType.class, typedInstance.getTypeName());
    return getQualifiedFieldName(dataType, attributeInfo.name);
  }

  public static String getQualifiedFieldName(IDataType dataType, String attributeName)
      throws AtlasException {
    return dataType.getTypeCategory() == DataTypes.TypeCategory.STRUCT
        ? dataType.getName() + "." + attributeName
        // else class or trait
        : ((HierarchicalType) dataType).getQualifiedName(attributeName);
  }

  public static String getTraitLabel(String typeName, String attrName) {
    return typeName + "." + attrName;
  }

  public static List<String> getTraitNames(Vertex entityVertex) {
    ArrayList<String> traits = new ArrayList<>();
    for (TitanProperty property :
        ((TitanVertex) entityVertex).getProperties(Constants.TRAIT_NAMES_PROPERTY_KEY)) {
      traits.add((String) property.getValue());
    }

    return traits;
  }

  public static String getEdgeLabel(ITypedInstance typedInstance, AttributeInfo aInfo)
      throws AtlasException {
    IDataType dataType = typeSystem.getDataType(IDataType.class, typedInstance.getTypeName());
    return getEdgeLabel(dataType, aInfo);
  }

  public static String getEdgeLabel(IDataType dataType, AttributeInfo aInfo) throws AtlasException {
    return GraphHelper.EDGE_LABEL_PREFIX + getQualifiedFieldName(dataType, aInfo.name);
  }

  public static Id getIdFromVertex(String dataTypeName, Vertex vertex) {
    return new Id(
        vertex.<String>getProperty(Constants.GUID_PROPERTY_KEY),
        vertex.<Integer>getProperty(Constants.VERSION_PROPERTY_KEY),
        dataTypeName);
  }

  public static String getTypeName(Vertex instanceVertex) {
    return instanceVertex.getProperty(Constants.ENTITY_TYPE_PROPERTY_KEY);
  }

  /**
   * For the given type, finds an unique attribute and checks if there is an existing instance with
   * the same unique value
   *
   * @param classType
   * @param instance
   * @return
   * @throws AtlasException
   */
  public Vertex getVertexForInstanceByUniqueAttribute(
      ClassType classType, IReferenceableInstance instance) throws AtlasException {

    Vertex result = null;
    for (AttributeInfo attributeInfo : classType.fieldMapping().fields.values()) {
      if (attributeInfo.isUnique) {
        String propertyKey = getQualifiedFieldName(classType, attributeInfo.name);
        try {
          result = getVertexForProperty(propertyKey, instance.get(attributeInfo.name));
          LOG.debug(
              "Found vertex by unique attribute : "
                  + propertyKey
                  + "="
                  + instance.get(attributeInfo.name));
        } catch (EntityNotFoundException e) {
          // Its ok if there is no entity with the same unique value
        }
      }
    }

    return result;
  }

  public static void dumpToLog(final Graph graph) {
    LOG.debug("*******************Graph Dump****************************");
    LOG.debug("Vertices of {}", graph);
    for (Vertex vertex : graph.getVertices()) {
      LOG.debug(vertexString(vertex));
    }

    LOG.debug("Edges of {}", graph);
    for (Edge edge : graph.getEdges()) {
      LOG.debug(edgeString(edge));
    }
    LOG.debug("*******************Graph Dump****************************");
  }
}
/** 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());
  }
}