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

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);
                        }
                    }
                }
            }
        }
    }
}
 

private void normalizeRelationshipAttributeValues(AtlasEntity entity, boolean isUpdate) {
    if (entity != null) {
        for (String attributeName : relationshipAttributes.keySet()) {
            if (entity.hasRelationshipAttribute(attributeName)) {
                Object         attributeValue   = entity.getRelationshipAttribute(attributeName);
                String         relationshipType = AtlasEntityUtil.getRelationshipType(attributeValue);
                AtlasAttribute attribute        = getRelationshipAttribute(attributeName, relationshipType);

                if (attribute != null) {
                    AtlasType attrType = attribute.getAttributeType();

                    if (isValidRelationshipType(attrType)) {
                        if (isUpdate) {
                            attributeValue = attrType.getNormalizedValueForUpdate(attributeValue);
                        } else {
                            attributeValue = attrType.getNormalizedValue(attributeValue);
                        }

                        entity.setRelationshipAttribute(attributeName, attributeValue);
                    }
                }
            }
        }
    }
}
 

void visitEntity(AtlasEntityType entityType, AtlasEntity entity) throws AtlasBaseException {
    visitStruct(entityType, entity);

    for (AtlasAttribute attribute : entityType.getRelationshipAttributes().values()) {
        AtlasType attrType = attribute.getAttributeType();
        Object    attrVal  = entity.getRelationshipAttribute(attribute.getName());

        visitAttribute(attrType, attrVal);
    }
}
 

@Test(dataProvider = "dup_col_data")
public void testImportDuplicateColumnsWithDifferentStatus(InputStream inputStream) throws IOException, AtlasBaseException {
    loadBaseModel();
    loadFsModel();
    loadHiveModel();

    runImportWithNoParameters(importService, inputStream);

    AtlasEntity.AtlasEntityWithExtInfo atlasEntityWithExtInfo = entityStore.getById("e18e15de-1810-4724-881a-5cb6b2160077");
    assertNotNull(atlasEntityWithExtInfo);

    AtlasEntity atlasEntity = atlasEntityWithExtInfo.getEntity();
    assertNotNull(atlasEntity);

    List<AtlasRelatedObjectId> columns = (List<AtlasRelatedObjectId>) atlasEntity.getRelationshipAttribute("columns");
    assertEquals( columns.size(), 4);

    for(AtlasRelatedObjectId id : columns){
        if(id.getGuid().equals("a3de3e3b-4bcd-4e57-a988-1101a2360200")){
            assertEquals(id.getEntityStatus(), AtlasEntity.Status.DELETED);
            assertEquals(id.getRelationshipStatus(), AtlasRelationship.Status.DELETED);
        }
        if(id.getGuid().equals("f7fa3768-f3de-48a8-92a5-38ec4070152c")) {
            assertEquals(id.getEntityStatus(), AtlasEntity.Status.ACTIVE);
            assertEquals(id.getRelationshipStatus(), AtlasRelationship.Status.ACTIVE);
        }
    }
}
 

public void testReactivation(String tableEntityGuid, int columnCount) throws AtlasBaseException, IOException {
    importSeedData();

    AtlasEntity.AtlasEntityWithExtInfo entity = entityStore.getById(tableEntityGuid);
    EntityMutationResponse response = createColumn(entity.getEntity());
    String columnGuid = response.getCreatedEntities().get(0).getGuid();
    assertNotNull(columnGuid);
    columnCount++;

    entityStore.deleteById(tableEntityGuid);
    entity = entityStore.getById(tableEntityGuid);
    assertEquals(entity.getEntity().getStatus(), AtlasEntity.Status.DELETED);

    importSeedData();

    AtlasEntity atlasEntity = entityStore.getById(tableEntityGuid).getEntity();

    assertEquals(atlasEntity.getStatus(), AtlasEntity.Status.ACTIVE);
    List<AtlasRelatedObjectId> columns = (List<AtlasRelatedObjectId>) atlasEntity.getRelationshipAttribute("columns");
    assertEquals(columns.size(), columnCount);

    int activeColumnCount = 0;
    int deletedColumnCount = 0;
    for (AtlasRelatedObjectId column : columns) {
        if (column.getGuid().equals(columnGuid)){
            assertEquals(column.getEntityStatus(), AtlasEntity.Status.DELETED);
            assertEquals(column.getRelationshipStatus(), AtlasRelationship.Status.DELETED);
            deletedColumnCount++;
        }else{
            assertEquals(column.getEntityStatus(), AtlasEntity.Status.ACTIVE);
            assertEquals(column.getRelationshipStatus(), AtlasRelationship.Status.ACTIVE);
            activeColumnCount++;
        }
    }
    assertEquals(activeColumnCount, --columnCount);
    assertEquals(deletedColumnCount, 1);
}
 

protected static void verifyRelationshipAttributeList(AtlasEntity entity, String relationshipAttrName, List<AtlasObjectId> expectedValues) {
    Object refValue = entity.getRelationshipAttribute(relationshipAttrName);
    assertTrue(refValue instanceof List);

    List<AtlasObjectId> refList = toAtlasObjectIds(refValue);
    assertEquals(refList.size(), expectedValues.size());

    if (expectedValues.size() > 0) {
        assertTrue(refList.containsAll(expectedValues));
    }
}
 

protected static void verifyRelationshipAttributeValue(AtlasEntity entity, String relationshipAttrName, String expectedGuid) {
    Object refValue = entity.getRelationshipAttribute(relationshipAttrName);
    if (expectedGuid == null) {
        assertNull(refValue);
    }
    else {
        assertTrue(refValue instanceof AtlasObjectId);
        AtlasObjectId referencedObjectId = (AtlasObjectId) refValue;
        assertEquals(referencedObjectId.getGuid(), expectedGuid);
    }
}
 

protected static void verifyRelationshipAttributeList(AtlasEntity entity, String relationshipAttrName, List<AtlasObjectId> expectedValues) {
    Object refValue = entity.getRelationshipAttribute(relationshipAttrName);
    assertTrue(refValue instanceof List);

    List<AtlasObjectId> refList = (List<AtlasObjectId>) refValue;
    assertEquals(refList.size(), expectedValues.size());

    if (expectedValues.size() > 0) {
        assertTrue(refList.containsAll(expectedValues));
    }
}
 

/**
 * This tests is for create view query with extra comment and spaces added in between:
 * 1) ImpalaLineageTool can parse one lineage file that contains " create   view" command lineage
 * 2) Lineage is sent to Atlas
 * 3) Atlas can get this lineage from Atlas
 */
@Test
public void testCreateViewWithCommentSpacesFromFile() {
    // this file contains a single lineage record for "create view".
    // It has table vertex with createTime
    String IMPALA = dir + "impalaCreateViewWithCommentSpaces.json";
    String IMPALA_WAL = dir + "WALimpala.wal";

    List<ImpalaQuery> lineageList = new ArrayList<>();
    ImpalaLineageHook impalaLineageHook = new ImpalaLineageHook();

    try {
        // create database and tables to simulate Impala behavior that Impala updates metadata
        // to HMS and HMSHook sends the metadata to Atlas, which has to happen before
        // Atlas can handle lineage notification
        String dbName = "db_8";
        createDatabase(dbName);

        String sourceTableName = "table_1";
        createTable(dbName, sourceTableName,"(id string, count int)", false);

        String targetTableName = "view_1";
        createTable(dbName, targetTableName,"(count int, id string)", false);

        // process lineage record, and send corresponding notification to Atlas
        String[] args = new String[]{"-d", "./", "-p", "impala"};
        ImpalaLineageTool toolInstance = new ImpalaLineageTool(args);
        toolInstance.importHImpalaEntities(impalaLineageHook, IMPALA, IMPALA_WAL);

        // verify the process is saved in Atlas
        // the value is from info in IMPALA_3
        String createTime = new Long((long)(1554750072)*1000).toString();
        String processQFName =
                "db_8.view_1" + AtlasImpalaHookContext.QNAME_SEP_METADATA_NAMESPACE +
                        CLUSTER_NAME + AtlasImpalaHookContext.QNAME_SEP_PROCESS + createTime;

        processQFName = processQFName.toLowerCase();

        String      queryString             = " create   /* comment1 */ view db_8.view_1 as   select /* comment2 */ count, id from db_8.table_1";
        AtlasEntity processEntity1          = validateProcess(processQFName, queryString);
        AtlasEntity processExecutionEntity1 = validateProcessExecution(processEntity1, queryString);
        AtlasObjectId process1              = toAtlasObjectId(processExecutionEntity1.getRelationshipAttribute(
                BaseImpalaEvent.ATTRIBUTE_PROCESS));
        Assert.assertEquals(process1.getGuid(), processEntity1.getGuid());
        Assert.assertEquals(numberOfProcessExecutions(processEntity1), 1);

        String      guid       = assertTableIsRegistered(dbName, targetTableName);
        AtlasEntity entity     = atlasClientV2.getEntityByGuid(guid).getEntity();
        List        ddlQueries = (List) entity.getRelationshipAttribute(ATTRIBUTE_DDL_QUERIES);

        assertNotNull(ddlQueries);
        assertEquals(ddlQueries.size(), 1);
    } catch (Exception e) {
        System.out.print("Appending file error");
    }
}
 

/**
 * This tests
 * 1) ImpalaLineageTool can parse one lineage file that contains "create table as select" command lineage,
 *    there is table vertex with createTime. The target vertex's vertexId does not contain db name and table name
 * 2) Lineage is sent to Atlas
 * 3) Atlas can get this lineage from Atlas
 */
@Test
public void testCreateTableAsSelectVertexIdNoTableNameFromFile() throws Exception {
    String IMPALA = dir + "impalaCreateTableAsSelectVertexIdNoTableName.json";
    String IMPALA_WAL = dir + "WALimpala.wal";

    ImpalaLineageHook impalaLineageHook = new ImpalaLineageHook();

    // create database and tables to simulate Impala behavior that Impala updates metadata
    // to HMS and HMSHook sends the metadata to Atlas, which has to happen before
    // Atlas can handle lineage notification
    String dbName = "sales_db";
    createDatabase(dbName);

    String sourceTableName = "sales_asia";
    createTable(dbName, sourceTableName,"(id string, name string)", false);

    String targetTableName = "sales_china";
    createTable(dbName, targetTableName,"(id string, name string)", false);

    // process lineage record, and send corresponding notification to Atlas
    String[] args = new String[]{"-d", "./", "-p", "impala"};
    ImpalaLineageTool toolInstance = new ImpalaLineageTool(args);
    toolInstance.importHImpalaEntities(impalaLineageHook, IMPALA, IMPALA_WAL);

    // verify the process is saved in Atlas
    // the value is from info in IMPALA_4.
    String createTime = new Long((long)1560885039*1000).toString();
    String processQFName =
        dbName + "." + targetTableName + AtlasImpalaHookContext.QNAME_SEP_METADATA_NAMESPACE +
            CLUSTER_NAME + AtlasImpalaHookContext.QNAME_SEP_PROCESS + createTime;

    processQFName = processQFName.toLowerCase();

    String queryString = "create table " + targetTableName + " as select * from " + sourceTableName;
    AtlasEntity processEntity1 = validateProcess(processQFName, queryString);
    AtlasEntity processExecutionEntity1 = validateProcessExecution(processEntity1, queryString);
    AtlasObjectId process1 = toAtlasObjectId(processExecutionEntity1.getRelationshipAttribute(
        BaseImpalaEvent.ATTRIBUTE_PROCESS));
    Assert.assertEquals(process1.getGuid(), processEntity1.getGuid());
    Assert.assertEquals(numberOfProcessExecutions(processEntity1), 1);

    String      guid       = assertTableIsRegistered(dbName, targetTableName);
    AtlasEntity entity     = atlasClientV2.getEntityByGuid(guid).getEntity();
    List        ddlQueries = (List) entity.getRelationshipAttribute(ATTRIBUTE_DDL_QUERIES);

    assertNotNull(ddlQueries);
    assertEquals(ddlQueries.size(), 1);
}
 

/**
 * This tests
 * 1) ImpalaLineageTool can parse one lineage file that contains multiple "insert into" command lineages,
 *    there is table vertex with createTime.
 * 2) Lineage is sent to Atlas
 * 3) Atlas can get these lineages from Atlas
 */
@Test
public void testMultipleInsertIntoAsSelectFromFile() throws Exception {
    String IMPALA = dir + "impalaMultipleInsertIntoAsSelect1.json";
    String IMPALA_WAL = dir + "WALimpala.wal";

    ImpalaLineageHook impalaLineageHook = new ImpalaLineageHook();

    // create database and tables to simulate Impala behavior that Impala updates metadata
    // to HMS and HMSHook sends the metadata to Atlas, which has to happen before
    // Atlas can handle lineage notification
    String dbName = "db_6";
    createDatabase(dbName);

    String sourceTableName = "table_1";
    createTable(dbName, sourceTableName,"(id string, count int)", false);

    String targetTableName = "table_2";
    createTable(dbName, targetTableName,"(count int, id string, int_col int)", false);

    // process lineage record, and send corresponding notification to Atlas
    String[] args = new String[]{"-d", "./", "-p", "impala"};
    ImpalaLineageTool toolInstance = new ImpalaLineageTool(args);
    toolInstance.importHImpalaEntities(impalaLineageHook, IMPALA, IMPALA_WAL);

    // re-run the same lineage record, should have the same process entity and another process execution entity
    Thread.sleep(5000);
    IMPALA = dir + "impalaMultipleInsertIntoAsSelect2.json";
    toolInstance.importHImpalaEntities(impalaLineageHook, IMPALA, IMPALA_WAL);
    Thread.sleep(5000);

    // verify the process is saved in Atlas
    // the value is from info in IMPALA_4.
    String createTime1 = new Long(TABLE_CREATE_TIME_SOURCE*1000).toString();
    String createTime2 = new Long(TABLE_CREATE_TIME*1000).toString();
    String sourceQFName = dbName + "." + sourceTableName + AtlasImpalaHookContext.QNAME_SEP_METADATA_NAMESPACE +
        CLUSTER_NAME + AtlasImpalaHookContext.QNAME_SEP_PROCESS + createTime1;
    String targetQFName = dbName + "." + targetTableName + AtlasImpalaHookContext.QNAME_SEP_METADATA_NAMESPACE +
        CLUSTER_NAME + AtlasImpalaHookContext.QNAME_SEP_PROCESS + createTime2;
    String processQFName = "QUERY:" + sourceQFName.toLowerCase() + "->:INSERT:" + targetQFName.toLowerCase();

    String queryString = "insert into table " + dbName + "." + targetTableName + " (count, id) select count, id from " + dbName + "." + sourceTableName;
    queryString = queryString.toLowerCase().trim();
    String queryString2 = queryString;

    Thread.sleep(5000);
    AtlasEntity processEntity1 = validateProcess(processQFName, queryString);

    List<AtlasObjectId> processExecutions = toAtlasObjectIdList(processEntity1.getRelationshipAttribute(
        BaseImpalaEvent.ATTRIBUTE_PROCESS_EXECUTIONS));
    Assert.assertEquals(processExecutions.size(), 2);
    for (AtlasObjectId processExecutionId : processExecutions) {
        AtlasEntity.AtlasEntityWithExtInfo atlasEntityWithExtInfo = atlasClientV2.
            getEntityByGuid(processExecutionId.getGuid());

        AtlasEntity processExecutionEntity = atlasEntityWithExtInfo.getEntity();
        String entityQueryText = String.valueOf(processExecutionEntity.getAttribute(ATTRIBUTE_QUERY_TEXT)).toLowerCase().trim();
        if (!(queryString.equalsIgnoreCase(entityQueryText) || queryString2.equalsIgnoreCase(entityQueryText))) {
            String errorMessage = String.format("process query text '%s' does not match expected value of '%s' or '%s'", entityQueryText, queryString, queryString2);
            Assert.assertTrue(false, errorMessage);
        }
    }

    String      guid       = assertTableIsRegistered(dbName, targetTableName);
    AtlasEntity entity     = atlasClientV2.getEntityByGuid(guid).getEntity();
    List        ddlQueries = (List) entity.getRelationshipAttribute(ATTRIBUTE_DDL_QUERIES);

    assertNotNull(ddlQueries);
    assertEquals(ddlQueries.size(), 0);
}
 

/**
 * This tests
 * 1) ImpalaLineageTool can parse one lineage file that contains "insert into" command lineage,
 *    there is table vertex with createTime.
 * 2) Lineage is sent to Atlas
 * 3) Atlas can get this lineage from Atlas
 */
@Test
public void testInsertIntoAsSelectFromFile() throws Exception {
    String IMPALA = dir + "impalaInsertIntoAsSelect.json";
    String IMPALA_WAL = dir + "WALimpala.wal";

    ImpalaLineageHook impalaLineageHook = new ImpalaLineageHook();

    // create database and tables to simulate Impala behavior that Impala updates metadata
    // to HMS and HMSHook sends the metadata to Atlas, which has to happen before
    // Atlas can handle lineage notification
    String dbName = "db_5";
    createDatabase(dbName);

    String sourceTableName = "table_1";
    createTable(dbName, sourceTableName,"(id string, count int)", false);

    String targetTableName = "table_2";
    createTable(dbName, targetTableName,"(count int, id string, int_col int)", false);

    // process lineage record, and send corresponding notification to Atlas
    String[] args = new String[]{"-d", "./", "-p", "impala"};
    ImpalaLineageTool toolInstance = new ImpalaLineageTool(args);
    toolInstance.importHImpalaEntities(impalaLineageHook, IMPALA, IMPALA_WAL);

    // verify the process is saved in Atlas
    // the value is from info in IMPALA_4.
    String createTime1 = new Long(TABLE_CREATE_TIME_SOURCE*1000).toString();
    String createTime2 = new Long(TABLE_CREATE_TIME*1000).toString();
    String sourceQFName = dbName + "." + sourceTableName + AtlasImpalaHookContext.QNAME_SEP_METADATA_NAMESPACE +
        CLUSTER_NAME + AtlasImpalaHookContext.QNAME_SEP_PROCESS + createTime1;
    String targetQFName = dbName + "." + targetTableName + AtlasImpalaHookContext.QNAME_SEP_METADATA_NAMESPACE +
        CLUSTER_NAME + AtlasImpalaHookContext.QNAME_SEP_PROCESS + createTime2;
    String processQFName = "QUERY:" + sourceQFName.toLowerCase() + "->:INSERT:" + targetQFName.toLowerCase();

    String queryString = "insert into table " + dbName + "." + targetTableName + " (count, id) select count, id from " + dbName + "." + sourceTableName;
    AtlasEntity processEntity1 = validateProcess(processQFName, queryString);
    AtlasEntity processExecutionEntity1 = validateProcessExecution(processEntity1, queryString);
    AtlasObjectId process1 = toAtlasObjectId(processExecutionEntity1.getRelationshipAttribute(
        BaseImpalaEvent.ATTRIBUTE_PROCESS));
    Assert.assertEquals(process1.getGuid(), processEntity1.getGuid());
    Assert.assertEquals(numberOfProcessExecutions(processEntity1), 1);

    String      guid       = assertTableIsRegistered(dbName, targetTableName);
    AtlasEntity entity     = atlasClientV2.getEntityByGuid(guid).getEntity();
    List        ddlQueries = (List) entity.getRelationshipAttribute(ATTRIBUTE_DDL_QUERIES);

    assertNotNull(ddlQueries);
    assertEquals(ddlQueries.size(), 0);
}
 

/**
 * This tests is for extra comment and spaces present in alter view as select query
 * 1) ImpalaLineageTool can parse one lineage file that contains "alter view as select" command lineage,
 *    there is table vertex with createTime.
 * 2) Lineage is sent to Atlas
 * 3) Atlas can get this lineage from Atlas
 */
@Test
public void testAlterViewAsSelectWithCommentSpacesFromFile() throws Exception {
    String IMPALA = dir + "impalaAlterViewAsSelectWithCommentSpaces.json";
    String IMPALA_WAL = dir + "WALimpala.wal";

    ImpalaLineageHook impalaLineageHook = new ImpalaLineageHook();

    // create database and tables to simulate Impala behavior that Impala updates metadata
    // to HMS and HMSHook sends the metadata to Atlas, which has to happen before
    // Atlas can handle lineage notification
    String dbName = "db_10";
    createDatabase(dbName);

    String sourceTableName = "table_1";
    createTable(dbName, sourceTableName,"(id string, count int)", false);

    String targetTableName = "view_1";
    createTable(dbName, targetTableName,"(count int, id string)", false);

    // process lineage record, and send corresponding notification to Atlas
    String[] args = new String[]{"-d", "./", "-p", "impala"};
    ImpalaLineageTool toolInstance = new ImpalaLineageTool(args);
    toolInstance.importHImpalaEntities(impalaLineageHook, IMPALA, IMPALA_WAL);

    // verify the process is saved in Atlas
    // the value is from info in IMPALA_4.
    String createTime = new Long(TABLE_CREATE_TIME*1000).toString();
    String processQFName =
            dbName + "." + targetTableName + AtlasImpalaHookContext.QNAME_SEP_METADATA_NAMESPACE +
                    CLUSTER_NAME + AtlasImpalaHookContext.QNAME_SEP_PROCESS + createTime;

    processQFName = processQFName.toLowerCase();

    String queryString = "alter   /* comment1 */ view " + dbName + "." + targetTableName
            + " as   select /* comment1 */ count, id from " + dbName + "." + sourceTableName;
    AtlasEntity processEntity1 = validateProcess(processQFName, queryString);
    AtlasEntity processExecutionEntity1 = validateProcessExecution(processEntity1, queryString);
    AtlasObjectId process1 = toAtlasObjectId(processExecutionEntity1.getRelationshipAttribute(
            BaseImpalaEvent.ATTRIBUTE_PROCESS));
    Assert.assertEquals(process1.getGuid(), processEntity1.getGuid());
    Assert.assertEquals(numberOfProcessExecutions(processEntity1), 1);

    String      guid       = assertTableIsRegistered(dbName, targetTableName);
    AtlasEntity entity     = atlasClientV2.getEntityByGuid(guid).getEntity();
    List        ddlQueries = (List) entity.getRelationshipAttribute(ATTRIBUTE_DDL_QUERIES);

    assertNotNull(ddlQueries);
    assertEquals(ddlQueries.size(), 1);
}
 

/**
 * This tests
 * 1) ImpalaLineageTool can parse one lineage file that contains "alter view as select" command lineage,
 *    there is table vertex with createTime.
 * 2) Lineage is sent to Atlas
 * 3) Atlas can get this lineage from Atlas
 */
@Test
public void testAlterViewAsSelectFromFile() throws Exception {
    String IMPALA = dir + "impalaAlterViewAsSelect.json";
    String IMPALA_WAL = dir + "WALimpala.wal";

    ImpalaLineageHook impalaLineageHook = new ImpalaLineageHook();

    // create database and tables to simulate Impala behavior that Impala updates metadata
    // to HMS and HMSHook sends the metadata to Atlas, which has to happen before
    // Atlas can handle lineage notification
    String dbName = "db_4";
    createDatabase(dbName);

    String sourceTableName = "table_1";
    createTable(dbName, sourceTableName,"(id string, count int)", false);

    String targetTableName = "view_1";
    createTable(dbName, targetTableName,"(count int, id string)", false);

    // process lineage record, and send corresponding notification to Atlas
    String[] args = new String[]{"-d", "./", "-p", "impala"};
    ImpalaLineageTool toolInstance = new ImpalaLineageTool(args);
    toolInstance.importHImpalaEntities(impalaLineageHook, IMPALA, IMPALA_WAL);

    // verify the process is saved in Atlas
    // the value is from info in IMPALA_4.
    String createTime = new Long(TABLE_CREATE_TIME*1000).toString();
    String processQFName =
        dbName + "." + targetTableName + AtlasImpalaHookContext.QNAME_SEP_METADATA_NAMESPACE +
            CLUSTER_NAME + AtlasImpalaHookContext.QNAME_SEP_PROCESS + createTime;

    processQFName = processQFName.toLowerCase();

    String queryString = "alter view " + dbName + "." + targetTableName + " as select count, id from " + dbName + "." + sourceTableName;
    AtlasEntity processEntity1 = validateProcess(processQFName, queryString);
    AtlasEntity processExecutionEntity1 = validateProcessExecution(processEntity1, queryString);
    AtlasObjectId process1 = toAtlasObjectId(processExecutionEntity1.getRelationshipAttribute(
        BaseImpalaEvent.ATTRIBUTE_PROCESS));
    Assert.assertEquals(process1.getGuid(), processEntity1.getGuid());
    Assert.assertEquals(numberOfProcessExecutions(processEntity1), 1);

    String      guid       = assertTableIsRegistered(dbName, targetTableName);
    AtlasEntity entity     = atlasClientV2.getEntityByGuid(guid).getEntity();
    List        ddlQueries = (List) entity.getRelationshipAttribute(ATTRIBUTE_DDL_QUERIES);

    assertNotNull(ddlQueries);
    assertEquals(ddlQueries.size(), 1);
}
 

/**
 * This tests is based on extra comment and spaces adding to create table as select query
 * 1) ImpalaLineageTool can parse one lineage file that contains "create   table   as   select" command lineage,
 *    there is table vertex with createTime.
 * 2) Lineage is sent to Atlas
 * 3) Atlas can get this lineage from Atlas
 */
@Test
public void testCreateTableAsSelectWithCommentSpacesFromFile() throws Exception {
    String IMPALA = dir + "impalaCreateTableAsSelectWithCommentSpaces.json";
    String IMPALA_WAL = dir + "WALimpala.wal";

    ImpalaLineageHook impalaLineageHook = new ImpalaLineageHook();

    // create database and tables to simulate Impala behavior that Impala updates metadata
    // to HMS and HMSHook sends the metadata to Atlas, which has to happen before
    // Atlas can handle lineage notification
    String dbName = "db_9";
    createDatabase(dbName);

    String sourceTableName = "table_1";
    createTable(dbName, sourceTableName,"(id string, count int)", false);

    String targetTableName = "table_2";
    createTable(dbName, targetTableName,"(count int, id string)", false);

    // process lineage record, and send corresponding notification to Atlas
    String[] args = new String[]{"-d", "./", "-p", "impala"};
    ImpalaLineageTool toolInstance = new ImpalaLineageTool(args);
    toolInstance.importHImpalaEntities(impalaLineageHook, IMPALA, IMPALA_WAL);

    // verify the process is saved in Atlas
    // the value is from info in IMPALA_4.
    String createTime = new Long(TABLE_CREATE_TIME*1000).toString();
    String processQFName =
            dbName + "." + targetTableName + AtlasImpalaHookContext.QNAME_SEP_METADATA_NAMESPACE +
                    CLUSTER_NAME + AtlasImpalaHookContext.QNAME_SEP_PROCESS + createTime;

    processQFName = processQFName.toLowerCase();

    String queryString = "create   /* Test */   table " + dbName + "."
            + targetTableName + "   as /* Test */ select count, id from " + dbName + "." + sourceTableName;
    AtlasEntity processEntity1 = validateProcess(processQFName, queryString);
    AtlasEntity processExecutionEntity1 = validateProcessExecution(processEntity1, queryString);
    AtlasObjectId process1 = toAtlasObjectId(processExecutionEntity1.getRelationshipAttribute(
            BaseImpalaEvent.ATTRIBUTE_PROCESS));
    Assert.assertEquals(process1.getGuid(), processEntity1.getGuid());
    Assert.assertEquals(numberOfProcessExecutions(processEntity1), 1);

    String      guid       = assertTableIsRegistered(dbName, targetTableName);
    AtlasEntity entity     = atlasClientV2.getEntityByGuid(guid).getEntity();
    List        ddlQueries = (List) entity.getRelationshipAttribute(ATTRIBUTE_DDL_QUERIES);

    assertNotNull(ddlQueries);
    assertEquals(ddlQueries.size(), 1);
}
 

/**
 * This tests
 * 1) ImpalaLineageTool can parse one lineage file that contains "create table as select" command lineage,
 *    there is table vertex with createTime.
 * 2) Lineage is sent to Atlas
 * 3) Atlas can get this lineage from Atlas
 */
@Test
public void testCreateTableAsSelectFromFile() throws Exception {
    String IMPALA = dir + "impalaCreateTableAsSelect.json";
    String IMPALA_WAL = dir + "WALimpala.wal";

    ImpalaLineageHook impalaLineageHook = new ImpalaLineageHook();

    // create database and tables to simulate Impala behavior that Impala updates metadata
    // to HMS and HMSHook sends the metadata to Atlas, which has to happen before
    // Atlas can handle lineage notification
    String dbName = "db_3";
    createDatabase(dbName);

    String sourceTableName = "table_1";
    createTable(dbName, sourceTableName,"(id string, count int)", false);

    String targetTableName = "table_2";
    createTable(dbName, targetTableName,"(count int, id string)", false);

    // process lineage record, and send corresponding notification to Atlas
    String[] args = new String[]{"-d", "./", "-p", "impala"};
    ImpalaLineageTool toolInstance = new ImpalaLineageTool(args);
    toolInstance.importHImpalaEntities(impalaLineageHook, IMPALA, IMPALA_WAL);

    // verify the process is saved in Atlas
    // the value is from info in IMPALA_4.
    String createTime = new Long(TABLE_CREATE_TIME*1000).toString();
    String processQFName =
        dbName + "." + targetTableName + AtlasImpalaHookContext.QNAME_SEP_METADATA_NAMESPACE +
            CLUSTER_NAME + AtlasImpalaHookContext.QNAME_SEP_PROCESS + createTime;

    processQFName = processQFName.toLowerCase();

    String queryString = "create table " + dbName + "." + targetTableName + " as select count, id from " + dbName + "." + sourceTableName;
    AtlasEntity processEntity1 = validateProcess(processQFName, queryString);
    AtlasEntity processExecutionEntity1 = validateProcessExecution(processEntity1, queryString);
    AtlasObjectId process1 = toAtlasObjectId(processExecutionEntity1.getRelationshipAttribute(
        BaseImpalaEvent.ATTRIBUTE_PROCESS));
    Assert.assertEquals(process1.getGuid(), processEntity1.getGuid());
    Assert.assertEquals(numberOfProcessExecutions(processEntity1), 1);

    String      guid       = assertTableIsRegistered(dbName, targetTableName);
    AtlasEntity entity     = atlasClientV2.getEntityByGuid(guid).getEntity();
    List        ddlQueries = (List) entity.getRelationshipAttribute(ATTRIBUTE_DDL_QUERIES);

    assertNotNull(ddlQueries);
    assertEquals(ddlQueries.size(), 1);
}
 

/**
 * This tests
 * 1) ImpalaLineageTool can parse one lineage file that contains "create view" command lineage,
 *    but there is no table vertex with createTime.
 * 2) Lineage is sent to Atlas
 * 3) Atlas can get this lineage from Atlas
 */
@Test
public void testCreateViewNoCreateTimeFromFile() {
    // this file contains a single lineage record for "create view".
    // there is no table vertex with createTime, which is lineage record generated by Impala
    // originally. The table create time is hard-coded before Impala fixes this issue.
    String IMPALA     = dir + "impalaCreateViewNoCreateTime.json";
    String IMPALA_WAL = dir + "WALimpala.wal";

    List<ImpalaQuery> lineageList       = new ArrayList<>();
    ImpalaLineageHook impalaLineageHook = new ImpalaLineageHook();

    try {
        // create database and tables to simulate Impala behavior that Impala updates metadata
        // to HMS and HMSHook sends the metadata to Atlas, which has to happen before
        // Atlas can handle lineage notification
        String dbName = "db_2";
        createDatabase(dbName);

        String sourceTableName = "table_1";
        createTable(dbName, sourceTableName,"(id string, count int)", false);

        String targetTableName = "view_1";
        createTable(dbName, targetTableName,"(count int, id string)", false);

        // process lineage record, and send corresponding notification to Atlas
        String[] args = new String[]{"-d", "./", "-p", "impala"};
        ImpalaLineageTool toolInstance = new ImpalaLineageTool(args);
        Long beforeCreateTime = System.currentTimeMillis() / BaseImpalaEvent.MILLIS_CONVERT_FACTOR;
        toolInstance.importHImpalaEntities(impalaLineageHook, IMPALA, IMPALA_WAL);
        Long afterCreateTime = System.currentTimeMillis() / BaseImpalaEvent.MILLIS_CONVERT_FACTOR;

        String processQFNameWithoutTime =
            dbName + "." + targetTableName + AtlasImpalaHookContext.QNAME_SEP_METADATA_NAMESPACE +
                CLUSTER_NAME + AtlasImpalaHookContext.QNAME_SEP_PROCESS;
        processQFNameWithoutTime = processQFNameWithoutTime.toLowerCase();

        List<String> processQFNames = new ArrayList<>();
        String createTime = new Long(beforeCreateTime.longValue()*1000).toString();
        processQFNames.add(processQFNameWithoutTime + createTime);

        if (beforeCreateTime != afterCreateTime) {
            createTime = new Long(afterCreateTime.longValue() * 1000).toString();
            processQFNames.add(processQFNameWithoutTime + createTime);
        }

        // verify the process is saved in Atlas. the value is from info in IMPALA_4.
        // There is no createTime in lineage record, so we don't know the process qualified name
        // And can only verify the process is created for the given query.
        String queryString = "create view " + dbName + "." + targetTableName + " as select count, id from " + dbName + "." + sourceTableName;
        AtlasEntity processEntity1 = validateProcess(processQFNames, queryString);
        AtlasEntity processExecutionEntity1 = validateProcessExecution(processEntity1, queryString);
        AtlasObjectId process1 = toAtlasObjectId(processExecutionEntity1.getRelationshipAttribute(
            BaseImpalaEvent.ATTRIBUTE_PROCESS));
        Assert.assertEquals(process1.getGuid(), processEntity1.getGuid());
        Assert.assertEquals(numberOfProcessExecutions(processEntity1), 1);

        String      guid       = assertTableIsRegistered(dbName, targetTableName);
        AtlasEntity entity     = atlasClientV2.getEntityByGuid(guid).getEntity();
        List        ddlQueries = (List) entity.getRelationshipAttribute(ATTRIBUTE_DDL_QUERIES);

        assertNotNull(ddlQueries);
        assertEquals(ddlQueries.size(), 1);
    } catch (Exception e) {
        System.out.print("Appending file error");
    }
}
 

/**
 * This tests
 * 1) ImpalaLineageTool can parse one lineage file that contains "create view" command lineage
 * 2) Lineage is sent to Atlas
 * 3) Atlas can get this lineage from Atlas
 */
@Test
public void testCreateViewFromFile() {
    // this file contains a single lineage record for "create view".
    // It has table vertex with createTime
    String IMPALA = dir + "impalaCreateView.json";
    String IMPALA_WAL = dir + "WALimpala.wal";

    List<ImpalaQuery> lineageList = new ArrayList<>();
    ImpalaLineageHook impalaLineageHook = new ImpalaLineageHook();

    try {
        // create database and tables to simulate Impala behavior that Impala updates metadata
        // to HMS and HMSHook sends the metadata to Atlas, which has to happen before
        // Atlas can handle lineage notification
        String dbName = "db_1";
        createDatabase(dbName);

        String sourceTableName = "table_1";
        createTable(dbName, sourceTableName,"(id string, count int)", false);

        String targetTableName = "view_1";
        createTable(dbName, targetTableName,"(count int, id string)", false);

        // process lineage record, and send corresponding notification to Atlas
        String[] args = new String[]{"-d", "./", "-p", "impala"};
        ImpalaLineageTool toolInstance = new ImpalaLineageTool(args);
        toolInstance.importHImpalaEntities(impalaLineageHook, IMPALA, IMPALA_WAL);

        // verify the process is saved in Atlas
        // the value is from info in IMPALA_3
        String createTime = new Long((long)(1554750072)*1000).toString();
        String processQFName =
            "db_1.view_1" + AtlasImpalaHookContext.QNAME_SEP_METADATA_NAMESPACE +
                CLUSTER_NAME + AtlasImpalaHookContext.QNAME_SEP_PROCESS + createTime;

        processQFName = processQFName.toLowerCase();

        String      queryString             = "create view db_1.view_1 as select count, id from db_1.table_1";
        AtlasEntity processEntity1          = validateProcess(processQFName, queryString);
        AtlasEntity processExecutionEntity1 = validateProcessExecution(processEntity1, queryString);
        AtlasObjectId process1              = toAtlasObjectId(processExecutionEntity1.getRelationshipAttribute(
            BaseImpalaEvent.ATTRIBUTE_PROCESS));
        Assert.assertEquals(process1.getGuid(), processEntity1.getGuid());
        Assert.assertEquals(numberOfProcessExecutions(processEntity1), 1);

        String      guid       = assertTableIsRegistered(dbName, targetTableName);
        AtlasEntity entity     = atlasClientV2.getEntityByGuid(guid).getEntity();
        List        ddlQueries = (List) entity.getRelationshipAttribute(ATTRIBUTE_DDL_QUERIES);

        assertNotNull(ddlQueries);
        assertEquals(ddlQueries.size(), 1);
    } catch (Exception e) {
        System.out.print("Appending file error");
    }
}
 

@Override
public AtlasGlossaryCategory from(final AtlasEntity entity) {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasGlossaryCategoryDTO.from()", entity);
    }
    Objects.requireNonNull(entity, "entity");

    AtlasGlossaryCategory ret = new AtlasGlossaryCategory();

    ret.setGuid(entity.getGuid());
    ret.setQualifiedName((String) entity.getAttribute("qualifiedName"));
    ret.setName((String) entity.getAttribute("name"));
    ret.setShortDescription((String) entity.getAttribute("shortDescription"));
    ret.setLongDescription((String) entity.getAttribute("longDescription"));
    ret.setAdditionalAttributes((Map) entity.getAttribute("additionalAttributes"));

    Object anchor = entity.getRelationshipAttribute("anchor");
    if (anchor instanceof AtlasRelatedObjectId) {
        LOG.debug("Processing anchor");
        ret.setAnchor(constructGlossaryId((AtlasRelatedObjectId) anchor));
    }

    Object parentCategory = entity.getRelationshipAttribute("parentCategory");
    if (parentCategory instanceof AtlasRelatedObjectId) {
        LOG.debug("Processing parentCategory");
        ret.setParentCategory(constructRelatedCategoryId((AtlasRelatedObjectId) parentCategory));
    }

    Object childrenCategories = entity.getRelationshipAttribute("childrenCategories");
    if (childrenCategories instanceof Collection) {
        LOG.debug("Processing childrenCategories");
        for (Object child : (Collection) childrenCategories) {
            if (child instanceof AtlasRelatedObjectId) {
                if (((AtlasRelatedObjectId) child).getRelationshipStatus() == AtlasRelationship.Status.ACTIVE) {
                    ret.addChild(constructRelatedCategoryId((AtlasRelatedObjectId) child));
                }
            }
        }
    }

    Object terms = entity.getRelationshipAttribute("terms");
    if (terms instanceof Collection) {
        LOG.debug("Processing terms");
        for (Object term : (Collection) terms) {
            if (term instanceof AtlasRelatedObjectId) {
                ret.addTerm(constructRelatedTermId((AtlasRelatedObjectId) term));
            }
        }
    }

    return ret;
}
 

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);
            }
        }
    }
}