/**
* This file is hereby placed into the Public Domain. This means anyone is
* free to do whatever they wish with this file.
*/
package mil.nga.giat.data.elasticsearch;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.logging.Level;
import java.util.logging.Logger;
import org.geotools.data.FeatureReader;
import org.geotools.data.FilteringFeatureReader;
import org.geotools.data.Query;
import org.geotools.data.store.ContentEntry;
import org.geotools.data.store.ContentFeatureSource;
import org.geotools.filter.visitor.ExtractBoundsFilterVisitor;
import org.geotools.geometry.jts.ReferencedEnvelope;
import org.geotools.util.logging.Logging;
import org.opengis.feature.simple.SimpleFeature;
import org.opengis.feature.simple.SimpleFeatureType;
import org.opengis.filter.sort.SortBy;
import org.opengis.filter.sort.SortOrder;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import com.vividsolutions.jts.geom.Envelope;
/**
* Provides access to a specific type within the Elasticsearch index described
* by the associated data store.
*
*/
public class ElasticFeatureSource extends ContentFeatureSource {
private final static Logger LOGGER = Logging.getLogger(ElasticFeatureSource.class);
private Boolean filterFullySupported;
public ElasticFeatureSource(ContentEntry entry, Query query) throws IOException {
super(entry, query);
final ElasticDataStore dataStore = getDataStore();
if (dataStore.getLayerConfigurations().get(entry.getName().getLocalPart()) == null) {
final List<ElasticAttribute> attributes = dataStore.getElasticAttributes(entry.getName());
final ElasticLayerConfiguration config = new ElasticLayerConfiguration(entry.getName().getLocalPart());
config.getAttributes().addAll(attributes);
dataStore.setLayerConfiguration(config);
}
}
/**
* Access parent datastore
*/
public ElasticDataStore getDataStore() {
return (ElasticDataStore) super.getDataStore();
}
/**
* Implementation that generates the total bounds
*/
@Override
protected ReferencedEnvelope getBoundsInternal(Query query) throws IOException {
LOGGER.fine("getBoundsInternal");
final CoordinateReferenceSystem crs = getSchema().getCoordinateReferenceSystem();
final ReferencedEnvelope bounds = new ReferencedEnvelope(crs);
try (FeatureReader<SimpleFeatureType, SimpleFeature> featureReader = getReaderInternal(query)) {
while (featureReader.hasNext()) {
final SimpleFeature feature = featureReader.next();
bounds.include(feature.getBounds());
}
}
return bounds;
}
@Override
protected int getCountInternal(Query query) throws IOException {
LOGGER.fine("getCountInternal");
int hits = 0;
final ElasticRequest searchRequest = prepareSearchRequest(query, false);
try {
if (!filterFullySupported) {
try (FeatureReader<SimpleFeatureType, SimpleFeature> reader = getReaderInternal(query)) {
while (reader.hasNext()) {
reader.next();
hits++;
}
}
} else {
searchRequest.setSize(0);
final ElasticDataStore dataStore = getDataStore();
final String docType = dataStore.getDocType(entry.getName());
final ElasticResponse sr = dataStore.getClient().search(dataStore.getSearchIndices(), docType, searchRequest);
final int totalHits = (int) sr.getTotalNumHits();
final int size = getSize(query);
final int from = getStartIndex(query);
hits = Math.max(0, Math.min(totalHits - from, size));
}
} catch (Exception e) {
LOGGER.log(Level.SEVERE, e.getMessage(), e);
throw new IOException("Error executing count search", e);
}
return hits;
}
@Override
protected FeatureReader<SimpleFeatureType, SimpleFeature> getReaderInternal(Query query) throws IOException {
LOGGER.fine("getReaderInternal");
FeatureReader<SimpleFeatureType, SimpleFeature> reader = null;
try {
final ElasticDataStore dataStore = getDataStore();
final String docType = dataStore.getDocType(entry.getName());
final boolean scroll = !useSortOrPagination(query) && dataStore.getScrollEnabled();
final ElasticRequest searchRequest = prepareSearchRequest(query, scroll);
final ElasticResponse sr = dataStore.getClient().search(dataStore.getSearchIndices(), docType, searchRequest);
if (LOGGER.isLoggable(Level.FINE)) {
LOGGER.fine("Search response: " + sr);
}
if (!scroll) {
reader = new ElasticFeatureReader(getState(), sr);
} else {
reader = new ElasticFeatureReaderScroll(getState(), sr, getSize(query));
}
if (!filterFullySupported) {
reader = new FilteringFeatureReader<SimpleFeatureType, SimpleFeature>(reader, query.getFilter());
}
} catch (Exception e) {
LOGGER.log(Level.SEVERE, e.getMessage(), e);
throw new IOException("Error executing query search", e);
}
return reader;
}
private ElasticRequest prepareSearchRequest(Query query, boolean scroll) throws IOException {
String naturalSortOrder = SortOrder.ASCENDING.toSQL().toLowerCase();
final ElasticRequest searchRequest = new ElasticRequest();
final ElasticDataStore dataStore = getDataStore();
final String docType = dataStore.getDocType(entry.getName());
LOGGER.fine("Preparing " + docType + " (" + entry.getName() + ") query");
if (!scroll) {
if (query.getSortBy()!=null){
for (final SortBy sort : query.getSortBy()) {
final String sortOrder = sort.getSortOrder().toSQL().toLowerCase();
if (sort.getPropertyName() != null) {
final String name = sort.getPropertyName().getPropertyName();
searchRequest.addSort(name, sortOrder);
} else {
naturalSortOrder = sortOrder;
}
}
}
// pagination
searchRequest.setSize(getSize(query));
searchRequest.setFrom(getStartIndex(query));
} else {
if (dataStore.getScrollSize() != null) {
searchRequest.setSize(dataStore.getScrollSize().intValue());
}
if (dataStore.getScrollTime() != null) {
searchRequest.setScroll(dataStore.getScrollTime());
}
}
if (dataStore.isSourceFilteringEnabled()) {
// add source includes
setSourceIncludes(searchRequest);
}
// add query and post filter
final FilterToElastic filterToElastic = new FilterToElastic();
filterToElastic.setFeatureType(buildFeatureType());
filterToElastic.encode(query);
filterFullySupported = filterToElastic.getFullySupported();
if (!filterFullySupported) {
LOGGER.fine("Filter is not fully supported by native Elasticsearch."
+ " Additional post-query filtering will be performed.");
}
final Map<String,Object> queryBuilder = filterToElastic.getQueryBuilder();
final Map<String,Object> nativeQueryBuilder = filterToElastic.getNativeQueryBuilder();
searchRequest.setQuery(queryBuilder);
if (isSort(query) && nativeQueryBuilder.equals(ElasticConstants.MATCH_ALL)) {
searchRequest.addSort("_uid", naturalSortOrder);
}
if (filterToElastic.getAggregations() != null) {
final Map<String, Map<String, Map<String, Object>>> aggregations = filterToElastic.getAggregations();
final Envelope envelope = (Envelope) query.getFilter().accept(ExtractBoundsFilterVisitor.BOUNDS_VISITOR, null);
final long gridSize;
if (dataStore.getGridSize() != null) {
gridSize = dataStore.getGridSize();
} else {
gridSize = (Long) ElasticDataStoreFactory.GRID_SIZE.getDefaultValue();
}
final double gridThreshold;
if (dataStore.getGridThreshold() != null) {
gridThreshold = dataStore.getGridThreshold();
} else {
gridThreshold = (Double) ElasticDataStoreFactory.GRID_THRESHOLD.getDefaultValue();
}
final int precision = GeohashUtil.computePrecision(envelope, gridSize, gridThreshold);
LOGGER.fine("Updating GeoHash grid aggregation precision to " + precision);
GeohashUtil.updateGridAggregationPrecision(aggregations, precision);
searchRequest.setAggregations(aggregations);
searchRequest.setSize(0);
}
return searchRequest;
}
private void setSourceIncludes(final ElasticRequest searchRequest) throws IOException {
final ElasticDataStore dataStore = getDataStore();
final List<ElasticAttribute> attributes = dataStore.getElasticAttributes(entry.getName());
for (final ElasticAttribute attribute : attributes) {
if (attribute.isUse() && attribute.isStored()) {
searchRequest.addField(attribute.getName());
} else if (attribute.isUse()) {
searchRequest.addSourceInclude(attribute.getName());
}
}
}
private boolean isSort(Query query) {
return query.getSortBy() != null && query.getSortBy().length > 0;
}
private boolean useSortOrPagination(Query query) {
return (query.getSortBy() != null && query.getSortBy().length > 0) ||
query.getStartIndex()!=null;
}
private int getSize(Query query) {
final int size;
if (!query.isMaxFeaturesUnlimited()) {
size = query.getMaxFeatures();
} else {
size = getDataStore().getDefaultMaxFeatures();
LOGGER.fine("Unlimited maxFeatures not supported. Using default: " + size);
}
return size;
}
private int getStartIndex(Query query) {
final int from;
if (query.getStartIndex() != null) {
from = query.getStartIndex();
} else {
from = 0;
}
return from;
}
@Override
protected SimpleFeatureType buildFeatureType() throws IOException {
final ElasticDataStore ds = getDataStore();
final ElasticLayerConfiguration layerConfig;
layerConfig = ds.getLayerConfigurations().get(entry.getTypeName());
final List<ElasticAttribute> attributes;
if (layerConfig != null) {
attributes = layerConfig.getAttributes();
} else {
attributes = null;
}
final ElasticFeatureTypeBuilder typeBuilder;
typeBuilder = new ElasticFeatureTypeBuilder(attributes, entry.getName());
return typeBuilder.buildFeatureType();
}
@Override
protected boolean canLimit() {
return true;
}
@Override
protected boolean canOffset() {
return true;
}
@Override
protected boolean canFilter() {
return true;
}
@Override
protected boolean canSort() {
return true;
}
}