org.apache.atlas.model.discovery.SearchParameters Java Examples

@Test(priority = -1)
public void searchWithNEQ_stringAttr() throws AtlasBaseException {
    String expectedEntityName = "hive_Table_Null_tableType";
    createDummyEntity(expectedEntityName,HIVE_TABLE_TYPE);
    SearchParameters params = new SearchParameters();
    params.setTypeName(HIVE_TABLE_TYPE);
    SearchParameters.FilterCriteria filterCriteria = getSingleFilterCondition("tableType", SearchParameters.Operator.NEQ, "Managed");
    params.setEntityFilters(filterCriteria);
    params.setLimit(20);

    SearchContext context = new SearchContext(params, typeRegistry, graph, indexer.getVertexIndexKeys());
    EntitySearchProcessor processor = new EntitySearchProcessor(context);
    List<AtlasVertex> vertices = processor.execute();

    assertEquals(vertices.size(), 3);

    List<String> nameList = new ArrayList<>();
    for (AtlasVertex vertex : vertices) {
        nameList.add((String) entityRetriever.toAtlasEntityHeader(vertex, Collections.singleton("name")).getAttribute("name"));
    }

    assertTrue(nameList.contains(expectedEntityName));
}
 

@Test(dependsOnMethods = "saveSearchesForUser", expectedExceptions = AtlasBaseException.class)
public void attemptToAddExistingSearch() throws AtlasBaseException {
    String           userName                = getIndexBasedUserName(0);
    SearchParameters expectedSearchParameter = getActualSearchParameters();

    for (int j = 0; j < NUM_SEARCHES; j++) {
        String               queryName = getIndexBasedQueryName(j);
        AtlasUserSavedSearch expected  = getDefaultSavedSearch(userName, queryName, expectedSearchParameter);
        AtlasUserSavedSearch actual    = userProfileService.addSavedSearch(expected);

        assertNotNull(actual);
        assertNotNull(actual.getGuid());
        assertEquals(actual.getOwnerName(), expected.getOwnerName());
        assertEquals(actual.getName(), expected.getName());
        assertEquals(actual.getSearchType(), expected.getSearchType());
        assertEquals(actual.getSearchParameters(), expected.getSearchParameters());
    }
}
 

@Test(dependsOnMethods = "verifyQueryConversionFromJSON")
public void addAdditionalSearchesForUser() throws AtlasBaseException {
    SearchParameters expectedSearchParameter = getActualSearchParameters();

    for (int i = 0; i < NUM_USERS; i++) {
        String userName = getIndexBasedUserName(i);

        for (int j = 0; j < 6; j++) {
            String               queryName = getIndexBasedQueryName(NUM_SEARCHES + j);
            AtlasUserSavedSearch actual    = userProfileService.addSavedSearch(getDefaultSavedSearch(userName, queryName, expectedSearchParameter));

            assertNotNull(actual);

            AtlasUserSavedSearch savedSearch = userProfileService.getSavedSearch(userName, queryName);

            assertNotNull(savedSearch);
            assertEquals(savedSearch.getSearchParameters(), expectedSearchParameter);
        }
    }
}
 

@Test(dependsOnMethods = { "addAdditionalSearchesForUser"})
public void updateSearch() throws AtlasBaseException {
    String               userName  = getIndexBasedUserName(0);
    String               queryName = getIndexBasedQueryName(0);
    AtlasUserSavedSearch expected  = userProfileService.getSavedSearch(userName, queryName);

    assertNotNull(expected);

    SearchParameters sp = expected.getSearchParameters();

    sp.setClassification("new-classification");

    AtlasUserSavedSearch actual = userProfileService.updateSavedSearch(expected);

    assertNotNull(actual);
    assertNotNull(actual.getSearchParameters());
    assertEquals(actual.getSearchParameters().getClassification(), expected.getSearchParameters().getClassification());
}
 

private String toIndexExpression(AtlasStructType type, String attrName, SearchParameters.Operator op, String attrVal) {
    String ret = EMPTY_STRING;

    try {
        if (OPERATOR_MAP.get(op) != null) {
            String qualifiedName         = type.getVertexPropertyName(attrName);
            String escapeIndexQueryValue = AtlasAttribute.escapeIndexQueryValue(attrVal);

            // map '__customAttributes' 'CONTAINS' operator to 'EQ' operator (solr limitation for json serialized string search)
            // map '__customAttributes' value from 'key1=value1' to '\"key1\":\"value1\"' (escape special characters and surround with quotes)
            if (attrName.equals(CUSTOM_ATTRIBUTES_PROPERTY_KEY) && op == SearchParameters.Operator.CONTAINS) {
                ret = String.format(OPERATOR_MAP.get(op), qualifiedName, getCustomAttributeIndexQueryValue(escapeIndexQueryValue, false));
            } else {
                ret = String.format(OPERATOR_MAP.get(op), qualifiedName, escapeIndexQueryValue);
            }
        }
    } catch (AtlasBaseException ex) {
        LOG.warn(ex.getMessage());
    }

    return ret;
}
 

@Test(dependsOnMethods = "createsNewProfile")
public void saveSearchesForUser() throws AtlasBaseException {
    String           userName         = getIndexBasedUserName(0);
    SearchParameters searchParameters = getActualSearchParameters();

    for (int i = 0; i < NUM_SEARCHES; i++) {
        userProfileService.addSavedSearch(getDefaultSavedSearch(userName, getIndexBasedQueryName(i), searchParameters));
    }

    for (int i = 0; i < NUM_SEARCHES; i++) {
        AtlasUserSavedSearch savedSearch = userProfileService.getSavedSearch(userName, getIndexBasedQueryName(i));

        assertEquals(savedSearch.getName(), getIndexBasedQueryName(i));
        assertEquals(savedSearch.getSearchParameters(), searchParameters);
    }
}
 

public FreeTextSearchProcessor(SearchContext context) {
    super(context);

    SearchParameters searchParameters = context.getSearchParameters();
    StringBuilder    queryString      = new StringBuilder();

    queryString.append(searchParameters.getQuery());

    if (CollectionUtils.isNotEmpty(context.getEntityTypeNames()) && context.getEntityTypesQryStr().length() <= MAX_QUERY_STR_LENGTH_TYPES) {
        queryString.append(AND_STR).append(context.getEntityTypesQryStr());
    }

    graphIndexQueryBuilder.addActiveStateQueryFilter(queryString);

    if (CollectionUtils.isNotEmpty(context.getClassificationTypeNames()) && context.getClassificationTypesQryStr().length() <= MAX_QUERY_STR_LENGTH_TYPES) {
        queryString.append(AND_STR).append(context.getClassificationTypesQryStr());
    }

    // just use the query string as is
    LOG.debug("Using query string  '{}'.", queryString);

    indexQuery = context.getGraph().indexQuery(prepareGraphIndexQueryParameters(context, queryString));
}
 

@Test(dependsOnMethods = "searchWithNEQ_stringAttr")
public void searchWithNEQ_doubleAttr() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setTypeName(HIVE_TABLE_TYPE);
    SearchParameters.FilterCriteria filterCriteria = getSingleFilterCondition("retention", SearchParameters.Operator.NEQ, "5");
    params.setEntityFilters(filterCriteria);
    params.setLimit(20);

    SearchContext context = new SearchContext(params, typeRegistry, graph, indexer.getVertexIndexKeys());
    EntitySearchProcessor processor = new EntitySearchProcessor(context);
    List<AtlasVertex> vertices = processor.execute();

    assertEquals(vertices.size(), 1);

    List<String> nameList = new ArrayList<>();
    for (AtlasVertex vertex : vertices) {
        nameList.add((String) entityRetriever.toAtlasEntityHeader(vertex, Collections.singleton("name")).getAttribute("name"));
    }

    assertTrue(nameList.contains("hive_Table_Null_tableType"));
}
 

@Test(dependsOnMethods = "searchWithNEQ_stringAttr")
public void searchWithNEQ_pipeSeperatedAttr() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setTypeName(HIVE_TABLE_TYPE);
    SearchParameters.FilterCriteria filterCriteria = getSingleFilterCondition("__classificationNames", SearchParameters.Operator.NEQ, METRIC_CLASSIFICATION);
    params.setEntityFilters(filterCriteria);
    params.setLimit(20);

    SearchContext context = new SearchContext(params, typeRegistry, graph, indexer.getVertexIndexKeys());
    EntitySearchProcessor processor = new EntitySearchProcessor(context);
    List<AtlasVertex> vertices = processor.execute();

    assertEquals(vertices.size(), 7);

    List<String> nameList = new ArrayList<>();
    for (AtlasVertex vertex : vertices) {
        nameList.add((String) entityRetriever.toAtlasEntityHeader(vertex, Collections.singleton("name")).getAttribute("name"));
    }

    assertTrue(nameList.contains("hive_Table_Null_tableType"));
}
 

@Test(dependsOnMethods = "testSearchByMultiTags")
public void testSearchByTerms() throws Exception {

    // database - phi, felt_classification
    // table1 - phi, classification, term: term1 | table2 - classification, term:term2
    // column - phi
    assignTermTo(term1, tableEntity);
    assignTermTo(term1, tableEntity2);

    searchParameters = new SearchParameters();
    searchParameters.setTermName(term1.getName() + "@testSearchGlossary");
    searchParameters.setClassification(CLASSIFICATION);

    AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);
    Assert.assertNotNull(res.getEntities());
    Assert.assertEquals(res.getEntities().size(), 2);

    searchParameters.setClassification(PII);
    res = discoveryREST.searchWithParameters(searchParameters);
    Assert.assertNull(res.getEntities());

    searchParameters.setClassification(PHI);
    res = discoveryREST.searchWithParameters(searchParameters);
    Assert.assertNotNull(res.getEntities());
    Assert.assertEquals(res.getEntities().size(), 1);
}
 

public SearchContext(SearchParameters searchParameters, AtlasTypeRegistry typeRegistry, AtlasGraph graph, Set<String> indexedKeys) {
    this.searchParameters   = searchParameters;
    this.typeRegistry       = typeRegistry;
    this.graph              = graph;
    this.indexedKeys        = indexedKeys;
    this.entityAttributes   = new HashSet<>();
    this.entityType         = typeRegistry.getEntityTypeByName(searchParameters.getTypeName());
    this.classificationType = typeRegistry.getClassificationTypeByName(searchParameters.getClassification());

    if (needFullTextrocessor()) {
        addProcessor(new FullTextSearchProcessor(this));
    }

    if (needClassificationProcessor()) {
        addProcessor(new ClassificationSearchProcessor(this));
    }

    if (needEntityProcessor()) {
        addProcessor(new EntitySearchProcessor(this));
    }
}
 

@Test
public void testSearchByMultiSystemAttributes() throws Exception {

    searchParameters = new SearchParameters();
    searchParameters.setTypeName("_ALL_ENTITY_TYPES");
    SearchParameters.FilterCriteria fc = new SearchParameters.FilterCriteria();
    SearchParameters.FilterCriteria subFc1 = new SearchParameters.FilterCriteria();
    SearchParameters.FilterCriteria subFc2 = new SearchParameters.FilterCriteria();

    subFc1.setAttributeName(MODIFICATION_TIMESTAMP_PROPERTY_KEY);
    subFc1.setOperator(SearchParameters.Operator.LT);
    subFc1.setAttributeValue(String.valueOf(System.currentTimeMillis()));

    subFc2.setAttributeName(TIMESTAMP_PROPERTY_KEY);
    subFc2.setOperator(SearchParameters.Operator.LT);
    subFc2.setAttributeValue(String.valueOf(System.currentTimeMillis()));

    fc.setCriterion(Arrays.asList(subFc1, subFc2));
    fc.setCondition(AND);
    searchParameters.setEntityFilters(fc);

    AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);
    Assert.assertNotNull(res.getEntities());
    Assert.assertTrue(res.getEntities().size() > 5);
}
 

private SearchParameters getSearchParameters(AuditSearchParameters auditSearchParameters) throws AtlasBaseException {
    SearchParameters searchParameters = new SearchParameters();
    searchParameters.setTypeName(ENTITY_TYPE_AUDIT_ENTRY);

    SearchParameters.FilterCriteria validFilter = getNonEmptyFilter(auditSearchParameters.getAuditFilters());
    searchParameters.setEntityFilters(validFilter);

    searchParameters.setLimit(auditSearchParameters.getLimit());
    searchParameters.setOffset(auditSearchParameters.getOffset());

    String sortBy = auditSearchParameters.getSortBy();
    validateSortByParameter(sortBy);

    searchParameters.setSortBy(auditSearchParameters.getSortBy());
    searchParameters.setSortOrder(auditSearchParameters.getSortOrder());

    return searchParameters;
}
 

private SearchParameters.FilterCriteria getNonEmptyFilter(SearchParameters.FilterCriteria auditFilter) throws AtlasBaseException {
    SearchParameters.FilterCriteria outCriteria = new SearchParameters.FilterCriteria();
    outCriteria.setCriterion(new ArrayList<>());

    if(auditFilter != null) {
        outCriteria.setCondition(auditFilter.getCondition());
        List<SearchParameters.FilterCriteria> givenFilterCriterion = auditFilter.getCriterion();

        for (SearchParameters.FilterCriteria each : givenFilterCriterion) {
            if (StringUtils.isNotEmpty(each.getAttributeName()) && !AtlasAuditEntryDTO.getAttributes().contains(each.getAttributeName())) {
                throw new AtlasBaseException(AtlasErrorCode.UNKNOWN_ATTRIBUTE, each.getAttributeName(), "Atlas Audit Entry");
            }

            addParameterIfValueNotEmpty(outCriteria, each.getAttributeName(), each.getOperator(), each.getAttributeValue());
        }
    }

    return outCriteria;
}
 

@Override
public AtlasUserSavedSearch from(AtlasEntity entity) {
    AtlasUserSavedSearch savedSearch = new AtlasUserSavedSearch();

    savedSearch.setGuid(entity.getGuid());
    savedSearch.setName((String) entity.getAttribute(PROPERTY_NAME));
    savedSearch.setOwnerName((String) entity.getAttribute(PROPERTY_OWNER_NAME));
    savedSearch.setSearchType(AtlasUserSavedSearch.SavedSearchType.to((String) entity.getAttribute(PROPERTY_SEARCH_TYPE)));

    String jsonSearchParams = (String) entity.getAttribute(PROPERTY_SEARCH_PARAMETERS);

    if (StringUtils.isNotEmpty(jsonSearchParams)) {
        savedSearch.setSearchParameters(AtlasType.fromJson(jsonSearchParams, SearchParameters.class));
    }

    savedSearch.setUiParameters((String) entity.getAttribute(PROPERTY_UI_PARAMETERS));

    return savedSearch;
}
 

@Test(priority = -1)
public void searchByALLTag() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setClassification(ALL_CLASSIFICATION_TYPES);

    List<AtlasEntityHeader> entityHeaders = discoveryService.searchWithParameters(params).getEntities();

    Assert.assertTrue(CollectionUtils.isNotEmpty(entityHeaders));

    totalEntities = getEntityCount();
    totalClassifiedEntities = entityHeaders.size();
    getTotalClassifiedEntitiesHistorical = getEntityWithTagCountHistorical();
}
 

@Test
public void searchByTag() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setClassification(DIMENSION_CLASSIFICATION);

    List<AtlasEntityHeader> entityHeaders = discoveryService.searchWithParameters(params).getEntities();

    assertEquals(entityHeaders.size(), dimensionTagEntities);
}
 

@Test
public void searchByNameSortBy() throws AtlasBaseException, InterruptedException {
    SearchParameters params = new SearchParameters();
    params.setTypeName("hive_table");
    params.setQuery("sales");
    params.setExcludeDeletedEntities(true);
    params.setLimit(3);
    params.setSortBy("owner");
    params.setSortOrder(SortOrder.ASCENDING);

    SearchContext context = new SearchContext(params, typeRegistry, graph, Collections.<String>emptySet());

    FreeTextSearchProcessor processor = new FreeTextSearchProcessor(context);

    List<AtlasVertex> vertices = processor.execute();

    assertEquals(processor.getResultCount(), 3);
    assertEquals(vertices.size(), 3);

    AtlasVertex firstVertex = vertices.get(0);

    String firstOwner = (String) entityRetriever.toAtlasEntityHeader(firstVertex, Sets.newHashSet(params.getSortBy())).getAttribute(params.getSortBy());

    AtlasVertex secondVertex = vertices.get(1);
    String secondOwner = (String) entityRetriever.toAtlasEntityHeader(secondVertex, Sets.newHashSet(params.getSortBy())).getAttribute(params.getSortBy());

    assertEquals(firstOwner, "Jane BI");
    assertEquals(secondOwner, "Joe");
}
 

private String toSolrExpression(AtlasStructType type, String attrName, SearchParameters.Operator op, String attrVal) {
    String ret = EMPTY_STRING;

    try {
        if (OPERATOR_MAP.get(op) != null) {
            String qualifiedName = type.getQualifiedAttributeName(attrName);

            ret = String.format(OPERATOR_MAP.get(op), qualifiedName, AtlasStructType.AtlasAttribute.escapeIndexQueryValue(attrVal));
        }
    } catch (AtlasBaseException ex) {
        LOG.warn(ex.getMessage());
    }

    return ret;
}
 

private void processSearchParametersForMultipleTypeNames(String fileName, AtlasSolrQueryBuilder underTest) throws IOException, AtlasBaseException {
    ObjectMapper mapper = new ObjectMapper();
    SearchParameters searchParameters =  mapper.readValue(new FileInputStream(fileName), SearchParameters.class);

    Set<AtlasEntityType> hiveTableEntityTypeMocks = new HashSet<>();
    hiveTableEntityTypeMocks.add(hiveTableEntityTypeMock);
    hiveTableEntityTypeMocks.add(hiveTableEntityTypeMock2);
    underTest.withEntityTypes(hiveTableEntityTypeMocks)
            .withQueryString(searchParameters.getQuery())
            .withCriteria(searchParameters.getEntityFilters())
            .withExcludedDeletedEntities(searchParameters.getExcludeDeletedEntities())
            .withCommonIndexFieldNames(indexFieldNamesMap);
}
 

@Test
public void emptySearch() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setTypeName("hive_table");
    params.setQuery("not_exists");
    params.setExcludeDeletedEntities(true);
    params.setLimit(3);

    SearchContext context = new SearchContext(params, typeRegistry, graph, Collections.<String>emptySet());

    FreeTextSearchProcessor processor = new FreeTextSearchProcessor(context);

    assertEquals(processor.getResultCount(), 0);
    assertEquals(processor.execute().size(), 0);
}
 

@Test
public void searchByALLTagAndIndexSysFilters() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setClassification(ALL_CLASSIFICATION_TYPES);
    FilterCriteria filterCriteria = getSingleFilterCondition("__timestamp", Operator.LT, String.valueOf(System.currentTimeMillis()));
    params.setTagFilters(filterCriteria);

    List<AtlasEntityHeader> entityHeaders = discoveryService.searchWithParameters(params).getEntities();

    assertEquals(entityHeaders.size(), totalClassifiedEntities);
}
 

@Test
public void searchByTagAndIndexSysFilters() throws AtlasBaseException {

    SearchParameters params = new SearchParameters();
    params.setClassification(DIMENSION_CLASSIFICATION);
    FilterCriteria filterCriteria = getSingleFilterCondition("__timestamp", Operator.LT, String.valueOf(System.currentTimeMillis()));
    params.setTagFilters(filterCriteria);

    List<AtlasEntityHeader> entityHeaders = discoveryService.searchWithParameters(params).getEntities();

    assertEquals(entityHeaders.size(), dimensionTagEntities);
}
 

@Test
public void entityType() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setTypeName(DATABASE_TYPE);
    params.setLimit(20);

    SearchContext context = new SearchContext(params, typeRegistry, graph, Collections.<String>emptySet());

    EntitySearchProcessor processor = new EntitySearchProcessor(context);
    assertEquals(processor.execute().size(), 3);
}
 

@Test
public void ALLEntityTypeWithTag() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setTypeName(SearchParameters.ALL_ENTITY_TYPES);
    params.setClassification(FACT_CLASSIFICATION);
    params.setLimit(20);

    SearchContext context = new SearchContext(params, typeRegistry, graph, Collections.<String>emptySet());

    EntitySearchProcessor processor = new EntitySearchProcessor(context);
    assertEquals(processor.execute().size(), 5);
}
 

@Test
public void ALLEntityType() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setTypeName(SearchParameters.ALL_ENTITY_TYPES);
    params.setLimit(20);

    SearchContext context = new SearchContext(params, typeRegistry, graph, Collections.<String>emptySet());

    EntitySearchProcessor processor = new EntitySearchProcessor(context);
    assertEquals(processor.execute().size(), 20);
}
 

@Test(expectedExceptions = AtlasBaseException.class, expectedExceptionsMessageRegExp = "NotExisting: Unknown/invalid classification")
public void searchByNonExistingClassification() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setTypeName("hive_process");
    params.setClassification("NotExisting");
    params.setLimit(10);

    SearchContext context = new SearchContext(params, typeRegistry, graph, Collections.<String>emptySet());
    new EntitySearchProcessor(context);
}
 

@Test
public void emptySearchByClassification() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setTypeName("hive_table");
    params.setClassification("PII");
    params.setLimit(10);

    SearchContext context = new SearchContext(params, typeRegistry, graph, Collections.<String>emptySet());

    EntitySearchProcessor processor = new EntitySearchProcessor(context);

    assertEquals(processor.getResultCount(), 0);
    assertEquals(processor.execute().size(), 0);
}
 

public void searchByTagAndGraphSysFilters() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setClassification(DIMENSION_CLASSIFICATION);
    FilterCriteria filterCriteria = getSingleFilterCondition("__entityStatus", Operator.EQ, "DELETED");
    params.setTagFilters(filterCriteria);
    params.setExcludeDeletedEntities(false);

    List<AtlasEntityHeader> entityHeaders = discoveryService.searchWithParameters(params).getEntities();

    assertEquals(entityHeaders.size(), 1);
    assertEquals(entityHeaders.get(0).getGuid(), dimensionTagDeleteGuid);

}
 

@Test
public void searchTablesByClassification() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setTypeName("hive_column");
    params.setClassification("PII");
    params.setLimit(10);

    SearchContext context = new SearchContext(params, typeRegistry, graph, Collections.<String>emptySet());

    EntitySearchProcessor processor = new EntitySearchProcessor(context);

    assertEquals(processor.getResultCount(), 4);
    assertEquals(processor.execute().size(), 4);
}