/**
* 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);
}
}
private void verifyRestoredClassType(TypesDef types, String typeName) throws AtlasException {
boolean clsTypeFound = false;
List<HierarchicalTypeDefinition<ClassType>> classTypes = types.classTypesAsJavaList();
for (HierarchicalTypeDefinition<ClassType> classType : classTypes) {
if (classType.typeName.equals(typeName)) {
ClassType expectedType = ts.getDataType(ClassType.class, classType.typeName);
Assert.assertEquals(
expectedType.immediateAttrs.size(), classType.attributeDefinitions.length);
Assert.assertEquals(expectedType.superTypes.size(), classType.superTypes.size());
Assert.assertEquals(classType.typeDescription, classType.typeName + DESCRIPTION);
clsTypeFound = true;
}
}
Assert.assertTrue(clsTypeFound, typeName + " type not restored");
}
@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);
}
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 getQualifiedFieldName(
ITypedInstance typedInstance, AttributeInfo attributeInfo) throws AtlasException {
IDataType dataType = typeSystem.getDataType(IDataType.class, typedInstance.getTypeName());
return getQualifiedFieldName(dataType, attributeInfo.name);
}
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());
}
/**
* 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);
}
}
/**
* 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);
}
@Test(dependsOnMethods = "testStore")
public void testTypeUpdate() throws Exception {
// Add enum value
String _description = "_description_updated";
EnumTypeDefinition orgLevelEnum =
new EnumTypeDefinition(
"OrgLevel",
"OrgLevel" + _description,
new EnumValue("L1", 1),
new EnumValue("L2", 2),
new EnumValue("L3", 3));
// Add attribute
StructTypeDefinition addressDetails =
createStructTypeDef(
"Address",
createRequiredAttrDef("street", DataTypes.STRING_TYPE),
createRequiredAttrDef("city", DataTypes.STRING_TYPE),
createOptionalAttrDef("state", DataTypes.STRING_TYPE));
HierarchicalTypeDefinition<ClassType> deptTypeDef =
createClassTypeDef(
"Department",
"Department" + _description,
ImmutableSet.<String>of(),
createRequiredAttrDef("name", DataTypes.STRING_TYPE),
new AttributeDefinition(
"employees",
String.format("array<%s>", "Person"),
Multiplicity.OPTIONAL,
true,
"department"));
TypesDef typesDef =
TypesUtil.getTypesDef(
ImmutableList.of(orgLevelEnum),
ImmutableList.of(addressDetails),
ImmutableList.<HierarchicalTypeDefinition<TraitType>>of(),
ImmutableList.of(deptTypeDef));
Map<String, IDataType> typesAdded = ts.updateTypes(typesDef);
typeStore.store(ts, ImmutableList.copyOf(typesAdded.keySet()));
verifyEdges();
// Validate the updated types
TypesDef types = typeStore.restore();
ts.reset();
ts.defineTypes(types);
// Assert new enum value
EnumType orgLevel = ts.getDataType(EnumType.class, orgLevelEnum.name);
Assert.assertEquals(orgLevel.name, orgLevelEnum.name);
Assert.assertEquals(orgLevel.description, orgLevelEnum.description);
Assert.assertEquals(orgLevel.values().size(), orgLevelEnum.enumValues.length);
Assert.assertEquals(orgLevel.fromValue("L3").ordinal, 3);
// Assert new attribute
StructType addressType = ts.getDataType(StructType.class, addressDetails.typeName);
Assert.assertEquals(addressType.numFields, 3);
Assert.assertEquals(
addressType.fieldMapping.fields.get("state").dataType(), DataTypes.STRING_TYPE);
// Updating the definition again shouldn't add another edge
typesDef =
TypesUtil.getTypesDef(
ImmutableList.<EnumTypeDefinition>of(),
ImmutableList.<StructTypeDefinition>of(),
ImmutableList.<HierarchicalTypeDefinition<TraitType>>of(),
ImmutableList.of(deptTypeDef));
typesAdded = ts.updateTypes(typesDef);
typeStore.store(ts, ImmutableList.copyOf(typesAdded.keySet()));
verifyEdges();
}
