package no.mnemonic.act.platform.dao.elastic;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.ObjectReader;
import com.fasterxml.jackson.databind.ObjectWriter;
import com.fasterxml.jackson.databind.json.JsonMapper;
import com.google.common.io.CharStreams;
import no.mnemonic.act.platform.dao.api.criteria.FactSearchCriteria;
import no.mnemonic.act.platform.dao.api.criteria.ObjectStatisticsCriteria;
import no.mnemonic.act.platform.dao.api.result.ObjectStatisticsContainer;
import no.mnemonic.act.platform.dao.elastic.criteria.FactExistenceSearchCriteria;
import no.mnemonic.act.platform.dao.elastic.document.FactDocument;
import no.mnemonic.act.platform.dao.elastic.document.ObjectDocument;
import no.mnemonic.act.platform.dao.elastic.result.ScrollingSearchResult;
import no.mnemonic.act.platform.dao.elastic.result.SearchResult;
import no.mnemonic.commons.component.Dependency;
import no.mnemonic.commons.component.LifecycleAspect;
import no.mnemonic.commons.logging.Logger;
import no.mnemonic.commons.logging.Logging;
import no.mnemonic.commons.utilities.ObjectUtils;
import no.mnemonic.commons.utilities.StringUtils;
import no.mnemonic.commons.utilities.collections.CollectionUtils;
import no.mnemonic.commons.utilities.collections.ListUtils;
import no.mnemonic.commons.utilities.collections.SetUtils;
import org.apache.lucene.search.join.ScoreMode;
import org.elasticsearch.ElasticsearchException;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.DocWriteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.*;
import org.elasticsearch.action.support.WriteRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.indices.CreateIndexRequest;
import org.elasticsearch.client.indices.CreateIndexResponse;
import org.elasticsearch.client.indices.GetIndexRequest;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.SimpleQueryStringBuilder;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.aggregations.Aggregation;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.aggregations.HasAggregations;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.metrics.cardinality.Cardinality;
import org.elasticsearch.search.aggregations.metrics.max.Max;
import org.elasticsearch.search.aggregations.metrics.tophits.TopHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import javax.inject.Inject;
import javax.inject.Singleton;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.List;
import java.util.Set;
import java.util.UUID;
import java.util.function.Function;
import java.util.stream.Collectors;
import static org.elasticsearch.common.bytes.BytesReference.toBytes;
import static org.elasticsearch.index.query.QueryBuilders.*;
import static org.elasticsearch.search.aggregations.AggregationBuilders.*;
/**
* Class for indexing Facts into ElasticSearch as well as for retrieving and searching indexed Facts.
*/
@Singleton
public class FactSearchManager implements LifecycleAspect {
private static final String INDEX_NAME = "act";
private static final String TYPE_NAME = "_doc";
private static final String MAPPINGS_JSON = "mappings.json";
private static final int MAX_RESULT_WINDOW = 10_000; // Must be the same value as specified in mappings.json.
private static final String FILTER_FACTS_AGGREGATION_NAME = "FilterFactsAggregation";
private static final String NESTED_OBJECTS_AGGREGATION_NAME = "NestedObjectsAggregation";
private static final String FILTER_OBJECTS_AGGREGATION_NAME = "FilterObjectsAggregation";
private static final String OBJECTS_COUNT_AGGREGATION_NAME = "ObjectsCountAggregation";
private static final String UNIQUE_OBJECTS_AGGREGATION_NAME = "UniqueObjectsAggregation";
private static final String UNIQUE_OBJECTS_SOURCE_AGGREGATION_NAME = "UniqueObjectsSourceAggregation";
private static final String REVERSED_FACTS_AGGREGATION_NAME = "ReversedFactsAggregation";
private static final String UNIQUE_FACT_TYPES_AGGREGATION_NAME = "UniqueFactTypesAggregation";
private static final String MAX_LAST_ADDED_TIMESTAMP_AGGREGATION_NAME = "MaxLastAddedTimestampAggregation";
private static final String MAX_LAST_SEEN_TIMESTAMP_AGGREGATION_NAME = "MaxLastSeenTimestampAggregation";
private static final float CONFIDENCE_EQUALITY_INTERVAL = 0.01f;
private static final Logger LOGGER = Logging.getLogger(FactSearchManager.class);
private static final ObjectMapper MAPPER = JsonMapper.builder().build();
private static final ObjectReader FACT_DOCUMENT_READER = MAPPER.readerFor(FactDocument.class);
private static final ObjectReader OBJECT_DOCUMENT_READER = MAPPER.readerFor(ObjectDocument.class);
private static final ObjectWriter FACT_DOCUMENT_WRITER = MAPPER.writerFor(FactDocument.class);
@Dependency
private final ClientFactory clientFactory;
private String searchScrollExpiration = "1m";
private int searchScrollSize = 1000;
private boolean isTestEnvironment = false;
@Inject
public FactSearchManager(ClientFactory clientFactory) {
this.clientFactory = clientFactory;
}
@Override
public void startComponent() {
if (!indexExists()) {
LOGGER.info("Index '%s' does not exist, create it.", INDEX_NAME);
createIndex();
}
}
@Override
public void stopComponent() {
// NOOP
}
/**
* Retrieve an indexed Fact by its UUID. Returns NULL if Fact cannot be fetched from ElasticSearch.
*
* @param id UUID of indexed Fact
* @return Indexed Fact or NULL if not available
*/
public FactDocument getFact(UUID id) {
if (id == null) return null;
GetResponse response;
try {
GetRequest request = new GetRequest(INDEX_NAME, TYPE_NAME, id.toString());
response = clientFactory.getClient().get(request, RequestOptions.DEFAULT);
} catch (ElasticsearchException | IOException ex) {
throw logAndExit(ex, String.format("Could not perform request to fetch Fact with id = %s.", id));
}
if (response.isExists()) {
LOGGER.info("Successfully fetched Fact with id = %s.", id);
return decodeFactDocument(id, response.getSourceAsBytes());
} else {
// Fact isn't indexed in ElasticSearch, log warning and return null.
LOGGER.warning("Could not fetch Fact with id = %s. Fact not indexed?", id);
return null;
}
}
/**
* Index a Fact into ElasticSearch.
*
* @param fact Fact to index
* @return Indexed Fact
*/
public FactDocument indexFact(FactDocument fact) {
if (fact == null || fact.getId() == null) return null;
IndexResponse response;
try {
IndexRequest request = new IndexRequest(INDEX_NAME, TYPE_NAME, fact.getId().toString())
.setRefreshPolicy(isTestEnvironment ? WriteRequest.RefreshPolicy.IMMEDIATE : WriteRequest.RefreshPolicy.NONE)
.source(FACT_DOCUMENT_WRITER.writeValueAsBytes(fact), XContentType.JSON);
response = clientFactory.getClient().index(request, RequestOptions.DEFAULT);
} catch (ElasticsearchException | IOException ex) {
throw logAndExit(ex, String.format("Could not perform request to index Fact with id = %s.", fact.getId()));
}
if (response.status() != RestStatus.OK && response.status() != RestStatus.CREATED) {
LOGGER.warning("Could not index Fact with id = %s.", fact.getId());
} else if (response.getResult() == DocWriteResponse.Result.CREATED) {
LOGGER.info("Successfully indexed Fact with id = %s.", fact.getId());
} else if (response.getResult() == DocWriteResponse.Result.UPDATED) {
LOGGER.info("Successfully re-indexed existing Fact with id = %s.", fact.getId());
}
return fact;
}
/**
* Retrieve all Facts which are considered logically the same when matched against a given search criteria, i.e. the
* following condition holds: an indexed Fact matches the search criteria and will be included in the returned result
* if and only if the value, FactType, organization, origin, confidence, access mode and all bound Objects (including
* direction) match exactly. In the case of a meta Fact the inReferenceTo must match instead of bound Objects (Objects
* should be omitted). If no Fact satisfies this condition an empty result container is returned.
* <p>
* This method can be used to determine if a Fact already logically exists in the system, e.g. before a new Fact is
* created. No access control will be performed. This must be done by the caller.
*
* @param criteria Criteria to retrieve existing Facts
* @return All Facts satisfying search criteria wrapped inside a result container
*/
public SearchResult<FactDocument> retrieveExistingFacts(FactExistenceSearchCriteria criteria) {
if (criteria == null) return SearchResult.<FactDocument>builder().build();
SearchResponse response;
try {
response = clientFactory.getClient().search(buildFactExistenceSearchRequest(criteria), RequestOptions.DEFAULT);
} catch (ElasticsearchException | IOException ex) {
throw logAndExit(ex, "Could not perform request to search for existing Facts.");
}
if (response.status() != RestStatus.OK) {
LOGGER.warning("Could not search for existing Facts (response code %s).", response.status());
return SearchResult.<FactDocument>builder().build();
}
List<FactDocument> result = retrieveFactDocuments(response);
LOGGER.info("Successfully retrieved %d existing Facts.", result.size());
return SearchResult.<FactDocument>builder()
.setCount((int) response.getHits().getTotalHits())
.setValues(result)
.build();
}
/**
* Search for Facts indexed in ElasticSearch by a given search criteria. Only Facts satisfying the search criteria
* will be returned. Returns a result container which will stream out the results from ElasticSearch. It will not
* limit the number of returned results. This must be done by the caller. Returns an empty result container if no
* Fact satisfies the search criteria.
* <p>
* Both 'currentUserID' (identifying the calling user) and 'availableOrganizationID' (identifying the Organizations
* the calling user has access to) must be set in the search criteria in order to apply access control to Facts. Only
* Facts accessible to the calling user will be returned.
*
* @param criteria Search criteria to match against Facts
* @return Facts satisfying search criteria wrapped inside a result container
*/
public ScrollingSearchResult<FactDocument> searchFacts(FactSearchCriteria criteria) {
if (criteria == null) return ScrollingSearchResult.<FactDocument>builder().build();
SearchResponse response;
try {
response = clientFactory.getClient().search(buildFactsSearchRequest(criteria), RequestOptions.DEFAULT);
} catch (ElasticsearchException | IOException ex) {
throw logAndExit(ex, "Could not perform request to search for Facts.");
}
if (response.status() != RestStatus.OK) {
LOGGER.warning("Could not search for Facts (response code %s).", response.status());
return ScrollingSearchResult.<FactDocument>builder().build();
}
LOGGER.info("Successfully initiated streaming of search results. Start fetching data.");
return ScrollingSearchResult.<FactDocument>builder()
.setInitialBatch(createFactsBatch(response))
.setFetchNextBatch(this::fetchNextFactsBatch)
.setCount((int) response.getHits().getTotalHits())
.build();
}
/**
* Search for Objects indexed in ElasticSearch by a given search criteria. Only Objects satisfying the search criteria
* will be returned. Returns an empty result container if no Object satisfies the search criteria.
* <p>
* First, the result will be reduced to only the Facts satisfying the search criteria. Then, for all matching Facts
* the bound Objects will be reduced to the unique Objects satisfying the search criteria.
* <p>
* Both 'currentUserID' (identifying the calling user) and 'availableOrganizationID' (identifying the Organizations
* the calling user has access to) must be set in the search criteria in order to apply access control to Facts. Only
* Objects bound to Facts accessible to the calling user will be returned.
*
* @param criteria Search criteria to match against Facts and their bound Objects
* @return Objects satisfying search criteria wrapped inside a result container
*/
public SearchResult<ObjectDocument> searchObjects(FactSearchCriteria criteria) {
if (criteria == null) return SearchResult.<ObjectDocument>builder().build();
SearchResponse response;
try {
response = clientFactory.getClient().search(buildObjectsSearchRequest(criteria), RequestOptions.DEFAULT);
} catch (ElasticsearchException | IOException ex) {
throw logAndExit(ex, "Could not perform request to search for Objects.");
}
if (response.status() != RestStatus.OK) {
LOGGER.warning("Could not search for Objects (response code %s).", response.status());
return SearchResult.<ObjectDocument>builder().setLimit(criteria.getLimit()).build();
}
int count = retrieveSearchObjectsResultCount(response);
List<ObjectDocument> result = retrieveSearchObjectsResultValues(response);
LOGGER.info("Successfully retrieved %d Objects from a total of %d matching Objects.", result.size(), count);
return SearchResult.<ObjectDocument>builder()
.setLimit(criteria.getLimit())
.setCount(count)
.setValues(result)
.build();
}
/**
* Calculate statistics about the Facts bound to Objects. For each Object specified in the statistics criteria it is
* calculated how many Facts of each FactType are bound to the Object and when a Fact of that FactType was last added
* and last seen.
* <p>
* Both 'currentUserID' (identifying the calling user) and 'availableOrganizationID' (identifying the Organizations
* the calling user has access to) must be set in the statistics criteria in order to apply access control to Facts.
* Only statistics for Objects bound to Facts accessible to the calling user will be returned, and only accessible
* Facts will be included in the returned statistics.
*
* @param criteria Criteria to specify for which Objects statistics should be calculated
* @return Result container with the calculated statistics for each Object
*/
public ObjectStatisticsContainer calculateObjectStatistics(ObjectStatisticsCriteria criteria) {
if (criteria == null) return ObjectStatisticsContainer.builder().build();
SearchResponse response;
try {
response = clientFactory.getClient().search(buildObjectStatisticsSearchRequest(criteria), RequestOptions.DEFAULT);
} catch (ElasticsearchException | IOException ex) {
throw logAndExit(ex, "Could not perform request to calculate Object statistics.");
}
if (response.status() != RestStatus.OK) {
LOGGER.warning("Could not calculate Object statistics (response code %s).", response.status());
return ObjectStatisticsContainer.builder().build();
}
ObjectStatisticsContainer result = retrieveObjectStatisticsResult(response);
LOGGER.info("Successfully retrieved statistics for %d Objects.", result.getStatisticsCount());
return result;
}
/**
* Specify if this class is executed during unit tests (defaults to false). This setting will make indexed documents
* available for search immediately.
*
* @param testEnvironment Whether this class is executed during unit tests
* @return Class instance, i.e. 'this'
*/
public FactSearchManager setTestEnvironment(boolean testEnvironment) {
this.isTestEnvironment = testEnvironment;
return this;
}
/**
* Specify how long the search context of a scrolling search will be kept open in ElasticSearch. Defaults to 1 minute.
* <p>
* Accepts an ElasticSearch time unit: https://www.elastic.co/guide/en/elasticsearch/reference/current/common-options.html#time-units
*
* @param searchScrollExpiration Expiration time of search context
* @return Class instance, i.e. 'this'
*/
public FactSearchManager setSearchScrollExpiration(String searchScrollExpiration) {
this.searchScrollExpiration = searchScrollExpiration;
return this;
}
/**
* Specify the batch size when fetching data from ElasticSearch using a scrolling search. Defaults to 1000.
*
* @param searchScrollSize Batch size
* @return Class instance, i.e. 'this'
*/
public FactSearchManager setSearchScrollSize(int searchScrollSize) {
this.searchScrollSize = searchScrollSize;
return this;
}
private boolean indexExists() {
try {
GetIndexRequest request = new GetIndexRequest(INDEX_NAME);
return clientFactory.getClient().indices().exists(request, RequestOptions.DEFAULT);
} catch (ElasticsearchException | IOException ex) {
throw logAndExit(ex, "Could not perform request to verify if index exists.");
}
}
private void createIndex() {
CreateIndexResponse response;
try (InputStream payload = FactSearchManager.class.getClassLoader().getResourceAsStream(MAPPINGS_JSON);
InputStreamReader reader = new InputStreamReader(payload)) {
CreateIndexRequest request = new CreateIndexRequest(INDEX_NAME)
.source(CharStreams.toString(reader), XContentType.JSON);
response = clientFactory.getClient().indices().create(request, RequestOptions.DEFAULT);
} catch (ElasticsearchException | IOException ex) {
throw logAndExit(ex, "Could not perform request to create index.");
}
if (!response.isAcknowledged()) {
String msg = String.format("Could not create index '%s'.", INDEX_NAME);
LOGGER.error(msg);
throw new IllegalStateException(msg);
}
LOGGER.info("Successfully created index '%s'.", INDEX_NAME);
}
private ScrollingSearchResult.ScrollingBatch<FactDocument> fetchNextFactsBatch(String scrollId) {
SearchResponse response;
try {
SearchScrollRequest request = new SearchScrollRequest()
.scrollId(scrollId)
.scroll(searchScrollExpiration);
response = clientFactory.getClient().scroll(request, RequestOptions.DEFAULT);
} catch (ElasticsearchException | IOException ex) {
LOGGER.warning(ex, "Could not perform request to retrieve next batch of search results. Stop scrolling.");
return ScrollingSearchResult.emptyBatch();
}
if (response.status() != RestStatus.OK) {
LOGGER.warning("Could not retrieve next batch of search results (response code %s). Stop scrolling.", response.status());
return ScrollingSearchResult.emptyBatch();
}
return createFactsBatch(response);
}
private ScrollingSearchResult.ScrollingBatch<FactDocument> createFactsBatch(SearchResponse response) {
List<FactDocument> values = retrieveFactDocuments(response);
LOGGER.debug("Successfully retrieved next batch of search results (batch: %d, total: %d).", values.size(), response.getHits().getTotalHits());
boolean finished = values.size() < searchScrollSize;
if (finished) {
LOGGER.info("Successfully retrieved all search results. No more data available.");
// Close search context when all results have been fetched. If the client doesn't consume all results the context
// will be kept open until ElasticSearch cleans it up automatically after the expiration time elapsed.
closeSearchContext(response.getScrollId());
}
return new ScrollingSearchResult.ScrollingBatch<>(response.getScrollId(), values.iterator(), finished);
}
private void closeSearchContext(String scrollId) {
ClearScrollRequest request = new ClearScrollRequest();
request.addScrollId(scrollId);
// Perform this clean-up asynchronously because the client doesn't require the result.
clientFactory.getClient().clearScrollAsync(request, RequestOptions.DEFAULT, new ActionListener<ClearScrollResponse>() {
@Override
public void onResponse(ClearScrollResponse response) {
if (!response.isSucceeded()) {
LOGGER.warning("Could not close search context (response code %s).", response.status());
} else {
LOGGER.debug("Successfully closed search context.");
}
}
@Override
public void onFailure(Exception ex) {
LOGGER.warning(ex, "Could not close search context.");
}
});
}
private SearchRequest buildFactExistenceSearchRequest(FactExistenceSearchCriteria criteria) {
SearchSourceBuilder sourceBuilder = new SearchSourceBuilder()
.size(MAX_RESULT_WINDOW) // Always return all matching documents, but usually this should be zero or one.
.query(buildFactExistenceQuery(criteria));
return new SearchRequest()
.indices(INDEX_NAME)
.types(TYPE_NAME)
.source(sourceBuilder);
}
private SearchRequest buildFactsSearchRequest(FactSearchCriteria criteria) {
SearchSourceBuilder sourceBuilder = new SearchSourceBuilder()
.size(searchScrollSize)
.query(buildFactsQuery(criteria));
return new SearchRequest()
.indices(INDEX_NAME)
.types(TYPE_NAME)
.scroll(searchScrollExpiration)
.source(sourceBuilder);
}
private SearchRequest buildObjectsSearchRequest(FactSearchCriteria criteria) {
SearchSourceBuilder sourceBuilder = new SearchSourceBuilder()
.size(0) // Not interested in the search hits as the search result is part of the returned aggregations.
.aggregation(buildObjectsAggregation(criteria));
return new SearchRequest()
.indices(INDEX_NAME)
.types(TYPE_NAME)
.source(sourceBuilder);
}
private SearchRequest buildObjectStatisticsSearchRequest(ObjectStatisticsCriteria criteria) {
SearchSourceBuilder sourceBuilder = new SearchSourceBuilder()
.size(0) // Not interested in the search hits as the search result is part of the returned aggregations.
.aggregation(buildObjectStatisticsAggregation(criteria));
return new SearchRequest()
.indices(INDEX_NAME)
.types(TYPE_NAME)
.source(sourceBuilder);
}
private QueryBuilder buildFactExistenceQuery(FactExistenceSearchCriteria criteria) {
// Define all filters on direct Fact fields. Every field from the criteria must match.
BoolQueryBuilder rootQuery = boolQuery()
.filter(termQuery("typeID", toString(criteria.getFactTypeID())))
.filter(termQuery("sourceID", toString(criteria.getOriginID())))
.filter(termQuery("organizationID", toString(criteria.getOrganizationID())))
.filter(termQuery("accessMode", toString(criteria.getAccessMode())))
// Consider 'confidence' to be equal inside a given interval.
.filter(rangeQuery("confidence")
.gt(criteria.getConfidence() - CONFIDENCE_EQUALITY_INTERVAL / 2)
.lt(criteria.getConfidence() + CONFIDENCE_EQUALITY_INTERVAL / 2));
if (criteria.getFactValue() != null) {
// Fact values must match exactly if given.
rootQuery.filter(termQuery("value", criteria.getFactValue()));
} else {
// If 'value' isn't given make sure that it's also not set on any existing Fact.
rootQuery.mustNot(existsQuery("value"));
}
// This clause is required when searching for existing meta Facts.
if (criteria.getInReferenceTo() != null) {
rootQuery.filter(termQuery("inReferenceTo", toString(criteria.getInReferenceTo())));
}
// These clauses are required when searching for regular Facts.
if (!CollectionUtils.isEmpty(criteria.getObjects())) {
// The number of bound Objects must match (stored as a de-normalized field).
rootQuery.filter(termQuery("objectCount", criteria.getObjects().size()));
// Define filters on nested Objects. Also all Objects must match.
for (FactExistenceSearchCriteria.ObjectExistence object : criteria.getObjects()) {
BoolQueryBuilder objectsQuery = boolQuery()
.filter(termQuery("objects.id", toString(object.getObjectID())))
.filter(termQuery("objects.direction", toString(object.getDirection())));
rootQuery.filter(nestedQuery("objects", objectsQuery, ScoreMode.None));
}
}
return rootQuery;
}
private QueryBuilder buildFactsQuery(FactSearchCriteria criteria) {
BoolQueryBuilder rootQuery = boolQuery();
applySimpleFilterQueries(criteria, rootQuery);
applyKeywordSearchQuery(criteria, rootQuery);
applyTimestampSearchQuery(criteria, rootQuery);
applyNumberSearchQuery(criteria, rootQuery);
// Always apply access control query.
return rootQuery.filter(createAccessControlQuery(criteria.getCurrentUserID(), criteria.getAvailableOrganizationID()));
}
private void applySimpleFilterQueries(FactSearchCriteria criteria, BoolQueryBuilder rootQuery) {
if (!CollectionUtils.isEmpty(criteria.getFactID())) {
rootQuery.filter(termsQuery("_id", toString(criteria.getFactID())));
}
if (!CollectionUtils.isEmpty(criteria.getFactTypeID())) {
rootQuery.filter(termsQuery("typeID", toString(criteria.getFactTypeID())));
}
if (!CollectionUtils.isEmpty(criteria.getFactValue())) {
rootQuery.filter(termsQuery("value", criteria.getFactValue()));
}
if (!CollectionUtils.isEmpty(criteria.getInReferenceTo())) {
rootQuery.filter(termsQuery("inReferenceTo", toString(criteria.getInReferenceTo())));
}
if (!CollectionUtils.isEmpty(criteria.getOrganizationID())) {
rootQuery.filter(termsQuery("organizationID", toString(criteria.getOrganizationID())));
}
if (!CollectionUtils.isEmpty(criteria.getOriginID())) {
rootQuery.filter(termsQuery("sourceID", toString(criteria.getOriginID())));
}
if (!CollectionUtils.isEmpty(criteria.getObjectID())) {
rootQuery.filter(nestedQuery("objects", termsQuery("objects.id", toString(criteria.getObjectID())), ScoreMode.None));
}
if (!CollectionUtils.isEmpty(criteria.getObjectTypeID())) {
rootQuery.filter(nestedQuery("objects", termsQuery("objects.typeID", toString(criteria.getObjectTypeID())), ScoreMode.None));
}
if (!CollectionUtils.isEmpty(criteria.getObjectValue())) {
rootQuery.filter(nestedQuery("objects", termsQuery("objects.value", criteria.getObjectValue()), ScoreMode.None));
}
if (criteria.getFactBinding() != null) {
rootQuery.filter(termQuery("objectCount", criteria.getFactBinding().getObjectCount()));
}
}
private void applyKeywordSearchQuery(FactSearchCriteria criteria, BoolQueryBuilder rootQuery) {
if (StringUtils.isBlank(criteria.getKeywords())) return;
applyFieldStrategy(rootQuery, field -> createFieldQuery(field, criteria.getKeywords()),
criteria.getKeywordFieldStrategy(), criteria.getKeywordMatchStrategy());
}
private void applyTimestampSearchQuery(FactSearchCriteria criteria, BoolQueryBuilder rootQuery) {
if (criteria.getStartTimestamp() == null && criteria.getEndTimestamp() == null) return;
applyFieldStrategy(rootQuery, field -> createFieldQuery(field, criteria.getStartTimestamp(), criteria.getEndTimestamp()),
criteria.getTimeFieldStrategy(), criteria.getTimeMatchStrategy());
}
private void applyNumberSearchQuery(FactSearchCriteria criteria, BoolQueryBuilder rootQuery) {
if (criteria.getMinNumber() == null && criteria.getMaxNumber() == null) return;
applyFieldStrategy(rootQuery, field -> rangeQuery(field).gte(criteria.getMinNumber()).lte(criteria.getMaxNumber()),
criteria.getNumberFieldStrategy(), criteria.getNumberMatchStrategy());
}
private void applyFieldStrategy(BoolQueryBuilder rootQuery, Function<String, QueryBuilder> fieldQueryResolver,
Set<? extends FactSearchCriteria.FieldStrategy> fieldStrategies,
FactSearchCriteria.MatchStrategy matchStrategy) {
// Determine all fields to query.
Set<String> fieldsToQuery = fieldStrategies.stream()
.flatMap(strategy -> strategy.getFields().stream())
.collect(Collectors.toSet());
BoolQueryBuilder strategyQuery = boolQuery();
for (String field : fieldsToQuery) {
if (matchStrategy == FactSearchCriteria.MatchStrategy.all) {
// Field query must match all fields.
strategyQuery.filter(fieldQueryResolver.apply(field));
} else {
// Field query should match at least one field.
strategyQuery.should(fieldQueryResolver.apply(field));
}
}
rootQuery.filter(strategyQuery);
}
private QueryBuilder createFieldQuery(String field, String keywords) {
SimpleQueryStringBuilder query = simpleQueryStringQuery(keywords)
.field(field)
// Values are indexed differently. Avoid errors when executing an IP search against a text field, for example.
.lenient(true);
// If field starts with the prefix 'objects.' it's part of the nested objects, thus, it must be wrapped inside a nested query.
return field.startsWith("objects.") ? nestedQuery("objects", query, ScoreMode.Avg) : query;
}
private QueryBuilder createFieldQuery(String field, Long startTimestamp, Long endTimestamp) {
// Negative timestamps are omitted by providing NULL to from() and to().
return rangeQuery(field)
.from(startTimestamp != null && startTimestamp > 0 ? startTimestamp : null)
.to(endTimestamp != null && endTimestamp > 0 ? endTimestamp : null);
}
private QueryBuilder createAccessControlQuery(UUID currentUserID, Set<UUID> availableOrganizationID) {
// Query to verify that user has access to Fact ...
return boolQuery()
// ... if Fact is public.
.should(termQuery("accessMode", toString(FactDocument.AccessMode.Public)))
// ... if AccessMode == Explicit user must be in ACL.
.should(boolQuery()
.filter(termQuery("accessMode", toString(FactDocument.AccessMode.Explicit)))
.filter(termQuery("acl", toString(currentUserID)))
)
// ... if AccessMode == RoleBased user must be in ACL or have access to the owning Organization.
.should(boolQuery()
.filter(termQuery("accessMode", toString(FactDocument.AccessMode.RoleBased)))
.filter(boolQuery()
.should(termQuery("acl", toString(currentUserID)))
.should(termsQuery("organizationID", toString(availableOrganizationID)))
)
);
}
private AggregationBuilder buildObjectsAggregation(FactSearchCriteria criteria) {
// 1. Reduce to Facts matching the search criteria.
return filter(FILTER_FACTS_AGGREGATION_NAME, buildFactsQuery(criteria))
// 2. Map to nested Object documents.
.subAggregation(nested(NESTED_OBJECTS_AGGREGATION_NAME, "objects")
// 3. Reduce to Objects matching the search criteria.
.subAggregation(filter(FILTER_OBJECTS_AGGREGATION_NAME, buildObjectsQuery(criteria))
// 4. Calculate the number of unique Objects by id. This will give the 'count' value.
// If 'count' is smaller than MAX_RESULT_WINDOW a correct value is expected, thus,
// the precision threshold is set to MAX_RESULT_WINDOW.
.subAggregation(cardinality(OBJECTS_COUNT_AGGREGATION_NAME)
.field("objects.id")
.precisionThreshold(MAX_RESULT_WINDOW)
)
// 5. Reduce to buckets of unique Objects by id, restricted to the search criteria's limit.
// This will give the actual search results.
.subAggregation(terms(UNIQUE_OBJECTS_AGGREGATION_NAME)
.field("objects.id")
.size(calculateMaximumSize(criteria))
// 6. Map to the unique Object's source. Set size to 1, because all Objects in one
// bucket are the same (ignoring 'direction' which isn't relevant for Object search).
.subAggregation(topHits(UNIQUE_OBJECTS_SOURCE_AGGREGATION_NAME)
.size(1)
)
)
)
);
}
private QueryBuilder buildObjectsQuery(FactSearchCriteria criteria) {
BoolQueryBuilder rootQuery = boolQuery();
// Apply all simple filter queries on Objects. It's not necessary to wrap them inside a nested query because the
// query is executed inside a nested aggregation which has direct access to the nested documents.
if (!CollectionUtils.isEmpty(criteria.getObjectID())) {
rootQuery.filter(termsQuery("objects.id", toString(criteria.getObjectID())));
}
if (!CollectionUtils.isEmpty(criteria.getObjectTypeID())) {
rootQuery.filter(termsQuery("objects.typeID", toString(criteria.getObjectTypeID())));
}
if (!CollectionUtils.isEmpty(criteria.getObjectValue())) {
rootQuery.filter(termsQuery("objects.value", criteria.getObjectValue()));
}
// Apply keyword search on Object values if necessary.
Set<FactSearchCriteria.KeywordFieldStrategy> objectFieldStrategy = onlyObjectFieldStrategy(criteria);
if (!StringUtils.isBlank(criteria.getKeywords()) && !CollectionUtils.isEmpty(objectFieldStrategy)) {
// Values are indexed differently. Avoid errors by setting 'lenient' to true.
applyFieldStrategy(rootQuery, field -> simpleQueryStringQuery(criteria.getKeywords()).field(field).lenient(true),
objectFieldStrategy, criteria.getKeywordMatchStrategy());
}
return rootQuery;
}
private Set<FactSearchCriteria.KeywordFieldStrategy> onlyObjectFieldStrategy(FactSearchCriteria criteria) {
Set<FactSearchCriteria.KeywordFieldStrategy> strategy = SetUtils.set();
// When keyword search is performed directly on Objects only consider the relevant field strategies.
if (criteria.getKeywordFieldStrategy().contains(FactSearchCriteria.KeywordFieldStrategy.all) ||
criteria.getKeywordFieldStrategy().contains(FactSearchCriteria.KeywordFieldStrategy.objectValueText)) {
strategy.add(FactSearchCriteria.KeywordFieldStrategy.objectValueText);
}
if (criteria.getKeywordFieldStrategy().contains(FactSearchCriteria.KeywordFieldStrategy.all) ||
criteria.getKeywordFieldStrategy().contains(FactSearchCriteria.KeywordFieldStrategy.objectValueIp)) {
strategy.add(FactSearchCriteria.KeywordFieldStrategy.objectValueIp);
}
if (criteria.getKeywordFieldStrategy().contains(FactSearchCriteria.KeywordFieldStrategy.all) ||
criteria.getKeywordFieldStrategy().contains(FactSearchCriteria.KeywordFieldStrategy.objectValueDomain)) {
strategy.add(FactSearchCriteria.KeywordFieldStrategy.objectValueDomain);
}
return strategy;
}
private int calculateMaximumSize(FactSearchCriteria criteria) {
return criteria.getLimit() > 0 && criteria.getLimit() < MAX_RESULT_WINDOW ? criteria.getLimit() : MAX_RESULT_WINDOW;
}
private AggregationBuilder buildObjectStatisticsAggregation(ObjectStatisticsCriteria criteria) {
QueryBuilder accessControlQuery = createAccessControlQuery(criteria.getCurrentUserID(), criteria.getAvailableOrganizationID());
QueryBuilder objectsQuery = termsQuery("objects.id", toString(criteria.getObjectID()));
// 1. Reduce to only the Facts the user has access to. Non-accessible Facts won't be available in sub aggregations!
return filter(FILTER_FACTS_AGGREGATION_NAME, accessControlQuery)
// 2. Map to nested Object documents.
.subAggregation(nested(NESTED_OBJECTS_AGGREGATION_NAME, "objects")
// 3. Reduce to only the Objects for which statistics should be calculated.
.subAggregation(filter(FILTER_OBJECTS_AGGREGATION_NAME, objectsQuery)
// 4. Reduce to buckets of unique Objects by id. There shouldn't be more buckets than the
// number of Objects for which statistics will be calculated ('size' parameter).
.subAggregation(terms(UNIQUE_OBJECTS_AGGREGATION_NAME)
.field("objects.id")
.size(criteria.getObjectID().size())
// 5. Reverse nested aggregation to have access to parent Facts.
.subAggregation(reverseNested(REVERSED_FACTS_AGGREGATION_NAME)
// 6. Create one bucket for each FactType. Set 'size' to MAX_RESULT_WINDOW
// in order to get the statistics for all FactTypes. The returned 'doc_count'
// will give the number of Facts per FactType.
.subAggregation(terms(UNIQUE_FACT_TYPES_AGGREGATION_NAME)
.field("typeID")
.size(MAX_RESULT_WINDOW)
// 7. Calculate the maximum lastAddedTimestamp per FactType.
.subAggregation(max(MAX_LAST_ADDED_TIMESTAMP_AGGREGATION_NAME)
.field("timestamp")
)
// 8. Calculate the maximum lastSeenTimestamp per FactType.
.subAggregation(max(MAX_LAST_SEEN_TIMESTAMP_AGGREGATION_NAME)
.field("lastSeenTimestamp")
)
)
)
)
)
);
}
private List<FactDocument> retrieveFactDocuments(SearchResponse response) {
List<FactDocument> result = ListUtils.list();
for (SearchHit hit : response.getHits()) {
FactDocument document = decodeFactDocument(UUID.fromString(hit.getId()), toBytes(hit.getSourceRef()));
if (document != null) {
result.add(document);
}
}
return result;
}
private int retrieveSearchObjectsResultCount(SearchResponse response) {
Aggregation objectsCountAggregation = resolveChildAggregation(response.getAggregations(), OBJECTS_COUNT_AGGREGATION_NAME);
if (!(objectsCountAggregation instanceof Cardinality)) {
LOGGER.warning("Could not retrieve result count when searching for Objects.");
return -1;
}
// Retrieve Object count from the cardinality aggregation.
return (int) Cardinality.class.cast(objectsCountAggregation).getValue();
}
private List<ObjectDocument> retrieveSearchObjectsResultValues(SearchResponse response) {
List<ObjectDocument> result = ListUtils.list();
Aggregation uniqueObjectsAggregation = resolveChildAggregation(response.getAggregations(), UNIQUE_OBJECTS_AGGREGATION_NAME);
if (!(uniqueObjectsAggregation instanceof Terms)) {
LOGGER.warning("Could not retrieve result values when searching for Objects.");
return result;
}
List<? extends Terms.Bucket> buckets = Terms.class.cast(uniqueObjectsAggregation).getBuckets();
if (CollectionUtils.isEmpty(buckets)) {
// No buckets means no results.
return result;
}
// Each bucket contains one unique Object, where the TopHits aggregation provides the document source.
for (Terms.Bucket bucket : buckets) {
Aggregation uniqueObjectsSourceAggregation = bucket.getAggregations().get(UNIQUE_OBJECTS_SOURCE_AGGREGATION_NAME);
if (!(uniqueObjectsSourceAggregation instanceof TopHits)) continue;
// Each bucket should contain only one hit with one unique Object.
SearchHits hits = TopHits.class.cast(uniqueObjectsSourceAggregation).getHits();
if (hits.getHits().length < 1) continue;
// Retrieve Object document from provided search hit.
ObjectDocument document = decodeObjectDocument(toBytes(hits.getAt(0).getSourceRef()));
if (document != null) {
result.add(document);
}
}
return result;
}
private ObjectStatisticsContainer retrieveObjectStatisticsResult(SearchResponse response) {
Aggregation uniqueObjectsAggregation = resolveChildAggregation(response.getAggregations(), UNIQUE_OBJECTS_AGGREGATION_NAME);
if (!(uniqueObjectsAggregation instanceof Terms)) {
LOGGER.warning("Could not retrieve results when calculating statistics for Objects.");
return ObjectStatisticsContainer.builder().build();
}
List<? extends Terms.Bucket> uniqueObjectBuckets = Terms.class.cast(uniqueObjectsAggregation).getBuckets();
if (CollectionUtils.isEmpty(uniqueObjectBuckets)) {
// No buckets means no results.
return ObjectStatisticsContainer.builder().build();
}
ObjectStatisticsContainer.Builder resultBuilder = ObjectStatisticsContainer.builder();
// Each bucket contains one unique Object. Calculate the statistics for each Object.
for (Terms.Bucket objectBucket : uniqueObjectBuckets) {
UUID objectID = UUID.fromString(objectBucket.getKeyAsString());
// Resolve buckets of unique FactTypes ...
Aggregation uniqueFactTypesAggregation = resolveChildAggregation(objectBucket.getAggregations(), UNIQUE_FACT_TYPES_AGGREGATION_NAME);
if (!(uniqueFactTypesAggregation instanceof Terms)) continue;
List<? extends Terms.Bucket> uniqueFactTypeBuckets = Terms.class.cast(uniqueFactTypesAggregation).getBuckets();
if (CollectionUtils.isEmpty(uniqueFactTypeBuckets)) continue;
// ... and add the statistics for each FactType to the result.
for (Terms.Bucket factTypeBucket : uniqueFactTypeBuckets) {
UUID factTypeID = UUID.fromString(factTypeBucket.getKeyAsString());
int factCount = (int) factTypeBucket.getDocCount();
long lastAddedTimestamp = retrieveMaxTimestamp(factTypeBucket, MAX_LAST_ADDED_TIMESTAMP_AGGREGATION_NAME);
long lastSeenTimestamp = retrieveMaxTimestamp(factTypeBucket, MAX_LAST_SEEN_TIMESTAMP_AGGREGATION_NAME);
resultBuilder.addStatistic(objectID, new ObjectStatisticsContainer.FactStatistic(factTypeID, factCount, lastAddedTimestamp, lastSeenTimestamp));
}
}
return resultBuilder.build();
}
private long retrieveMaxTimestamp(Terms.Bucket bucket, String targetAggregationName) {
Aggregation maxAggregation = bucket.getAggregations().get(targetAggregationName);
if (!(maxAggregation instanceof Max)) {
LOGGER.warning("Could not retrieve maximum timestamp when calculating statistics for Objects.");
return -1;
}
// Retrieve maximum timestamp from the max aggregation.
return Math.round(Max.class.cast(maxAggregation).getValue());
}
private Aggregation resolveChildAggregation(Aggregations aggregations, String targetAggregationName) {
if (aggregations == null) return null;
for (Aggregation aggregation : aggregations) {
// Check if 'aggregation' is already the target aggregation.
if (aggregation.getName().equals(targetAggregationName)) {
return aggregation;
}
// Otherwise check all sub aggregations if applicable.
if (HasAggregations.class.isAssignableFrom(aggregation.getClass())) {
Aggregation target = resolveChildAggregation(HasAggregations.class.cast(aggregation).getAggregations(), targetAggregationName);
if (target != null) return target;
}
}
// Couldn't find target aggregation.
return null;
}
private FactDocument decodeFactDocument(UUID factID, byte[] source) {
try {
FactDocument fact = FACT_DOCUMENT_READER.readValue(source);
// Need to set ID manually because it's not indexed as an own field.
return fact.setId(factID);
} catch (IOException ex) {
LOGGER.warning(ex, "Could not deserialize Fact with id = %s. Source document not stored?", factID);
return null;
}
}
private ObjectDocument decodeObjectDocument(byte[] source) {
try {
return OBJECT_DOCUMENT_READER.readValue(source);
} catch (IOException ex) {
LOGGER.warning(ex, "Could not deserialize Object. Source of nested document not returned?");
return null;
}
}
private String toString(Object object) {
return ObjectUtils.ifNotNull(object, Object::toString);
}
private Set<String> toString(Set<?> collection) {
return SetUtils.set(collection, Object::toString);
}
private RuntimeException logAndExit(Exception ex, String msg) {
LOGGER.error(ex, msg);
return new IllegalStateException(msg, ex);
}
}
