Examples with AtlasSearchResult org.apache.atlas.model.discovery.AtlasSearchResult used on opensource projects

Example 1 with AtlasSearchResult

use of org.apache.atlas.model.discovery.AtlasSearchResult in project ranger by apache.

the class AtlasRESTTagSource method getAtlasActiveEntities.

private List<RangerAtlasEntityWithTags> getAtlasActiveEntities() {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> getAtlasActiveEntities()");
    }
    List<RangerAtlasEntityWithTags> ret = new ArrayList<>();
    AtlasClientV2 atlasClient = null;
    try {
        atlasClient = getAtlasClient();
    } catch (IOException exception) {
        LOG.error("Failed to get Atlas client.", exception);
    }
    if (atlasClient != null) {
        SearchParameters searchParams = new SearchParameters();
        searchParams.setExcludeDeletedEntities(true);
        searchParams.setClassification("*");
        // searchParams.setIncludeSubClassifications(true);
        // searchParams.setIncludeSubTypes(true);
        searchParams.setIncludeClassificationAttributes(true);
        searchParams.setLimit(REQUESTED_ENTITIES_LIMIT_MAX);
        boolean isMoreData;
        int nextStartIndex = 0;
        do {
            AtlasTypeRegistry typeRegistry = new AtlasTypeRegistry();
            AtlasTypeRegistry.AtlasTransientTypeRegistry tty = null;
            AtlasSearchResult searchResult = null;
            boolean commitUpdates = false;
            searchParams.setOffset(nextStartIndex);
            isMoreData = false;
            try {
                searchResult = atlasClient.facetedSearch(searchParams);
                AtlasTypesDef typesDef = atlasClient.getAllTypeDefs(new SearchFilter());
                tty = typeRegistry.lockTypeRegistryForUpdate();
                tty.addTypes(typesDef);
                commitUpdates = true;
            } catch (AtlasServiceException | AtlasBaseException excp) {
                LOG.error("failed to download tags from Atlas", excp);
                ret = null;
            } catch (Exception unexpectedException) {
                LOG.error("Failed to download tags from Atlas due to unexpected exception", unexpectedException);
                ret = null;
            } finally {
                if (tty != null) {
                    typeRegistry.releaseTypeRegistryForUpdate(tty, commitUpdates);
                }
            }
            if (commitUpdates && searchResult != null) {
                if (LOG.isDebugEnabled()) {
                    LOG.debug(AtlasType.toJson(searchResult));
                }
                List<AtlasEntityHeader> entityHeaders = searchResult.getEntities();
                if (CollectionUtils.isNotEmpty(entityHeaders)) {
                    nextStartIndex += entityHeaders.size();
                    isMoreData = true;
                    for (AtlasEntityHeader header : entityHeaders) {
                        if (!header.getStatus().equals(AtlasEntity.Status.ACTIVE)) {
                            if (LOG.isDebugEnabled()) {
                                LOG.debug("Skipping entity because it is not ACTIVE, header:[" + header + "]");
                            }
                            continue;
                        }
                        String typeName = header.getTypeName();
                        if (!AtlasResourceMapperUtil.isEntityTypeHandled(typeName)) {
                            if (LOG.isDebugEnabled()) {
                                LOG.debug("Not fetching Atlas entities of type:[" + typeName + "]");
                            }
                            continue;
                        }
                        List<EntityNotificationWrapper.RangerAtlasClassification> allTagsForEntity = new ArrayList<>();
                        for (AtlasClassification classification : header.getClassifications()) {
                            List<EntityNotificationWrapper.RangerAtlasClassification> tags = resolveTag(typeRegistry, classification);
                            if (tags != null) {
                                allTagsForEntity.addAll(tags);
                            }
                        }
                        if (CollectionUtils.isNotEmpty(allTagsForEntity)) {
                            RangerAtlasEntity entity = new RangerAtlasEntity(typeName, header.getGuid(), header.getAttributes());
                            RangerAtlasEntityWithTags entityWithTags = new RangerAtlasEntityWithTags(entity, allTagsForEntity, typeRegistry);
                            ret.add(entityWithTags);
                        }
                    }
                }
            }
        } while (isMoreData);
    }
    if (LOG.isDebugEnabled()) {
        LOG.debug("<== getAtlasActiveEntities()");
    }
    return ret;
}

Example 2 with AtlasSearchResult

use of org.apache.atlas.model.discovery.AtlasSearchResult in project atlas by apache.

the class DSLQueriesTest method pollForData.

private void pollForData() throws InterruptedException {
    Object[][] basicVerificationQueries = new Object[][] { { "hive_db", 3 }, { "hive_process", 7 }, { "hive_table", 10 }, { "hive_column", 17 }, { "hive_storagedesc", 1 }, { "Manager", 2 }, { "Employee", 4 } };
    int pollingAttempts = 5;
    // in msecs
    int pollingBackoff = 0;
    boolean success;
    for (int attempt = 0; attempt < pollingAttempts; attempt++, pollingBackoff += attempt * 5000) {
        LOG.debug("Polling -- Attempt {}, Backoff {}", attempt, pollingBackoff);
        success = false;
        for (Object[] verificationQuery : basicVerificationQueries) {
            String query = (String) verificationQuery[0];
            int expected = (int) verificationQuery[1];
            try {
                AtlasSearchResult result = discoveryService.searchUsingDslQuery(query, 25, 0);
                if (result.getEntities() == null || result.getEntities().isEmpty()) {
                    LOG.warn("DSL {} returned no entities", query);
                    success = false;
                } else if (result.getEntities().size() != expected) {
                    LOG.warn("DSL {} returned unexpected number of entities. Expected {} Actual {}", query, expected, result.getEntities().size());
                    success = false;
                } else {
                    success = true;
                }
            } catch (AtlasBaseException e) {
                LOG.error("Got exception for DSL {}, errorCode: {}", query, e.getAtlasErrorCode());
                waitOrBailout(pollingAttempts, pollingBackoff, attempt);
            }
        }
        // DSL queries were successful
        if (success) {
            LOG.info("Polling was success");
            break;
        } else {
            waitOrBailout(pollingAttempts, pollingBackoff, attempt);
        }
    }
}

Example 3 with AtlasSearchResult

use of org.apache.atlas.model.discovery.AtlasSearchResult in project atlas by apache.

the class DSLQueriesTest method comparison.

@Test(dataProvider = "comparisonQueriesProvider")
public void comparison(String query, int expected) throws AtlasBaseException {
    LOG.debug(query);
    AtlasSearchResult searchResult = discoveryService.searchUsingDslQuery(query, DEFAULT_LIMIT, 0);
    assertSearchResult(searchResult, expected, query);
    AtlasSearchResult searchResult2 = discoveryService.searchUsingDslQuery(query.replace("where", " "), DEFAULT_LIMIT, 0);
    assertSearchResult(searchResult2, expected, query);
}

Example 4 with AtlasSearchResult

use of org.apache.atlas.model.discovery.AtlasSearchResult in project atlas by apache.

the class EntityDiscoveryService method searchWithParameters.

@Override
@GraphTransaction
public AtlasSearchResult searchWithParameters(SearchParameters searchParameters) throws AtlasBaseException {
    AtlasSearchResult ret = new AtlasSearchResult(searchParameters);
    final QueryParams params = QueryParams.getNormalizedParams(searchParameters.getLimit(), searchParameters.getOffset());
    searchParameters.setLimit(params.limit());
    searchParameters.setOffset(params.offset());
    SearchContext context = new SearchContext(searchParameters, typeRegistry, graph, indexer.getVertexIndexKeys());
    // For future cancellations
    String searchID = searchTracker.add(context);
    try {
        List<AtlasVertex> resultList = context.getSearchProcessor().execute();
        // By default any attribute that shows up in the search parameter should be sent back in the response
        // If additional values are requested then the entityAttributes will be a superset of the all search attributes
        // and the explicitly requested attribute(s)
        Set<String> resultAttributes = new HashSet<>();
        Set<String> entityAttributes = new HashSet<>();
        if (CollectionUtils.isNotEmpty(searchParameters.getAttributes())) {
            resultAttributes.addAll(searchParameters.getAttributes());
        }
        if (CollectionUtils.isNotEmpty(context.getEntityAttributes())) {
            resultAttributes.addAll(context.getEntityAttributes());
        }
        AtlasEntityType entityType = context.getEntityType();
        if (entityType != null) {
            for (String resultAttribute : resultAttributes) {
                AtlasAttribute attribute = entityType.getAttribute(resultAttribute);
                if (attribute != null) {
                    AtlasType attributeType = attribute.getAttributeType();
                    if (attributeType instanceof AtlasArrayType) {
                        attributeType = ((AtlasArrayType) attributeType).getElementType();
                    }
                    if (attributeType instanceof AtlasEntityType || attributeType instanceof AtlasObjectIdType) {
                        entityAttributes.add(resultAttribute);
                    }
                }
            }
        }
        for (AtlasVertex atlasVertex : resultList) {
            AtlasEntityHeader entity = entityRetriever.toAtlasEntityHeader(atlasVertex, resultAttributes);
            if (searchParameters.getIncludeClassificationAttributes()) {
                entity.setClassifications(entityRetriever.getAllClassifications(atlasVertex));
            }
            ret.addEntity(entity);
            // populate ret.referredEntities
            for (String entityAttribute : entityAttributes) {
                Object attrValue = entity.getAttribute(entityAttribute);
                if (attrValue instanceof AtlasObjectId) {
                    AtlasObjectId objId = (AtlasObjectId) attrValue;
                    if (ret.getReferredEntities() == null) {
                        ret.setReferredEntities(new HashMap<String, AtlasEntityHeader>());
                    }
                    if (!ret.getReferredEntities().containsKey(objId.getGuid())) {
                        ret.getReferredEntities().put(objId.getGuid(), entityRetriever.toAtlasEntityHeader(objId.getGuid()));
                    }
                } else if (attrValue instanceof Collection) {
                    Collection objIds = (Collection) attrValue;
                    for (Object obj : objIds) {
                        if (obj instanceof AtlasObjectId) {
                            AtlasObjectId objId = (AtlasObjectId) obj;
                            if (ret.getReferredEntities() == null) {
                                ret.setReferredEntities(new HashMap<String, AtlasEntityHeader>());
                            }
                            if (!ret.getReferredEntities().containsKey(objId.getGuid())) {
                                ret.getReferredEntities().put(objId.getGuid(), entityRetriever.toAtlasEntityHeader(objId.getGuid()));
                            }
                        }
                    }
                }
            }
        }
    } finally {
        searchTracker.remove(searchID);
    }
    return ret;
}

Example 5 with AtlasSearchResult

use of org.apache.atlas.model.discovery.AtlasSearchResult in project atlas by apache.

the class EntityDiscoveryService method searchUsingDslQuery.

@Override
@GraphTransaction
public AtlasSearchResult searchUsingDslQuery(String dslQuery, int limit, int offset) throws AtlasBaseException {
    AtlasSearchResult ret = new AtlasSearchResult(dslQuery, AtlasQueryType.DSL);
    GremlinQuery gremlinQuery = toGremlinQuery(dslQuery, limit, offset);
    String queryStr = gremlinQuery.queryStr();
    if (LOG.isDebugEnabled()) {
        LOG.debug("Executing DSL: query={}, gremlinQuery={}", dslQuery, queryStr);
    }
    Object result = graph.executeGremlinScript(queryStr, false);
    if (result instanceof List && CollectionUtils.isNotEmpty((List) result)) {
        List queryResult = (List) result;
        Object firstElement = queryResult.get(0);
        if (firstElement instanceof AtlasVertex) {
            for (Object element : queryResult) {
                if (element instanceof AtlasVertex) {
                    ret.addEntity(entityRetriever.toAtlasEntityHeader((AtlasVertex) element));
                } else {
                    LOG.warn("searchUsingDslQuery({}): expected an AtlasVertex; found unexpected entry in result {}", dslQuery, element);
                }
            }
        } else if (gremlinQuery.hasSelectList()) {
            ret.setAttributes(toAttributesResult(queryResult, gremlinQuery));
        } else if (firstElement instanceof Map) {
            for (Object element : queryResult) {
                if (element instanceof Map) {
                    Map map = (Map) element;
                    for (Object key : map.keySet()) {
                        Object value = map.get(key);
                        if (value instanceof List && CollectionUtils.isNotEmpty((List) value)) {
                            for (Object o : (List) value) {
                                Object entry = o;
                                if (entry instanceof AtlasVertex) {
                                    ret.addEntity(entityRetriever.toAtlasEntityHeader((AtlasVertex) entry));
                                }
                            }
                        }
                    }
                }
            }
        } else {
            LOG.warn("searchUsingDslQuery({}/{}): found unexpected entry in result {}", dslQuery, dslQuery, gremlinQuery.queryStr());
        }
    }
    return ret;
}