Java Code Examples for org.apache.atlas.model.instance.AtlasEntity#setRelationshipAttribute()

AtlasEntity createProcess(String name, String description, String user, List<AtlasEntity> inputs, List<AtlasEntity> outputs,
        String queryText, String queryPlan, String queryId, String queryGraph, String... classificationNames) throws Exception {

    AtlasEntity entity = new AtlasEntity(LOAD_PROCESS_TYPE);

    // set attributes
    entity.setAttribute("name", name);
    entity.setAttribute(REFERENCEABLE_ATTRIBUTE_NAME, name + CLUSTER_SUFFIX);
    entity.setAttribute("description", description);
    entity.setAttribute("user", user);
    entity.setAttribute("startTime", System.currentTimeMillis());
    entity.setAttribute("endTime", System.currentTimeMillis() + 10000);
    entity.setAttribute("queryText", queryText);
    entity.setAttribute("queryPlan", queryPlan);
    entity.setAttribute("queryId", queryId);
    entity.setAttribute("queryGraph", queryGraph);

    // set relationship attributes
    entity.setRelationshipAttribute("inputs", toAtlasRelatedObjectIds(inputs));
    entity.setRelationshipAttribute("outputs", toAtlasRelatedObjectIds(outputs));

    // set classifications
    entity.setClassifications(toAtlasClassifications(classificationNames));

    return createInstance(entity);
}
 

private void addTopologyOutputs(StormTopology stormTopology, String topologyOwner, Map stormConf, AtlasEntity topology, AtlasEntityExtInfo entityExtInfo) {
    List<AtlasEntity> outputs   = new ArrayList<>();
    Map<String, Bolt> bolts     = stormTopology.get_bolts();
    Set<String>       boltNames = StormTopologyUtil.getTerminalUserBoltNames(stormTopology);

    for (String boltName : boltNames) {
        Serializable instance = Utils.javaDeserialize(bolts.get(boltName).get_bolt_object().get_serialized_java(), Serializable.class);
        String       dsType   = instance.getClass().getSimpleName();
        AtlasEntity  dsEntity = addDataSet(dsType, topologyOwner, instance, stormConf, entityExtInfo);

        if (dsEntity != null) {
            outputs.add(dsEntity);
        }
    }

    topology.setRelationshipAttribute("outputs", AtlasTypeUtil.getAtlasRelatedObjectIds(outputs, RELATIONSHIP_PROCESS_DATASET_OUTPUTS));
}
 

private List<AtlasEntity> toColumns(List<FieldSchema> schemaList, AtlasEntity table, String relationshipType) throws AtlasHookException {
    List<AtlasEntity> ret = new ArrayList<>();

    int columnPosition = 0;
    for (FieldSchema fs : schemaList) {
        LOG.debug("Processing field {}", fs);

        AtlasEntity column = new AtlasEntity(HiveDataTypes.HIVE_COLUMN.getName());

        column.setRelationshipAttribute(ATTRIBUTE_TABLE, AtlasTypeUtil.getAtlasRelatedObjectId(table, relationshipType));
        column.setAttribute(ATTRIBUTE_QUALIFIED_NAME, getColumnQualifiedName((String) table.getAttribute(ATTRIBUTE_QUALIFIED_NAME), fs.getName()));
        column.setAttribute(ATTRIBUTE_NAME, fs.getName());
        column.setAttribute(ATTRIBUTE_OWNER, table.getAttribute(ATTRIBUTE_OWNER));
        column.setAttribute(ATTRIBUTE_COL_TYPE, fs.getType());
        column.setAttribute(ATTRIBUTE_COL_POSITION, columnPosition++);
        column.setAttribute(ATTRIBUTE_COMMENT, fs.getComment());

        ret.add(column);
    }
    return ret;
}
 

private void compactAttributes(AtlasEntity entity, AtlasEntityType entityType) {
    if (entity != null) {
        for (String attrName : entityType.getRelationshipAttributes().keySet()) {
            if (entity.hasAttribute(attrName)) { // relationship attribute is present in 'attributes'
                Object attrValue = entity.removeAttribute(attrName);

                if (attrValue != null) {
                    // if the attribute doesn't exist in relationshipAttributes, add it
                    Object relationshipAttrValue = entity.getRelationshipAttribute(attrName);

                    if (relationshipAttrValue == null) {
                        entity.setRelationshipAttribute(attrName, attrValue);

                        if (LOG.isDebugEnabled()) {
                            LOG.debug("moved attribute {}.{} from attributes to relationshipAttributes", entityType.getTypeName(), attrName);
                        }
                    } else {
                        if (LOG.isDebugEnabled()) {
                            LOG.debug("attribute {}.{} is present in attributes and relationshipAttributes. Removed from attributes", entityType.getTypeName(), attrName);
                        }
                    }
                }
            }
        }
    }
}
 

@Override
public void preprocess(AtlasEntity entity, PreprocessorContext context) {
    super.preprocess(entity, context);

    // try auto-fix when 'db' attribute is not present in relationshipAttribute & attributes
    Object db = entity.getRelationshipAttribute(ATTRIBUTE_DB);

    if (db == null) {
        db = entity.getAttribute(ATTRIBUTE_DB);
    }

    if (db == null) {
        String dbQualifiedName = getDbQualifiedName(entity);

        if (dbQualifiedName != null) {
            AtlasObjectId dbId = new AtlasObjectId(TYPE_RDBMS_DB, Collections.singletonMap(ATTRIBUTE_QUALIFIED_NAME, dbQualifiedName));

            LOG.info("missing attribute {}.{} is set to {}", TYPE_RDBMS_TABLE, ATTRIBUTE_DB, dbId);

            entity.setRelationshipAttribute(ATTRIBUTE_DB, dbId);
        }
    }
}
 

AtlasEntity createView(String name, AtlasEntity database, List<AtlasEntity> inputTables, String... classificationNames) throws Exception {
    AtlasEntity entity = new AtlasEntity(VIEW_TYPE);

    // set attributes
    entity.setAttribute("name", name);
    entity.setAttribute(REFERENCEABLE_ATTRIBUTE_NAME, name + CLUSTER_SUFFIX);

    // set relationship attributes
    entity.setRelationshipAttribute("db", toAtlasRelatedObjectId(database));
    entity.setRelationshipAttribute("inputTables", toAtlasRelatedObjectIds(inputTables));

    // set classifications
    entity.setClassifications(toAtlasClassifications(classificationNames));

    return createInstance(entity);
}
 

private EntityMutationResponse createColumn(AtlasEntity tblEntity) throws AtlasBaseException {
    AtlasEntity ret = new AtlasEntity(ENTITY_TYPE_COL);
    String name = "new_column";

    ret.setAttribute("name", name);
    ret.setAttribute(REFERENCEABLE_ATTRIBUTE_NAME, name + REPL_FROM);
    ret.setAttribute("type", "int");
    ret.setAttribute("comment", name);

    ret.setRelationshipAttribute("table", toAtlasRelatedObjectId(tblEntity));
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(ret), false);

    return response;
}
 

private void addTopologyInputs(Map<String, SpoutSpec> spouts, Map stormConf, String topologyOwner, AtlasEntity topology, AtlasEntityExtInfo entityExtInfo) {
    List<AtlasEntity> inputs = new ArrayList<>();

    for (Map.Entry<String, SpoutSpec> entry : spouts.entrySet()) {
        Serializable instance = Utils.javaDeserialize(entry.getValue().get_spout_object().get_serialized_java(), Serializable.class);
        String       dsType   = instance.getClass().getSimpleName();
        AtlasEntity  dsEntity = addDataSet(dsType, topologyOwner, instance, stormConf, entityExtInfo);

        if (dsEntity != null) {
            inputs.add(dsEntity);
        }
    }

    topology.setRelationshipAttribute("inputs", AtlasTypeUtil.getAtlasRelatedObjectIds(inputs, RELATIONSHIP_DATASET_PROCESS_INPUTS));
}
 

private AtlasEntity toHiveTableEntity(AtlasEntity entHiveDb, String tableName) {
    AtlasEntity entHiveTable  = new AtlasEntity(HiveDataTypes.HIVE_TABLE.getName());
    String      qualifiedName = HiveMetaStoreBridge.getTableQualifiedName((String)entHiveDb.getAttribute(AtlasConstants.CLUSTER_NAME_ATTRIBUTE), (String)entHiveDb.getAttribute(AtlasClient.NAME), tableName);

    entHiveTable.setAttribute(AtlasClient.NAME, tableName.toLowerCase());
    entHiveTable.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, qualifiedName);
    entHiveTable.setRelationshipAttribute(ATTRIBUTE_DB, AtlasTypeUtil.getAtlasRelatedObjectId(entHiveDb, RELATIONSHIP_HIVE_TABLE_DB));

    return entHiveTable;
}
 

private AtlasEntity toSqoopProcessEntity(AtlasEntity entDbStore, AtlasEntity entHiveDb, AtlasEntity entHiveTable,
                                         SqoopJobDataPublisher.Data data, String metadataNamespace) {
    AtlasEntity         entProcess       = new AtlasEntity(SqoopDataTypes.SQOOP_PROCESS.getName());
    String              sqoopProcessName = getSqoopProcessName(data, metadataNamespace);
    Map<String, String> sqoopOptionsMap  = new HashMap<>();
    Properties          options          = data.getOptions();

    for (Object k : options.keySet()) {
        sqoopOptionsMap.put((String)k, (String) options.get(k));
    }

    entProcess.setAttribute(AtlasClient.NAME, sqoopProcessName);
    entProcess.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, sqoopProcessName);
    entProcess.setAttribute(SqoopHook.OPERATION, data.getOperation());

    List<AtlasObjectId> sqoopObjects = Collections.singletonList(AtlasTypeUtil.getAtlasObjectId(entDbStore));
    List<AtlasObjectId> hiveObjects  = Collections.singletonList(AtlasTypeUtil.getAtlasObjectId(entHiveTable != null ? entHiveTable : entHiveDb));

    if (isImportOperation(data)) {
        entProcess.setRelationshipAttribute(SqoopHook.INPUTS, AtlasTypeUtil.getAtlasRelatedObjectIdList(sqoopObjects, RELATIONSHIP_DATASET_PROCESS_INPUTS));
        entProcess.setRelationshipAttribute(SqoopHook.OUTPUTS, AtlasTypeUtil.getAtlasRelatedObjectIdList(hiveObjects, RELATIONSHIP_PROCESS_DATASET_OUTPUTS));
    } else {
        entProcess.setRelationshipAttribute(SqoopHook.INPUTS, AtlasTypeUtil.getAtlasRelatedObjectIdList(hiveObjects, RELATIONSHIP_DATASET_PROCESS_INPUTS));
        entProcess.setRelationshipAttribute(SqoopHook.OUTPUTS, AtlasTypeUtil.getAtlasRelatedObjectIdList(sqoopObjects, RELATIONSHIP_PROCESS_DATASET_OUTPUTS));
    }

    entProcess.setAttribute(SqoopHook.USER, data.getUser());
    entProcess.setAttribute(SqoopHook.START_TIME, new Date(data.getStartTime()));
    entProcess.setAttribute(SqoopHook.END_TIME, new Date(data.getEndTime()));
    entProcess.setAttribute(SqoopHook.CMD_LINE_OPTS, sqoopOptionsMap);

    return entProcess;
}
 

private void createOrUpdateTableInstance(HBaseOperationContext hbaseOperationContext) {
    AtlasEntity       nameSpace      = buildNameSpace(hbaseOperationContext);
    AtlasEntity       table          = buildTable(hbaseOperationContext, nameSpace);
    List<AtlasEntity> columnFamilies = buildColumnFamilies(hbaseOperationContext, nameSpace, table);

    table.setRelationshipAttribute(ATTR_COLUMNFAMILIES, AtlasTypeUtil.getAtlasRelatedObjectIds(columnFamilies, RELATIONSHIP_HBASE_TABLE_COLUMN_FAMILIES));

    AtlasEntitiesWithExtInfo entities = new AtlasEntitiesWithExtInfo(table);

    entities.addReferredEntity(nameSpace);

    if (CollectionUtils.isNotEmpty(columnFamilies)) {
        for (AtlasEntity columnFamily : columnFamilies) {
            entities.addReferredEntity(columnFamily);
        }
    }

    switch (hbaseOperationContext.getOperation()) {
        case CREATE_TABLE:
            LOG.info("Create Table {}", table.getAttribute(REFERENCEABLE_ATTRIBUTE_NAME));

            hbaseOperationContext.addMessage(new EntityCreateRequestV2(hbaseOperationContext.getUser(), entities));
            break;

        case ALTER_TABLE:
            LOG.info("Modify Table {}", table.getAttribute(REFERENCEABLE_ATTRIBUTE_NAME));

            hbaseOperationContext.addMessage(new EntityUpdateRequestV2(hbaseOperationContext.getUser(), entities));
            break;
    }
}
 

public void removeRefAttributeAndRegisterToMove(AtlasEntity entity, String attrName, String relationshipType, String refAttrName) {
    Object attrVal = entity.removeAttribute(attrName);

    if (attrVal != null) {
        AtlasRelatedObjectId entityId = null;
        Set<String>          guids    = new HashSet<>();

        collectGuids(attrVal, guids);

        // removed attrVal might have elements removed (e.g. removed column); to handle this case register the entity for partial update
        addToPostUpdate(entity, attrName, attrVal);

        for (String guid : guids) {
            AtlasEntity refEntity = getEntity(guid);

            if (refEntity != null) {
                Object refAttr = null;

                if (refEntity.hasRelationshipAttribute(refAttrName)) {
                    refAttr = refEntity.getRelationshipAttribute(refAttrName);
                } else if (refEntity.hasAttribute(refAttrName)) {
                    refAttr = refEntity.getAttribute(refAttrName);
                } else {
                    if (entityId == null) {
                        entityId = AtlasTypeUtil.toAtlasRelatedObjectId(entity, typeRegistry);
                    }

                    refAttr = entityId;
                }

                if (refAttr != null) {
                    refAttr = setRelationshipType(refAttr, relationshipType);
                }

                if (refAttr != null) {
                    refEntity.setRelationshipAttribute(refAttrName, refAttr);
                }

                addToReferredEntitiesToMove(guid);
            }
        }
    }
}
 

private Object mapVertexToRelationshipAttribute(AtlasVertex entityVertex, AtlasEntityType entityType, String attributeName, AtlasEntity entity, AtlasEntityExtInfo entityExtInfo, boolean isMinExtInfo) throws AtlasBaseException {
    Object                ret                  = null;
    String                relationshipTypeName = graphHelper.getRelationshipTypeName(entityVertex, entityType, attributeName);
    AtlasRelationshipType relationshipType     = relationshipTypeName != null ? typeRegistry.getRelationshipTypeByName(relationshipTypeName) : null;

    if (relationshipType == null) {
        throw new AtlasBaseException(AtlasErrorCode.RELATIONSHIPDEF_INVALID, "relationshipDef is null");
    }

    AtlasAttribute          attribute       = entityType.getRelationshipAttribute(attributeName, relationshipTypeName);
    AtlasRelationshipDef    relationshipDef = relationshipType.getRelationshipDef();
    AtlasRelationshipEndDef endDef1         = relationshipDef.getEndDef1();
    AtlasRelationshipEndDef endDef2         = relationshipDef.getEndDef2();
    AtlasEntityType         endDef1Type     = typeRegistry.getEntityTypeByName(endDef1.getType());
    AtlasEntityType         endDef2Type     = typeRegistry.getEntityTypeByName(endDef2.getType());
    AtlasRelationshipEndDef attributeEndDef = null;

    if (endDef1Type.isTypeOrSuperTypeOf(entityType.getTypeName()) && StringUtils.equals(endDef1.getName(), attributeName)) {
        attributeEndDef = endDef1;
    } else if (endDef2Type.isTypeOrSuperTypeOf(entityType.getTypeName()) && StringUtils.equals(endDef2.getName(), attributeName)) {
        attributeEndDef = endDef2;
    }

    if (attributeEndDef == null) {
        throw new AtlasBaseException(AtlasErrorCode.RELATIONSHIPDEF_INVALID, relationshipDef.toString());
    }

    switch (attributeEndDef.getCardinality()) {
        case SINGLE:
            ret = mapRelatedVertexToObjectId(entityVertex, attribute, entityExtInfo, isMinExtInfo);
            break;

        case LIST:
        case SET:
            ret = mapRelationshipArrayAttribute(entityVertex, attribute, entityExtInfo, isMinExtInfo);
            break;
    }

    if (ret != null) {
        entity.setRelationshipAttribute(attributeName, ret);

        if (attributeEndDef.getIsLegacyAttribute() && !entity.hasAttribute(attributeName)) {
            entity.setAttribute(attributeName, toLegacyAttribute(ret));
        }
    }

    return ret;
}
 

@Test
public void testRelationshipAttributeUpdate_NonComposite_ManyToOne() throws Exception {
    AtlasEntity a1 = new AtlasEntity("A");
    a1.setAttribute(NAME, "a1_name");

    AtlasEntity a2 = new AtlasEntity("A");
    a2.setAttribute(NAME, "a2_name");

    AtlasEntity a3 = new AtlasEntity("A");
    a3.setAttribute(NAME, "a3_name");

    AtlasEntity b = new AtlasEntity("B");
    b.setAttribute(NAME, "b_name");

    AtlasEntitiesWithExtInfo entitiesWithExtInfo = new AtlasEntitiesWithExtInfo();
    entitiesWithExtInfo.addEntity(a1);
    entitiesWithExtInfo.addEntity(a2);
    entitiesWithExtInfo.addEntity(a3);
    entitiesWithExtInfo.addEntity(b);
    entityStore.createOrUpdate(new AtlasEntityStream(entitiesWithExtInfo) , false);

    AtlasEntity bPartialUpdate = new AtlasEntity("B");
    bPartialUpdate.setRelationshipAttribute("manyA", ImmutableList.of(getAtlasObjectId(a1), getAtlasObjectId(a2)));

    init();
    EntityMutationResponse response = entityStore.updateByUniqueAttributes(typeRegistry.getEntityTypeByName("B"),
                                                                           Collections.singletonMap(NAME, b.getAttribute(NAME)),
                                                                           new AtlasEntityWithExtInfo(bPartialUpdate));
    // Verify 3 entities were updated:
    // * set b.manyA reference to a1 and a2
    // * set inverse a1.oneB reference to b
    // * set inverse a2.oneB reference to b
    assertEquals(response.getPartialUpdatedEntities().size(), 3);
    AtlasEntitiesWithExtInfo updatedEntities = entityStore.getByIds(ImmutableList.of(a1.getGuid(), a2.getGuid(), b.getGuid()));

    AtlasEntity a1Entity = updatedEntities.getEntity(a1.getGuid());
    verifyRelationshipAttributeValue(a1Entity, "oneB", b.getGuid());

    AtlasEntity a2Entity = updatedEntities.getEntity(a2.getGuid());
    verifyRelationshipAttributeValue(a2Entity, "oneB", b.getGuid());

    AtlasEntity bEntity = updatedEntities.getEntity(b.getGuid());
    verifyRelationshipAttributeList(bEntity, "manyA", ImmutableList.of(getAtlasObjectId(a1), getAtlasObjectId(a2)));


    bPartialUpdate.setRelationshipAttribute("manyA", ImmutableList.of(getAtlasObjectId(a3)));
    init();
    response = entityStore.updateByUniqueAttributes(typeRegistry.getEntityTypeByName("B"),
                                                    Collections.singletonMap(NAME, b.getAttribute(NAME)),
                                                    new AtlasEntityWithExtInfo(bPartialUpdate));
    // Verify 4 entities were updated:
    // * set b.manyA reference to a3
    // * set inverse a3.oneB reference to b
    // * disconnect inverse a1.oneB reference to b
    // * disconnect inverse a2.oneB reference to b
    assertEquals(response.getPartialUpdatedEntities().size(), 4);
    init();

    updatedEntities = entityStore.getByIds(ImmutableList.of(a1.getGuid(), a2.getGuid(), a3.getGuid(), b.getGuid()));
    a1Entity        = updatedEntities.getEntity(a1.getGuid());
    a2Entity        = updatedEntities.getEntity(a2.getGuid());
    bEntity         = updatedEntities.getEntity(b.getGuid());

    AtlasEntity a3Entity = updatedEntities.getEntity(a3.getGuid());
    verifyRelationshipAttributeValue(a3Entity, "oneB", b.getGuid());

    verifyRelationshipAttributeUpdate_NonComposite_ManyToOne(a1Entity, a2Entity, a3Entity, bEntity);
}
 

private AtlasEntity toStorageDescEntity(StorageDescriptor storageDesc, String tableQualifiedName, String sdQualifiedName, AtlasObjectId tableId ) throws AtlasHookException {
    AtlasEntity ret = new AtlasEntity(HiveDataTypes.HIVE_STORAGEDESC.getName());

    ret.setRelationshipAttribute(ATTRIBUTE_TABLE, AtlasTypeUtil.getAtlasRelatedObjectId(tableId, RELATIONSHIP_HIVE_TABLE_STORAGE_DESC));
    ret.setAttribute(ATTRIBUTE_QUALIFIED_NAME, sdQualifiedName);
    ret.setAttribute(ATTRIBUTE_PARAMETERS, storageDesc.getParameters());
    ret.setAttribute(ATTRIBUTE_LOCATION, HdfsNameServiceResolver.getPathWithNameServiceID(storageDesc.getLocation()));
    ret.setAttribute(ATTRIBUTE_INPUT_FORMAT, storageDesc.getInputFormat());
    ret.setAttribute(ATTRIBUTE_OUTPUT_FORMAT, storageDesc.getOutputFormat());
    ret.setAttribute(ATTRIBUTE_COMPRESSED, storageDesc.isCompressed());
    ret.setAttribute(ATTRIBUTE_NUM_BUCKETS, storageDesc.getNumBuckets());
    ret.setAttribute(ATTRIBUTE_STORED_AS_SUB_DIRECTORIES, storageDesc.isStoredAsSubDirectories());

    if (storageDesc.getBucketCols().size() > 0) {
        ret.setAttribute(ATTRIBUTE_BUCKET_COLS, storageDesc.getBucketCols());
    }

    if (storageDesc.getSerdeInfo() != null) {
        SerDeInfo serdeInfo = storageDesc.getSerdeInfo();

        LOG.debug("serdeInfo = {}", serdeInfo);
        // SkewedInfo skewedInfo = storageDesc.getSkewedInfo();

        AtlasStruct serdeInfoStruct = new AtlasStruct(HiveDataTypes.HIVE_SERDE.getName());

        serdeInfoStruct.setAttribute(ATTRIBUTE_NAME, serdeInfo.getName());
        serdeInfoStruct.setAttribute(ATTRIBUTE_SERIALIZATION_LIB, serdeInfo.getSerializationLib());
        serdeInfoStruct.setAttribute(ATTRIBUTE_PARAMETERS, serdeInfo.getParameters());

        ret.setAttribute(ATTRIBUTE_SERDE_INFO, serdeInfoStruct);
    }

    if (CollectionUtils.isNotEmpty(storageDesc.getSortCols())) {
        List<AtlasStruct> sortColsStruct = new ArrayList<>();

        for (Order sortcol : storageDesc.getSortCols()) {
            String hiveOrderName = HiveDataTypes.HIVE_ORDER.getName();
            AtlasStruct colStruct = new AtlasStruct(hiveOrderName);
            colStruct.setAttribute("col", sortcol.getCol());
            colStruct.setAttribute("order", sortcol.getOrder());

            sortColsStruct.add(colStruct);
        }

        ret.setAttribute(ATTRIBUTE_SORT_COLS, sortColsStruct);
    }

    return ret;
}
 

@Test
public void testRelationshipAttributeUpdate_NonComposite_ManyToMany() throws Exception {
    AtlasEntity a1 = new AtlasEntity("A");
    a1.setAttribute(NAME, "a1_name");

    AtlasEntity a2 = new AtlasEntity("A");
    a2.setAttribute(NAME, "a2_name");

    AtlasEntity a3 = new AtlasEntity("A");
    a3.setAttribute(NAME, "a3_name");

    AtlasEntity b1 = new AtlasEntity("B");
    b1.setAttribute(NAME, "b1_name");

    AtlasEntity b2 = new AtlasEntity("B");
    b2.setAttribute(NAME, "b2_name");

    AtlasEntitiesWithExtInfo entitiesWithExtInfo = new AtlasEntitiesWithExtInfo();
    entitiesWithExtInfo.addEntity(a1);
    entitiesWithExtInfo.addEntity(a2);
    entitiesWithExtInfo.addEntity(a3);
    entitiesWithExtInfo.addEntity(b1);
    entitiesWithExtInfo.addEntity(b2);
    entityStore.createOrUpdate(new AtlasEntityStream(entitiesWithExtInfo) , false);

    AtlasEntity b1PartialUpdate = new AtlasEntity("B");
    b1PartialUpdate.setRelationshipAttribute("manyToManyA", ImmutableList.of(getAtlasObjectId(a1), getAtlasObjectId(a2)));

    init();
    EntityMutationResponse response = entityStore.updateByUniqueAttributes(typeRegistry.getEntityTypeByName("B"),
                                                                           Collections.singletonMap(NAME, b1.getAttribute(NAME)),
                                                                           new AtlasEntityWithExtInfo(b1PartialUpdate));

    List<AtlasEntityHeader> updatedEntityHeaders = response.getPartialUpdatedEntities();
    assertEquals(updatedEntityHeaders.size(), 3);

    AtlasEntitiesWithExtInfo updatedEntities = entityStore.getByIds(ImmutableList.of(a1.getGuid(), a2.getGuid(), b1.getGuid()));

    AtlasEntity b1Entity = updatedEntities.getEntity(b1.getGuid());
    verifyRelationshipAttributeList(b1Entity, "manyToManyA", ImmutableList.of(getAtlasObjectId(a1), getAtlasObjectId(a2)));

    AtlasEntity a1Entity = updatedEntities.getEntity(a1.getGuid());
    verifyRelationshipAttributeList(a1Entity, "manyB", ImmutableList.of(getAtlasObjectId(b1)));

    AtlasEntity a2Entity = updatedEntities.getEntity(a2.getGuid());
    verifyRelationshipAttributeList(a2Entity, "manyB", ImmutableList.of(getAtlasObjectId(b1)));
}
 

@VisibleForTesting
public int importTable(AtlasEntity dbEntity, String databaseName, String tableName, final boolean failOnError) throws Exception {
    try {
        Table                  table       = hiveClient.getTable(databaseName, tableName);
        AtlasEntityWithExtInfo tableEntity = registerTable(dbEntity, table);

        if (table.getTableType() == TableType.EXTERNAL_TABLE) {
            String                 processQualifiedName = getTableProcessQualifiedName(metadataNamespace, table);
            AtlasEntityWithExtInfo processEntity        = findProcessEntity(processQualifiedName);

            if (processEntity == null) {
                String      tableLocation = isConvertHdfsPathToLowerCase() ? lower(table.getDataLocation().toString()) : table.getDataLocation().toString();
                String      query         = getCreateTableString(table, tableLocation);
                AtlasEntity pathInst      = toHdfsPathEntity(tableLocation);
                AtlasEntity tableInst     = tableEntity.getEntity();
                AtlasEntity processInst   = new AtlasEntity(HiveDataTypes.HIVE_PROCESS.getName());
                long        now           = System.currentTimeMillis();

                processInst.setAttribute(ATTRIBUTE_QUALIFIED_NAME, processQualifiedName);
                processInst.setAttribute(ATTRIBUTE_NAME, query);
                processInst.setAttribute(ATTRIBUTE_CLUSTER_NAME, metadataNamespace);
                processInst.setRelationshipAttribute(ATTRIBUTE_INPUTS, Collections.singletonList(AtlasTypeUtil.getAtlasRelatedObjectId(pathInst, RELATIONSHIP_DATASET_PROCESS_INPUTS)));
                processInst.setRelationshipAttribute(ATTRIBUTE_OUTPUTS, Collections.singletonList(AtlasTypeUtil.getAtlasRelatedObjectId(tableInst, RELATIONSHIP_PROCESS_DATASET_OUTPUTS)));
                processInst.setAttribute(ATTRIBUTE_USER_NAME, table.getOwner());
                processInst.setAttribute(ATTRIBUTE_START_TIME, now);
                processInst.setAttribute(ATTRIBUTE_END_TIME, now);
                processInst.setAttribute(ATTRIBUTE_OPERATION_TYPE, "CREATETABLE");
                processInst.setAttribute(ATTRIBUTE_QUERY_TEXT, query);
                processInst.setAttribute(ATTRIBUTE_QUERY_ID, query);
                processInst.setAttribute(ATTRIBUTE_QUERY_PLAN, "{}");
                processInst.setAttribute(ATTRIBUTE_RECENT_QUERIES, Collections.singletonList(query));

                AtlasEntitiesWithExtInfo createTableProcess = new AtlasEntitiesWithExtInfo();

                createTableProcess.addEntity(processInst);
                createTableProcess.addEntity(pathInst);

                registerInstances(createTableProcess);
            } else {
                LOG.info("Process {} is already registered", processQualifiedName);
            }
        }

        return 1;
    } catch (Exception e) {
        LOG.error("Import failed for hive_table {}", tableName, e);

        if (failOnError) {
            throw e;
        }

        return 0;
    }
}
 

private AtlasEntity buildTable(HBaseOperationContext hbaseOperationContext, AtlasEntity nameSpace) {
    AtlasEntity table         = new AtlasEntity(HBaseDataTypes.HBASE_TABLE.getName());
    String      tableName     = getTableName(hbaseOperationContext);
    String      nameSpaceName = (String) nameSpace.getAttribute(ATTR_NAME);
    String      tableQName    = getTableQualifiedName(getMetadataNamespace(), nameSpaceName, tableName);
    OPERATION   operation     = hbaseOperationContext.getOperation();
    Date        now           = new Date(System.currentTimeMillis());

    table.setAttribute(REFERENCEABLE_ATTRIBUTE_NAME, tableQName);
    table.setAttribute(ATTR_NAME, tableName);
    table.setAttribute(ATTR_URI, tableName);
    table.setAttribute(ATTR_OWNER, hbaseOperationContext.getOwner());
    table.setAttribute(ATTR_DESCRIPTION, tableName);
    table.setAttribute(ATTR_PARAMETERS, hbaseOperationContext.getHbaseConf());
    table.setRelationshipAttribute(ATTR_NAMESPACE, AtlasTypeUtil.getAtlasRelatedObjectId(nameSpace, RELATIONSHIP_HBASE_TABLE_NAMESPACE));

    TableDescriptor tableDescriptor = hbaseOperationContext.gethTableDescriptor();
    if (tableDescriptor != null) {
        table.setAttribute(ATTR_TABLE_MAX_FILESIZE, tableDescriptor.getMaxFileSize());
        table.setAttribute(ATTR_TABLE_REPLICATION_PER_REGION, tableDescriptor.getRegionReplication());
        table.setAttribute(ATTR_TABLE_ISREADONLY, tableDescriptor.isReadOnly());
        table.setAttribute(ATTR_TABLE_ISNORMALIZATION_ENABLED, tableDescriptor.isNormalizationEnabled());
        table.setAttribute(ATTR_TABLE_ISCOMPACTION_ENABLED, tableDescriptor.isCompactionEnabled());
        table.setAttribute(ATTR_TABLE_DURABLILITY, (tableDescriptor.getDurability() != null ? tableDescriptor.getDurability().name() : null));
        table.setAttribute(ATTR_TABLE_NORMALIZATION_ENABLED, tableDescriptor.isNormalizationEnabled());
    }

    switch (operation) {
        case CREATE_TABLE:
            table.setAttribute(ATTR_CREATE_TIME, now);
            table.setAttribute(ATTR_MODIFIED_TIME, now);
            break;
        case CREATE_COLUMN_FAMILY:
            table.setAttribute(ATTR_MODIFIED_TIME, now);
            break;
        case ALTER_TABLE:
        case ALTER_COLUMN_FAMILY:
            table.setAttribute(ATTR_MODIFIED_TIME, now);
            break;
        default:
            break;
    }

    return table;
}
 

private AtlasEntity buildColumnFamily(HBaseOperationContext hbaseOperationContext, ColumnFamilyDescriptor columnFamilyDescriptor, AtlasEntity nameSpace, AtlasEntity table) {
    AtlasEntity columnFamily      = new AtlasEntity(HBaseDataTypes.HBASE_COLUMN_FAMILY.getName());
    String      columnFamilyName  = columnFamilyDescriptor.getNameAsString();
    String      tableName         = (String) table.getAttribute(ATTR_NAME);
    String      nameSpaceName     = (String) nameSpace.getAttribute(ATTR_NAME);
    String      columnFamilyQName = getColumnFamilyQualifiedName(getMetadataNamespace(), nameSpaceName, tableName, columnFamilyName);
    Date        now               = new Date(System.currentTimeMillis());

    columnFamily.setAttribute(ATTR_NAME, columnFamilyName);
    columnFamily.setAttribute(ATTR_DESCRIPTION, columnFamilyName);
    columnFamily.setAttribute(REFERENCEABLE_ATTRIBUTE_NAME, columnFamilyQName);
    columnFamily.setAttribute(ATTR_OWNER, hbaseOperationContext.getOwner());
    columnFamily.setRelationshipAttribute(ATTR_TABLE, AtlasTypeUtil.getAtlasRelatedObjectId(table, RELATIONSHIP_HBASE_TABLE_COLUMN_FAMILIES));

    if (columnFamilyDescriptor!= null) {
        columnFamily.setAttribute(ATTR_CF_BLOCK_CACHE_ENABLED, columnFamilyDescriptor.isBlockCacheEnabled());
        columnFamily.setAttribute(ATTR_CF_BLOOMFILTER_TYPE, (columnFamilyDescriptor.getBloomFilterType() != null ? columnFamilyDescriptor.getBloomFilterType().name():null));
        columnFamily.setAttribute(ATTR_CF_CACHED_BLOOM_ON_WRITE, columnFamilyDescriptor.isCacheBloomsOnWrite());
        columnFamily.setAttribute(ATTR_CF_CACHED_DATA_ON_WRITE, columnFamilyDescriptor.isCacheDataOnWrite());
        columnFamily.setAttribute(ATTR_CF_CACHED_INDEXES_ON_WRITE, columnFamilyDescriptor.isCacheIndexesOnWrite());
        columnFamily.setAttribute(ATTR_CF_COMPACTION_COMPRESSION_TYPE, (columnFamilyDescriptor.getCompactionCompressionType() != null ? columnFamilyDescriptor.getCompactionCompressionType().name():null));
        columnFamily.setAttribute(ATTR_CF_COMPRESSION_TYPE, (columnFamilyDescriptor.getCompressionType() != null ? columnFamilyDescriptor.getCompressionType().name():null));
        columnFamily.setAttribute(ATTR_CF_DATA_BLOCK_ENCODING, (columnFamilyDescriptor.getDataBlockEncoding() != null ? columnFamilyDescriptor.getDataBlockEncoding().name():null));
        columnFamily.setAttribute(ATTR_CF_ENCRYPTION_TYPE, columnFamilyDescriptor.getEncryptionType());
        columnFamily.setAttribute(ATTR_CF_EVICT_BLOCK_ONCLOSE, columnFamilyDescriptor.isEvictBlocksOnClose());
        columnFamily.setAttribute(ATTR_CF_INMEMORY_COMPACTION_POLICY, (columnFamilyDescriptor.getInMemoryCompaction() != null ? columnFamilyDescriptor.getInMemoryCompaction().name():null));
        columnFamily.setAttribute(ATTR_CF_KEEP_DELETE_CELLS, ( columnFamilyDescriptor.getKeepDeletedCells() != null ? columnFamilyDescriptor.getKeepDeletedCells().name():null));
        columnFamily.setAttribute(ATTR_CF_MAX_VERSIONS, columnFamilyDescriptor.getMaxVersions());
        columnFamily.setAttribute(ATTR_CF_MIN_VERSIONS, columnFamilyDescriptor.getMinVersions());
        columnFamily.setAttribute(ATTR_CF_NEW_VERSION_BEHAVIOR, columnFamilyDescriptor.isNewVersionBehavior());
        columnFamily.setAttribute(ATTR_CF_MOB_ENABLED, columnFamilyDescriptor.isMobEnabled());
        columnFamily.setAttribute(ATTR_CF_MOB_COMPATCTPARTITION_POLICY, ( columnFamilyDescriptor.getMobCompactPartitionPolicy() != null ? columnFamilyDescriptor.getMobCompactPartitionPolicy().name():null));
        columnFamily.setAttribute(ATTR_CF_PREFETCH_BLOCK_ONOPEN, columnFamilyDescriptor.isPrefetchBlocksOnOpen());
        columnFamily.setAttribute(ATTR_CF_STORAGE_POLICY, columnFamilyDescriptor.getStoragePolicy());
        columnFamily.setAttribute(ATTR_CF_TTL, columnFamilyDescriptor.getTimeToLive());
    }

    switch (hbaseOperationContext.getOperation()) {
        case CREATE_COLUMN_FAMILY:
        case CREATE_TABLE:
            columnFamily.setAttribute(ATTR_CREATE_TIME, now);
            columnFamily.setAttribute(ATTR_MODIFIED_TIME, now);
            break;

        case ALTER_COLUMN_FAMILY:
            columnFamily.setAttribute(ATTR_MODIFIED_TIME, now);
            break;

        default:
            break;
    }

    return columnFamily;
}
 

@Test
public void testRelationshipAttributeUpdate_NonComposite_OneToOne() throws Exception {
    AtlasEntity a1 = new AtlasEntity("A");
    a1.setAttribute(NAME, "a1_name");

    AtlasEntity a2 = new AtlasEntity("A");
    a2.setAttribute(NAME, "a2_name");

    AtlasEntity b = new AtlasEntity("B");
    b.setAttribute(NAME, "b_name");

    AtlasEntitiesWithExtInfo entitiesWithExtInfo = new AtlasEntitiesWithExtInfo();
    entitiesWithExtInfo.addEntity(a1);
    entitiesWithExtInfo.addEntity(a2);
    entitiesWithExtInfo.addEntity(b);

    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesWithExtInfo) , false);

    AtlasEntity partialUpdateB = new AtlasEntity("B");
    partialUpdateB.setRelationshipAttribute("a", getAtlasObjectId(a1));

    init();
    AtlasEntityType bType = typeRegistry.getEntityTypeByName("B");

    response = entityStore.updateByUniqueAttributes(bType, Collections.singletonMap(NAME, b.getAttribute(NAME)), new AtlasEntityWithExtInfo(partialUpdateB));
    List<AtlasEntityHeader> partialUpdatedEntitiesHeader = response.getPartialUpdatedEntities();
    // Verify 2 entities were updated:
    // * set b.a reference to a1
    // * set inverse a1.b reference to b
    assertEquals(partialUpdatedEntitiesHeader.size(), 2);
    AtlasEntitiesWithExtInfo partialUpdatedEntities = entityStore.getByIds(ImmutableList.of(a1.getGuid(), b.getGuid()));

    AtlasEntity a1Entity = partialUpdatedEntities.getEntity(a1.getGuid());
    verifyRelationshipAttributeValue(a1Entity, "b", b.getGuid());

    AtlasEntity bEntity = partialUpdatedEntities.getEntity(b.getGuid());
    verifyRelationshipAttributeValue(bEntity, "a", a1.getGuid());

    init();

    // Update b.a to reference a2.
    partialUpdateB.setRelationshipAttribute("a", getAtlasObjectId(a2));
    response = entityStore.updateByUniqueAttributes(bType, Collections.<String, Object>singletonMap(NAME, b.getAttribute(NAME)), new AtlasEntityWithExtInfo(partialUpdateB));
    partialUpdatedEntitiesHeader = response.getPartialUpdatedEntities();
    // Verify 3 entities were updated:
    // * set b.a reference to a2
    // * set a2.b reference to b
    // * disconnect a1.b reference
    assertEquals(partialUpdatedEntitiesHeader.size(), 3);
    partialUpdatedEntities = entityStore.getByIds(ImmutableList.of(a1.getGuid(), a2.getGuid(), b.getGuid()));

    bEntity = partialUpdatedEntities.getEntity(b.getGuid());
    verifyRelationshipAttributeValue(bEntity, "a", a2.getGuid());

    AtlasEntity a2Entity = partialUpdatedEntities.getEntity(a2.getGuid());
    verifyRelationshipAttributeValue(a2Entity, "b", b.getGuid());

    a1Entity = partialUpdatedEntities.getEntity(a1.getGuid());
    verifyRelationshipAttributeUpdate_NonComposite_OneToOne(a1Entity, bEntity);
}