Java Code Examples for org.apache.atlas.model.instance.AtlasClassification#setEntityGuid()

@Test(dependsOnMethods = {"addClassification_PropagateFalse"})
public void updateClassification_PropagateFalseToTrue() throws AtlasBaseException {
    AtlasEntity         hdfs_employees = getEntity(HDFS_PATH_EMPLOYEES);
    AtlasClassification tag2           = new AtlasClassification("tag2"); tag2.setEntityGuid(hdfs_employees.getGuid());

    //update tag2 propagate to 'true'
    tag2 = getClassification(hdfs_employees, tag2); tag2.setPropagate(true);

    updateClassifications(hdfs_employees, tag2);

    List<String> propagatedToEntities = Arrays.asList(EMPLOYEES1_PROCESS, EMPLOYEES2_PROCESS, EMPLOYEES1_TABLE,
                                                      EMPLOYEES2_TABLE, EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);

    assertClassificationExistInEntities(propagatedToEntities, tag2);

    deleteClassification(hdfs_employees, tag2);
}
 

@Test(dependsOnMethods = {"removeBlockedPropagatedClassifications"})
public void addClassification_removePropagationsTrue_DeleteCase() throws AtlasBaseException {
    AtlasEntity         orders = getEntity(ORDERS_TABLE);
    AtlasClassification tag2      = new AtlasClassification("tag2");

    tag2.setEntityGuid(orders.getGuid());
    tag2.setPropagate(true);
    tag2.setRemovePropagationsOnEntityDelete(true);

    addClassification(orders, tag2);

    List<String> propagatedEntities = Arrays.asList(EMPLOYEES_PROCESS, US_EMPLOYEES_TABLE);

    assertClassificationExistInEntities(propagatedEntities, tag2);

    AtlasEntity orders_process = getEntity(ORDERS_PROCESS);
    AtlasEntity us_orders      = getEntity(US_ORDERS_TABLE);

    deletePropagatedClassificationExpectFail(orders_process, tag2);
    deletePropagatedClassificationExpectFail(us_orders, tag2);

    deleteEntity(ORDERS_TABLE);
    assertClassificationNotExistInEntity(ORDERS_PROCESS, tag2);
    assertClassificationNotExistInEntity(US_ORDERS_TABLE, tag2);
}
 

@Test(dependsOnMethods = {"addClassification_removePropagationsTrue_DeleteCase"})
public void addClassification_removePropagationsFalse_DeleteCase() throws AtlasBaseException {
    AtlasEntity         employees = getEntity(EMPLOYEES_TABLE);
    AtlasClassification tag1      = new AtlasClassification("tag1");

    tag1.setEntityGuid(employees.getGuid());
    tag1.setPropagate(true);
    tag1.setRemovePropagationsOnEntityDelete(false);

    addClassification(employees, tag1);

    List<String> propagatedEntities = Arrays.asList(EMPLOYEES_PROCESS, US_EMPLOYEES_TABLE);

    assertClassificationExistInEntities(propagatedEntities, tag1);

    AtlasEntity employees_process = getEntity(EMPLOYEES_PROCESS);
    AtlasEntity us_employees      = getEntity(US_EMPLOYEES_TABLE);

    deletePropagatedClassificationExpectFail(employees_process, tag1);
    deletePropagatedClassificationExpectFail(us_employees, tag1);

    deleteEntity(EMPLOYEES_PROCESS);
    assertClassificationExistInEntity(EMPLOYEES_PROCESS, tag1);
    assertClassificationExistInEntity(US_EMPLOYEES_TABLE, tag1);
}
 

private void updateClassificationsForHeader(AtlasEntityHeader header, AtlasEntityHeader currentHeader, boolean keyFound) {
    for (AtlasClassification c : header.getClassifications()) {
        c.setEntityGuid(null);

        if (keyFound) {
            boolean found =
                    currentHeader.getClassifications().stream().anyMatch(ox -> ox.getTypeName().equals(c.getTypeName()));
            if (!found) {
                currentHeader.getClassifications().add(c);
            } else {
                displayCrLf("Ignoring: " + c.toString());
                LOG.warn("Ignoring: {}", AtlasJson.toJson(c));
            }
        }
    }
}
 

@Test(dependsOnMethods = {"updateClassification_PropagateFalseToTrue"})
public void addClassification_PropagateTrue() throws AtlasBaseException {
    AtlasEntity         hdfs_employees = getEntity(HDFS_PATH_EMPLOYEES);
    AtlasClassification tag1           = new AtlasClassification("tag1"); tag1.setPropagate(true); tag1.setEntityGuid(hdfs_employees.getGuid());

    // add classification with propagate flag to 'true'
    addClassification(hdfs_employees, tag1);

    List<String> propagatedToEntities = Arrays.asList(EMPLOYEES1_PROCESS, EMPLOYEES2_PROCESS, EMPLOYEES1_TABLE,
                                                      EMPLOYEES2_TABLE, EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);

    assertClassificationExistInEntities(propagatedToEntities, tag1);
}
 

@Test(dependsOnMethods = {"addClassification_PropagateTrue"})
public void updateClassification_PropagateTrueToFalse() throws AtlasBaseException {
    AtlasEntity         hdfs_employees = getEntity(HDFS_PATH_EMPLOYEES);
    AtlasClassification tag1           = new AtlasClassification("tag1"); tag1.setEntityGuid(hdfs_employees.getGuid());

    List<String> propagatedToEntities = Arrays.asList(EMPLOYEES1_PROCESS, EMPLOYEES2_PROCESS, EMPLOYEES1_TABLE,
                                                      EMPLOYEES2_TABLE, EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);

    // update propagate flag to 'false'
    tag1 = getClassification(hdfs_employees, tag1); tag1.setPropagate(false);

    updateClassifications(hdfs_employees, tag1);

    assertClassificationNotExistInEntities(propagatedToEntities, tag1);
}
 

@Test(dependsOnMethods = {"updateClassification_PropagateTrueToFalse"})
public void deleteClassification_PropagateTrue() throws AtlasBaseException {
    AtlasEntity         hdfs_employees = getEntity(HDFS_PATH_EMPLOYEES);
    AtlasClassification tag1           = new AtlasClassification("tag1"); tag1.setPropagate(true); tag1.setEntityGuid(hdfs_employees.getGuid());

    deleteClassification(hdfs_employees, tag1);

    List<String> propagatedToEntities = Arrays.asList(EMPLOYEES1_PROCESS, EMPLOYEES2_PROCESS, EMPLOYEES1_TABLE,
                                                      EMPLOYEES2_TABLE, EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);

    assertClassificationNotExistInEntities(propagatedToEntities, tag1);
}
 

private List<AtlasClassification> getAllClassifications(List<AtlasClassification> classifications) {
    List<AtlasClassification> ret = new ArrayList<>();

    if (CollectionUtils.isNotEmpty(classifications)) {
        for (AtlasClassification classification : classifications) {
            AtlasClassificationType classificationType = typeRegistry.getClassificationTypeByName(classification.getTypeName());
            Set<String>             superTypeNames     = classificationType != null ? classificationType.getAllSuperTypes() : null;

            ret.add(classification);

            if (CollectionUtils.isNotEmpty(superTypeNames)) {
                for (String superTypeName : superTypeNames) {
                    AtlasClassification superTypeClassification = new AtlasClassification(superTypeName);

                    superTypeClassification.setEntityGuid(classification.getEntityGuid());
                    superTypeClassification.setPropagate(classification.isPropagate());

                    if (MapUtils.isNotEmpty(classification.getAttributes())) {
                        AtlasClassificationType superType = typeRegistry.getClassificationTypeByName(superTypeName);

                        if (superType != null && MapUtils.isNotEmpty(superType.getAllAttributes())) {
                            Map<String, Object> superTypeClassificationAttributes = new HashMap<>();

                            for (Map.Entry<String, Object> attrEntry : classification.getAttributes().entrySet()) {
                                String attrName = attrEntry.getKey();

                                if (superType.getAllAttributes().containsKey(attrName)) {
                                    superTypeClassificationAttributes.put(attrName, attrEntry.getValue());
                                }
                            }

                            superTypeClassification.setAttributes(superTypeClassificationAttributes);
                        }
                    }

                    ret.add(superTypeClassification);
                }
            }
        }
    }

    return ret;
}
 

@Test(dependsOnMethods = {"propagateSameTagFromDifferentEntities"})
public void updatePropagateTagsValue() throws AtlasBaseException {
    AtlasEntity hdfs_employees          = getEntity(HDFS_PATH_EMPLOYEES);
    AtlasEntity employees2_table        = getEntity(EMPLOYEES2_TABLE);
    AtlasEntity employees_union_process = getEntity(EMPLOYEES_UNION_PROCESS);
    AtlasEntity employees_union_table   = getEntity(EMPLOYEES_UNION_TABLE);

    AtlasClassification tag1 = new AtlasClassification("tag1"); tag1.setPropagate(true); tag1.setEntityGuid(hdfs_employees.getGuid());
    AtlasClassification tag2 = new AtlasClassification("tag2"); tag2.setPropagate(true); tag2.setEntityGuid(employees2_table.getGuid());
    AtlasClassification tag3 = new AtlasClassification("tag3"); tag3.setPropagate(true); tag3.setEntityGuid(employees_union_process.getGuid());
    AtlasClassification tag4 = new AtlasClassification("tag4"); tag4.setPropagate(true); tag4.setEntityGuid(employees_union_table.getGuid());

    // add tag1 to hdfs_employees, tag2 to employees2, tag3 to process3, tag4 to employees_union
    addClassification(hdfs_employees, tag1);
    addClassification(employees2_table, tag2);
    addClassification(employees_union_process, tag3);
    addClassification(employees_union_table, tag4);

    //validate if tag1, tag2, tag3 propagated to employees_union table
    assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, tag1);
    assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, tag2);
    assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, tag3);
    assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, tag4);

    // change propagation between employees2 -> process3 from TWO_TO_ONE to NONE
    AtlasRelationship employees2_process_relationship = getRelationship(EMPLOYEES2_TABLE, EMPLOYEES_UNION_PROCESS);
    assertEquals(employees2_process_relationship.getPropagateTags(), TWO_TO_ONE);
    employees2_process_relationship.setPropagateTags(NONE);
    relationshipStore.update(employees2_process_relationship);

    // validate tag1 propagated to employees_union through other path
    assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, tag1);

    // validate tag2 is no more propagated to employees_union
    assertClassificationNotExistInEntity(EMPLOYEES_UNION_TABLE, tag2);

    // change propagation between employees2 -> process3 from NONE to TWO_TO_ONE
    employees2_process_relationship = getRelationship(EMPLOYEES2_TABLE, EMPLOYEES_UNION_PROCESS);
    assertEquals(employees2_process_relationship.getPropagateTags(), NONE);
    employees2_process_relationship.setPropagateTags(TWO_TO_ONE);
    relationshipStore.update(employees2_process_relationship);

    // validate tag2 is propagated to employees_union
    assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, tag2);

    //update propagation to BOTH for edge process3 --> employee_union. This should fail
    AtlasRelationship process3_employee_union_relationship = getRelationship(EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);
    assertEquals(process3_employee_union_relationship.getPropagateTags(), ONE_TO_TWO);
    process3_employee_union_relationship.setPropagateTags(BOTH);

    try {
        relationshipStore.update(process3_employee_union_relationship);
    } catch (AtlasBaseException ex) {
        assertEquals(ex.getAtlasErrorCode(), AtlasErrorCode.INVALID_PROPAGATION_TYPE);
    }

    //cleanup
    deleteClassification(hdfs_employees, tag1);
    deleteClassification(employees2_table, tag2);
    deleteClassification(employees_union_process, tag3);
    deleteClassification(employees_union_table, tag4);
}
 

@Test(dependsOnMethods = {"updatePropagateTagsValue"})
public void addBlockedPropagatedClassifications() throws AtlasBaseException {
    AtlasEntity hdfs_path       = getEntity(HDFS_PATH_EMPLOYEES);
    AtlasEntity employees1      = getEntity(EMPLOYEES1_TABLE);
    AtlasEntity employees2      = getEntity(EMPLOYEES2_TABLE);
    AtlasEntity employees_union = getEntity(EMPLOYEES_UNION_TABLE);

    AtlasClassification PII_tag1 = new AtlasClassification("PII"); PII_tag1.setPropagate(true);
        PII_tag1.setAttribute("type", "from hdfs_path entity");
        PII_tag1.setAttribute("valid", true);

    AtlasClassification PII_tag2 = new AtlasClassification("PII"); PII_tag2.setPropagate(true);
        PII_tag2.setAttribute("type", "from employees1 entity");
        PII_tag2.setAttribute("valid", true);

    AtlasClassification PII_tag3 = new AtlasClassification("PII"); PII_tag3.setPropagate(true);
        PII_tag3.setAttribute("type", "from employees2 entity");
        PII_tag3.setAttribute("valid", true);

    AtlasClassification PII_tag4 = new AtlasClassification("PII"); PII_tag4.setPropagate(true);
        PII_tag4.setAttribute("type", "from employees_union entity");
        PII_tag4.setAttribute("valid", true);

    // add PII to hdfs_path, employees1, employees2 and employee_union
    addClassification(hdfs_path, PII_tag1);
    addClassification(employees1, PII_tag2);
    addClassification(employees2, PII_tag3);

    // check 4 PII tags exists in employee_union table
    assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, PII_tag1.getTypeName(), hdfs_path.getGuid());
    assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, PII_tag2.getTypeName(), employees1.getGuid());
    assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, PII_tag3.getTypeName(), employees2.getGuid());

    AtlasRelationship process3_employee_union_relationship = getRelationship(EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);
    Set<AtlasClassification> propagatedClassifications     = process3_employee_union_relationship.getPropagatedClassifications();
    Set<AtlasClassification> blockedClassifications        = process3_employee_union_relationship.getBlockedPropagatedClassifications();

    assertNotNull(propagatedClassifications);
    assertClassificationEquals(propagatedClassifications, PII_tag1);
    assertClassificationEquals(propagatedClassifications, PII_tag2);
    assertClassificationEquals(propagatedClassifications, PII_tag3);
    assertTrue(blockedClassifications.isEmpty());

    // block PII tag propagating from employees1 and employees2
    PII_tag2.setEntityGuid(employees1.getGuid());
    PII_tag3.setEntityGuid(employees2.getGuid());

    process3_employee_union_relationship.setBlockedPropagatedClassifications(new HashSet<>(Arrays.asList(PII_tag2, PII_tag3)));
    relationshipStore.update(process3_employee_union_relationship);

    process3_employee_union_relationship = getRelationship(EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);
    propagatedClassifications            = process3_employee_union_relationship.getPropagatedClassifications();
    blockedClassifications               = process3_employee_union_relationship.getBlockedPropagatedClassifications();

    assertClassificationEquals(propagatedClassifications, PII_tag1);
    assertTrue(!blockedClassifications.isEmpty());
    assertClassificationEquals(blockedClassifications, PII_tag2);
    assertClassificationEquals(blockedClassifications, PII_tag3);

    assertClassificationNotExistInEntity(EMPLOYEES_UNION_TABLE, PII_tag2);
    assertClassificationNotExistInEntity(EMPLOYEES_UNION_TABLE, PII_tag3);

    // assert only PII from hdfs_path is propagated to employees_union, PII from employees1 and employees2 is blocked.
    assertEquals(getEntity(EMPLOYEES_UNION_TABLE).getClassifications().size(), 1);
    assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, PII_tag1.getTypeName(), hdfs_path.getGuid());
}
 

@Test(dependsOnMethods = {"addBlockedPropagatedClassifications"})
public void removeBlockedPropagatedClassifications  () throws AtlasBaseException {
    AtlasEntity hdfs_path  = getEntity(HDFS_PATH_EMPLOYEES);
    AtlasEntity employees1 = getEntity(EMPLOYEES1_TABLE);
    AtlasEntity employees2 = getEntity(EMPLOYEES2_TABLE);

    AtlasClassification PII_tag1 = new AtlasClassification("PII"); PII_tag1.setPropagate(true); PII_tag1.setEntityGuid(hdfs_path.getGuid());
        PII_tag1.setAttribute("type", "from hdfs_path entity");
        PII_tag1.setAttribute("valid", true);

    AtlasClassification PII_tag2 = new AtlasClassification("PII"); PII_tag2.setPropagate(true); PII_tag2.setEntityGuid(employees1.getGuid());
        PII_tag2.setAttribute("type", "from employees1 entity");
        PII_tag2.setAttribute("valid", true);

    AtlasClassification PII_tag3 = new AtlasClassification("PII"); PII_tag3.setPropagate(true); PII_tag3.setEntityGuid(employees2.getGuid());
        PII_tag3.setAttribute("type", "from employees2 entity");
        PII_tag3.setAttribute("valid", true);

    AtlasRelationship process3_employee_union_relationship = getRelationship(EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);

    // remove blocked propagated classification entry for PII (from employees2) - allow PII from employees2 to propagate to employee_union
    process3_employee_union_relationship.setBlockedPropagatedClassifications(new HashSet<>(Arrays.asList(PII_tag3)));
    relationshipStore.update(process3_employee_union_relationship);

    process3_employee_union_relationship               = getRelationship(EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);
    Set<AtlasClassification> propagatedClassifications = process3_employee_union_relationship.getPropagatedClassifications();
    Set<AtlasClassification> blockedClassifications    = process3_employee_union_relationship.getBlockedPropagatedClassifications();

    assertClassificationExistInList(propagatedClassifications, PII_tag1);
    assertClassificationExistInList(propagatedClassifications, PII_tag2);
    assertClassificationExistInList(blockedClassifications, PII_tag3);

    // remove all blocked propagated classification entry
    process3_employee_union_relationship.setBlockedPropagatedClassifications(Collections.emptySet());
    relationshipStore.update(process3_employee_union_relationship);

    process3_employee_union_relationship = getRelationship(EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);
    propagatedClassifications            = process3_employee_union_relationship.getPropagatedClassifications();
    blockedClassifications               = process3_employee_union_relationship.getBlockedPropagatedClassifications();

    assertClassificationExistInList(propagatedClassifications, PII_tag1);
    assertClassificationExistInList(propagatedClassifications, PII_tag2);
    assertClassificationExistInList(propagatedClassifications, PII_tag3);
    assertTrue(blockedClassifications.isEmpty());
}
 

@Test
public void updateClassificationsShouldNotThrowExceptionIfResponseIs204() {

    AtlasClientV2 atlasClient = new AtlasClientV2(service, configuration);


    AtlasClassification atlasClassification = new AtlasClassification("Testdb");
    atlasClassification.setEntityGuid("abb672b1-e4bd-402d-a98f-73cd8f775e2a");
    WebResource.Builder builder = setupBuilder(AtlasClientV2.API_V2.UPDATE_CLASSIFICATIONS, service);


    ClientResponse response = mock(ClientResponse.class);
    when(response.getStatus()).thenReturn(Response.Status.NO_CONTENT.getStatusCode());


    when(builder.method(anyString(), Matchers.<Class>any(), anyString())).thenReturn(response);

    try {
        atlasClient.updateClassifications("abb672b1-e4bd-402d-a98f-73cd8f775e2a", Collections.singletonList(atlasClassification));

    } catch (AtlasServiceException e) {
        Assert.fail("Failed with Exception");
    }

}
 

@Test
public void updateClassificationsShouldThrowExceptionIfResponseIsNot204() {

    AtlasClientV2 atlasClient = new AtlasClientV2(service, configuration);


    AtlasClassification atlasClassification = new AtlasClassification("Testdb");
    atlasClassification.setEntityGuid("abb672b1-e4bd-402d-a98f-73cd8f775e2a");
    WebResource.Builder builder = setupBuilder(AtlasClientV2.API_V2.UPDATE_CLASSIFICATIONS, service);


    ClientResponse response = mock(ClientResponse.class);
    when(response.getStatus()).thenReturn(Response.Status.OK.getStatusCode());


    when(builder.method(anyString(), Matchers.<Class>any(), anyString())).thenReturn(response);

    try {
        atlasClient.updateClassifications("abb672b1-e4bd-402d-a98f-73cd8f775e2a", Collections.singletonList(atlasClassification));
        Assert.fail("Failed with Exception");
    } catch (AtlasServiceException e) {
        Assert.assertTrue(e.getMessage().contains(" failed with status 200 "));
    }

}
 

/** This test uses the lineage graph:
 *

 Lineage - 1
 -----------
 [Process1] ----> [Employees1]
 /                              \
 /                                \
 [hdfs_employees]                                  [Process3] ----> [ EmployeesUnion ]
 \                                /
 \                             /
 [Process2] ----> [Employees2]


 Lineage - 2
 -----------

 [Employees] ----> [Process] ----> [ US_Employees ]


 Lineage - 3
 -----------

 [Orders] ----> [Process] ----> [ US_Orders ]

 */

@Test
public void addClassification_PropagateFalse() throws AtlasBaseException {
    AtlasEntity         hdfs_employees = getEntity(HDFS_PATH_EMPLOYEES);
    AtlasClassification tag2           = new AtlasClassification("tag2"); tag2.setPropagate(false); tag2.setEntityGuid(hdfs_employees.getGuid());

    // add classification with propagate to 'false'
    addClassification(hdfs_employees, tag2);

    List<String> propagatedToEntities = Arrays.asList(EMPLOYEES1_PROCESS, EMPLOYEES2_PROCESS, EMPLOYEES1_TABLE,
                                                      EMPLOYEES2_TABLE, EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);

    assertClassificationNotExistInEntities(propagatedToEntities, tag2);
}