@BeforeClass
public void setUp() throws Exception {
ts = TypeSystem.getInstance();
ts.reset();
TestUtils.defineDeptEmployeeTypes(ts);
}
/** 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());
}
}