ATLAS-1818: basic-search enhancements to improve search performance
Signed-off-by: Madhan Neethiraj <madhan@apache.org>
diff --git a/distro/src/conf/atlas-application.properties b/distro/src/conf/atlas-application.properties
index b2b8e74..5e59528 100755
--- a/distro/src/conf/atlas-application.properties
+++ b/distro/src/conf/atlas-application.properties
@@ -62,6 +62,8 @@
${titan.index.properties}
+# Solr-specific configuration property
+atlas.graph.index.search.max-result-set-size=150
######### Notification Configs #########
atlas.notification.embedded=true
diff --git a/distro/src/conf/solr/solrconfig.xml b/distro/src/conf/solr/solrconfig.xml
index ce2e20b..1d414f7 100644
--- a/distro/src/conf/solr/solrconfig.xml
+++ b/distro/src/conf/solr/solrconfig.xml
@@ -277,19 +277,19 @@
and old cache.
-->
<filterCache class="solr.FastLRUCache"
- size="512"
- initialSize="512"
- autowarmCount="0"/>
+ size="2000"
+ initialSize="2000"
+ autowarmCount="1000"/>
<!-- Query Result Cache
Caches results of searches - ordered lists of document ids
(DocList) based on a query, a sort, and the range of documents requested.
-->
- <queryResultCache class="solr.LRUCache"
- size="512"
- initialSize="512"
- autowarmCount="0"/>
+ <queryResultCache class="solr.FastLRUCache"
+ size="26000"
+ initialSize="26000"
+ autowarmCount="400"/>
<!-- Document Cache
@@ -297,10 +297,10 @@
document). Since Lucene internal document ids are transient,
this cache will not be autowarmed.
-->
- <documentCache class="solr.LRUCache"
- size="512"
- initialSize="512"
- autowarmCount="0"/>
+ <documentCache class="solr.FastLRUCache"
+ size="26000"
+ initialSize="26000"
+ autowarmCount="400"/>
<!-- custom cache currently used by block join -->
<cache name="perSegFilter"
diff --git a/graphdb/api/src/main/java/org/apache/atlas/repository/graphdb/AtlasGraph.java b/graphdb/api/src/main/java/org/apache/atlas/repository/graphdb/AtlasGraph.java
index a3a27bf..dded76f 100644
--- a/graphdb/api/src/main/java/org/apache/atlas/repository/graphdb/AtlasGraph.java
+++ b/graphdb/api/src/main/java/org/apache/atlas/repository/graphdb/AtlasGraph.java
@@ -158,6 +158,20 @@
AtlasIndexQuery<V, E> indexQuery(String indexName, String queryString);
/**
+ * Creates an index query.
+ *
+ * @param indexName index name
+ * @param queryString the query
+ * @param offset specify the offset that should be applied for the query. This is useful for paging through
+ * list of results
+ *
+ * @see <a
+ * href="https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl.html">
+ * Elastic Search Reference</a> for query syntax
+ */
+ AtlasIndexQuery<V, E> indexQuery(String indexName, String queryString, int offset);
+
+ /**
* Gets the management object associated with this graph and opens a transaction
* for changes that are made.
* @return
diff --git a/graphdb/titan0/src/main/java/org/apache/atlas/repository/graphdb/titan0/Titan0Graph.java b/graphdb/titan0/src/main/java/org/apache/atlas/repository/graphdb/titan0/Titan0Graph.java
index 9624c99..2408287 100644
--- a/graphdb/titan0/src/main/java/org/apache/atlas/repository/graphdb/titan0/Titan0Graph.java
+++ b/graphdb/titan0/src/main/java/org/apache/atlas/repository/graphdb/titan0/Titan0Graph.java
@@ -17,39 +17,6 @@
*/
package org.apache.atlas.repository.graphdb.titan0;
-import java.io.IOException;
-import java.io.OutputStream;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-
-import javax.script.Bindings;
-import javax.script.ScriptContext;
-import javax.script.ScriptEngine;
-import javax.script.ScriptEngineManager;
-import javax.script.ScriptException;
-
-import org.apache.atlas.AtlasErrorCode;
-import org.apache.atlas.exception.AtlasBaseException;
-import org.apache.atlas.groovy.GroovyExpression;
-import org.apache.atlas.repository.graphdb.AtlasEdge;
-import org.apache.atlas.repository.graphdb.AtlasGraph;
-import org.apache.atlas.repository.graphdb.AtlasGraphManagement;
-import org.apache.atlas.repository.graphdb.AtlasGraphQuery;
-import org.apache.atlas.repository.graphdb.AtlasIndexQuery;
-import org.apache.atlas.repository.graphdb.AtlasSchemaViolationException;
-import org.apache.atlas.repository.graphdb.AtlasVertex;
-import org.apache.atlas.repository.graphdb.GremlinVersion;
-import org.apache.atlas.repository.graphdb.titan0.query.Titan0GraphQuery;
-import org.apache.atlas.repository.graphdb.utils.IteratorToIterableAdapter;
-import org.apache.atlas.typesystem.types.IDataType;
-
import com.google.common.base.Function;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
@@ -65,9 +32,40 @@
import com.tinkerpop.blueprints.Vertex;
import com.tinkerpop.blueprints.util.io.graphson.GraphSONWriter;
import com.tinkerpop.pipes.util.structures.Row;
+import org.apache.atlas.AtlasErrorCode;
+import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.groovy.GroovyExpression;
+import org.apache.atlas.repository.graphdb.AtlasEdge;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
+import org.apache.atlas.repository.graphdb.AtlasGraphManagement;
+import org.apache.atlas.repository.graphdb.AtlasGraphQuery;
+import org.apache.atlas.repository.graphdb.AtlasIndexQuery;
+import org.apache.atlas.repository.graphdb.AtlasSchemaViolationException;
+import org.apache.atlas.repository.graphdb.AtlasVertex;
+import org.apache.atlas.repository.graphdb.GremlinVersion;
+import org.apache.atlas.repository.graphdb.titan0.query.Titan0GraphQuery;
+import org.apache.atlas.repository.graphdb.utils.IteratorToIterableAdapter;
+import org.apache.atlas.typesystem.types.IDataType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import javax.script.Bindings;
+import javax.script.ScriptContext;
+import javax.script.ScriptEngine;
+import javax.script.ScriptEngineManager;
+import javax.script.ScriptException;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
/**
* Titan 0.5.4 implementation of AtlasGraph.
@@ -163,7 +161,12 @@
@Override
public AtlasIndexQuery<Titan0Vertex, Titan0Edge> indexQuery(String fulltextIndex, String graphQuery) {
- TitanIndexQuery query = getGraph().indexQuery(fulltextIndex, graphQuery);
+ return indexQuery(fulltextIndex, graphQuery, 0);
+ }
+
+ @Override
+ public AtlasIndexQuery<Titan0Vertex, Titan0Edge> indexQuery(String fulltextIndex, String graphQuery, int offset) {
+ TitanIndexQuery query = getGraph().indexQuery(fulltextIndex, graphQuery).offset(offset);
return new Titan0IndexQuery(this, query);
}
diff --git a/graphdb/titan0/src/test/resources/atlas-application.properties b/graphdb/titan0/src/test/resources/atlas-application.properties
index 636a9ff..3058330 100644
--- a/graphdb/titan0/src/test/resources/atlas-application.properties
+++ b/graphdb/titan0/src/test/resources/atlas-application.properties
@@ -50,6 +50,8 @@
atlas.graph.index.search.solr.mode=cloud
atlas.graph.index.search.solr.zookeeper-url=${solr.zk.address}
+# Solr-specific configuration property
+atlas.graph.index.search.max-result-set-size=150
######### Hive Lineage Configs #########
# This models reflects the base super types for Data and Process
diff --git a/graphdb/titan1/src/main/java/org/apache/atlas/repository/graphdb/titan1/Titan1Graph.java b/graphdb/titan1/src/main/java/org/apache/atlas/repository/graphdb/titan1/Titan1Graph.java
index 6a61075..e829d91 100644
--- a/graphdb/titan1/src/main/java/org/apache/atlas/repository/graphdb/titan1/Titan1Graph.java
+++ b/graphdb/titan1/src/main/java/org/apache/atlas/repository/graphdb/titan1/Titan1Graph.java
@@ -17,19 +17,18 @@
*/
package org.apache.atlas.repository.graphdb.titan1;
-import java.io.IOException;
-import java.io.OutputStream;
-import java.util.Collection;
-import java.util.HashSet;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-
-import javax.script.Bindings;
-import javax.script.ScriptEngine;
-import javax.script.ScriptException;
-
+import com.google.common.base.Function;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
+import com.thinkaurelius.titan.core.Cardinality;
+import com.thinkaurelius.titan.core.PropertyKey;
+import com.thinkaurelius.titan.core.SchemaViolationException;
+import com.thinkaurelius.titan.core.TitanGraph;
+import com.thinkaurelius.titan.core.TitanIndexQuery;
+import com.thinkaurelius.titan.core.schema.TitanGraphIndex;
+import com.thinkaurelius.titan.core.schema.TitanManagement;
+import com.thinkaurelius.titan.core.util.TitanCleanup;
import org.apache.atlas.AtlasErrorCode;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.groovy.GroovyExpression;
@@ -57,18 +56,17 @@
import org.apache.tinkerpop.gremlin.structure.io.graphson.GraphSONMapper;
import org.apache.tinkerpop.gremlin.structure.io.graphson.GraphSONWriter;
-import com.google.common.base.Function;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Lists;
-import com.google.common.collect.Maps;
-import com.thinkaurelius.titan.core.Cardinality;
-import com.thinkaurelius.titan.core.PropertyKey;
-import com.thinkaurelius.titan.core.SchemaViolationException;
-import com.thinkaurelius.titan.core.TitanGraph;
-import com.thinkaurelius.titan.core.TitanIndexQuery;
-import com.thinkaurelius.titan.core.schema.TitanGraphIndex;
-import com.thinkaurelius.titan.core.schema.TitanManagement;
-import com.thinkaurelius.titan.core.util.TitanCleanup;
+import javax.script.Bindings;
+import javax.script.ScriptEngine;
+import javax.script.ScriptException;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
/**
* Titan 1.0.0 implementation of AtlasGraph.
@@ -179,7 +177,12 @@
@Override
public AtlasIndexQuery<Titan1Vertex, Titan1Edge> indexQuery(String fulltextIndex, String graphQuery) {
- TitanIndexQuery query = getGraph().indexQuery(fulltextIndex, graphQuery);
+ return indexQuery(fulltextIndex, graphQuery, 0);
+ }
+
+ @Override
+ public AtlasIndexQuery<Titan1Vertex, Titan1Edge> indexQuery(String fulltextIndex, String graphQuery, int offset) {
+ TitanIndexQuery query = getGraph().indexQuery(fulltextIndex, graphQuery).offset(offset);
return new Titan1IndexQuery(this, query);
}
diff --git a/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java b/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java
index e6a06c3..1b4583a 100644
--- a/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java
+++ b/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java
@@ -17,14 +17,16 @@
*/
package org.apache.atlas.discovery;
+import org.apache.atlas.ApplicationProperties;
import org.apache.atlas.AtlasConfiguration;
import org.apache.atlas.AtlasErrorCode;
-import org.apache.atlas.model.discovery.AtlasSearchResult.AtlasFullTextResult;
-import org.apache.atlas.model.discovery.AtlasSearchResult.AtlasQueryType;
-import org.apache.atlas.model.discovery.AtlasSearchResult.AttributeSearchResult;
+import org.apache.atlas.AtlasException;
import org.apache.atlas.discovery.graph.DefaultGraphPersistenceStrategy;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.model.discovery.AtlasSearchResult;
+import org.apache.atlas.model.discovery.AtlasSearchResult.AtlasFullTextResult;
+import org.apache.atlas.model.discovery.AtlasSearchResult.AtlasQueryType;
+import org.apache.atlas.model.discovery.AtlasSearchResult.AttributeSearchResult;
import org.apache.atlas.model.instance.AtlasEntity.Status;
import org.apache.atlas.model.instance.AtlasEntityHeader;
import org.apache.atlas.query.Expressions.AliasExpression;
@@ -83,14 +85,20 @@
private final EntityGraphRetriever entityRetriever;
private final AtlasGremlinQueryProvider gremlinQueryProvider;
private final AtlasTypeRegistry typeRegistry;
+ private final int maxResultSetSize;
+ private final int maxTypesCountInIdxQuery;
+ private final int maxTagsCountInIdxQuery;
@Inject
- EntityDiscoveryService(MetadataRepository metadataRepository, AtlasTypeRegistry typeRegistry) {
+ EntityDiscoveryService(MetadataRepository metadataRepository, AtlasTypeRegistry typeRegistry) throws AtlasException {
this.graph = AtlasGraphProvider.getGraphInstance();
this.graphPersistenceStrategy = new DefaultGraphPersistenceStrategy(metadataRepository);
this.entityRetriever = new EntityGraphRetriever(typeRegistry);
this.gremlinQueryProvider = AtlasGremlinQueryProvider.INSTANCE;
this.typeRegistry = typeRegistry;
+ this.maxResultSetSize = ApplicationProperties.get().getInt("atlas.graph.index.search.max-result-set-size", 150);
+ this.maxTypesCountInIdxQuery = ApplicationProperties.get().getInt("atlas.graph.index.search.max-types-count", 10);
+ this.maxTagsCountInIdxQuery = ApplicationProperties.get().getInt("atlas.graph.index.search.max-tags-count", 10);
}
@Override
@@ -217,7 +225,7 @@
if (attribute == null) {
throw new AtlasBaseException(AtlasErrorCode.UNKNOWN_ATTRIBUTE, attrName, typeName);
}
-
+
} else {
// if attrName is null|empty iterate defaultAttrNames to get attribute value
final List<String> defaultAttrNames = new ArrayList<>(Arrays.asList("qualifiedName", "name"));
@@ -238,7 +246,7 @@
} else {
attrQualifiedName = attribute.getQualifiedName();
- String attrQuery = String.format("%s AND (%s *)", attrName, attrValuePrefix.replaceAll("\\.", " "));
+ String attrQuery = String.format("%s AND (%s *)", attrName, attrValuePrefix.replaceAll("\\.", " "));
query = StringUtils.isEmpty(query) ? attrQuery : String.format("(%s) AND (%s)", query, attrQuery);
}
@@ -252,59 +260,72 @@
// if query was provided, perform indexQuery and filter for typeName & classification in memory; this approach
// results in a faster and accurate results than using CONTAINS/CONTAINS_PREFIX filter on entityText property
if (StringUtils.isNotEmpty(query)) {
- final String idxQuery = String.format("v.\"%s\":(%s)", Constants.ENTITY_TEXT_PROPERTY_KEY, query);
- final Iterator<Result<?,?>> qryResult = graph.indexQuery(Constants.FULLTEXT_INDEX, idxQuery).vertices();
- final int startIdx = params.offset();
- final int resultSize = params.limit();
+ final String idxQuery = getQueryForFullTextSearch(query, typeName, classification);
+ final int startIdx = params.offset();
+ final int resultSize = params.limit();
+ int resultIdx = 0;
- int resultIdx = 0;
+ for (int indexQueryOffset = 0; ; indexQueryOffset += getMaxResultSetSize()) {
+ final Iterator<Result<?, ?>> qryResult = graph.indexQuery(Constants.FULLTEXT_INDEX, idxQuery, indexQueryOffset).vertices();
- while (qryResult.hasNext()) {
- AtlasVertex<?,?> vertex = qryResult.next().getVertex();
-
- String vertexTypeName = GraphHelper.getTypeName(vertex);
-
- // skip non-entity vertices
- if (StringUtils.isEmpty(vertexTypeName) || StringUtils.isEmpty(GraphHelper.getGuid(vertex))) {
- continue;
+ if (LOG.isDebugEnabled()) {
+ LOG.debug("indexQuery: query=" + idxQuery + "; offset=" + indexQueryOffset);
}
- if (typeNames != null && !typeNames.contains(vertexTypeName)) {
- continue;
+ if(!qryResult.hasNext()) {
+ break;
}
- if (classificationNames != null) {
- List<String> traitNames = GraphHelper.getTraitNames(vertex);
+ while (qryResult.hasNext()) {
+ AtlasVertex<?, ?> vertex = qryResult.next().getVertex();
+ String vertexTypeName = GraphHelper.getTypeName(vertex);
- if (CollectionUtils.isEmpty(traitNames) ||
- !CollectionUtils.containsAny(classificationNames, traitNames)) {
+ // skip non-entity vertices
+ if (StringUtils.isEmpty(vertexTypeName) || StringUtils.isEmpty(GraphHelper.getGuid(vertex))) {
continue;
}
+
+ if (typeNames != null && !typeNames.contains(vertexTypeName)) {
+ continue;
+ }
+
+ if (classificationNames != null) {
+ List<String> traitNames = GraphHelper.getTraitNames(vertex);
+
+ if (CollectionUtils.isEmpty(traitNames) ||
+ !CollectionUtils.containsAny(classificationNames, traitNames)) {
+ continue;
+ }
+ }
+
+ if (isAttributeSearch) {
+ String vertexAttrValue = vertex.getProperty(attrQualifiedName, String.class);
+
+ if (StringUtils.isNotEmpty(vertexAttrValue) && !vertexAttrValue.startsWith(attrValuePrefix)) {
+ continue;
+ }
+ }
+
+ if (skipDeletedEntities(excludeDeletedEntities, vertex)) {
+ continue;
+ }
+
+ resultIdx++;
+
+ if (resultIdx <= startIdx) {
+ continue;
+ }
+
+ AtlasEntityHeader header = entityRetriever.toAtlasEntityHeader(vertex);
+
+ ret.addEntity(header);
+
+ if (ret.getEntities().size() == resultSize) {
+ break;
+ }
}
- if (isAttributeSearch) {
- String vertexAttrValue = vertex.getProperty(attrQualifiedName, String.class);
-
- if (StringUtils.isNotEmpty(vertexAttrValue) && !vertexAttrValue.startsWith(attrValuePrefix)) {
- continue;
- }
- }
-
- if (skipDeletedEntities(excludeDeletedEntities, vertex)) {
- continue;
- }
-
- resultIdx++;
-
- if (resultIdx <= startIdx) {
- continue;
- }
-
- AtlasEntityHeader header = entityRetriever.toAtlasEntityHeader(vertex);
-
- ret.addEntity(header);
-
- if (ret.getEntities().size() == resultSize) {
+ if (ret.getEntities() != null && ret.getEntities().size() == resultSize) {
break;
}
}
@@ -347,7 +368,7 @@
Object result = graph.executeGremlinScript(scriptEngine, bindings, basicQuery, false);
if (result instanceof List && CollectionUtils.isNotEmpty((List) result)) {
- List queryResult = (List) result;
+ List queryResult = (List) result;
Object firstElement = queryResult.get(0);
if (firstElement instanceof AtlasVertex) {
@@ -371,6 +392,57 @@
return ret;
}
+ private String getQueryForFullTextSearch(String userKeyedString, String typeName, String classification) {
+ String typeFilter = getTypeFilter(typeRegistry, typeName, maxTypesCountInIdxQuery);
+ String classficationFilter = getClassificationFilter(typeRegistry, classification, maxTagsCountInIdxQuery);
+
+ StringBuilder queryText = new StringBuilder();
+
+ if (! StringUtils.isEmpty(userKeyedString)) {
+ queryText.append(userKeyedString);
+ }
+
+ if (! StringUtils.isEmpty(typeFilter)) {
+ if (queryText.length() > 0) {
+ queryText.append(" AND ");
+ }
+
+ queryText.append(typeFilter);
+ }
+
+ if (! StringUtils.isEmpty(classficationFilter)) {
+ if (queryText.length() > 0) {
+ queryText.append(" AND ");
+ }
+
+ queryText.append(classficationFilter);
+ }
+
+ return String.format("v.\"%s\":(%s)", Constants.ENTITY_TEXT_PROPERTY_KEY, queryText.toString());
+ }
+
+ private static String getClassificationFilter(AtlasTypeRegistry typeRegistry, String classificationName, int maxTypesCountInIdxQuery) {
+ AtlasClassificationType classification = typeRegistry.getClassificationTypeByName(classificationName);
+ Set<String> typeAndSubTypes = classification != null ? classification.getTypeAndAllSubTypes() : null;
+
+ if(CollectionUtils.isNotEmpty(typeAndSubTypes) && typeAndSubTypes.size() <= maxTypesCountInIdxQuery) {
+ return String.format("(%s)", StringUtils.join(typeAndSubTypes, " "));
+ }
+
+ return "";
+ }
+
+ private static String getTypeFilter(AtlasTypeRegistry typeRegistry, String typeName, int maxTypesCountInIdxQuery) {
+ AtlasEntityType type = typeRegistry.getEntityTypeByName(typeName);
+ Set<String> typeAndSubTypes = type != null ? type.getTypeAndAllSubTypes() : null;
+
+ if(CollectionUtils.isNotEmpty(typeAndSubTypes) && typeAndSubTypes.size() <= maxTypesCountInIdxQuery) {
+ return String.format("(%s)", StringUtils.join(typeAndSubTypes, " "));
+ }
+
+ return "";
+ }
+
private List<AtlasFullTextResult> getIndexQueryResults(AtlasIndexQuery query, QueryParams params, boolean excludeDeletedEntities) throws AtlasBaseException {
List<AtlasFullTextResult> ret = new ArrayList<>();
Iterator<Result> iter = query.vertices();
@@ -471,4 +543,8 @@
private boolean skipDeletedEntities(boolean excludeDeletedEntities, AtlasVertex<?, ?> vertex) {
return excludeDeletedEntities && GraphHelper.getStatus(vertex) == Status.DELETED;
}
+
+ public int getMaxResultSetSize() {
+ return maxResultSetSize;
+ }
}
diff --git a/repository/src/test/java/org/apache/atlas/services/EntityDiscoveryServiceTest.java b/repository/src/test/java/org/apache/atlas/services/EntityDiscoveryServiceTest.java
new file mode 100644
index 0000000..d503ef7
--- /dev/null
+++ b/repository/src/test/java/org/apache/atlas/services/EntityDiscoveryServiceTest.java
@@ -0,0 +1,124 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.atlas.services;
+
+import org.apache.atlas.TestOnlyModule;
+import org.apache.atlas.discovery.EntityDiscoveryService;
+import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.model.typedef.AtlasEntityDef;
+import org.apache.atlas.type.AtlasTypeRegistry;
+import org.apache.commons.lang.StringUtils;
+import org.powermock.reflect.Whitebox;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Guice;
+import org.testng.annotations.Test;
+
+
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+import static org.testng.Assert.assertEquals;
+
+@Guice(modules = TestOnlyModule.class)
+public class EntityDiscoveryServiceTest {
+
+ private final String TEST_TYPE = "test";
+ private final String TEST_TYPE1 = "test1";
+ private final String TEST_TYPE2 = "test2";
+ private final String TEST_TYPE3 = "test3";
+ private final String TEST_TYPE_WITH_SUB_TYPES = "testTypeWithSubTypes";
+ private AtlasTypeRegistry typeRegistry = new AtlasTypeRegistry();
+
+ AtlasEntityDef typeTest = null;
+ AtlasEntityDef typeTest1 = null;
+ AtlasEntityDef typeTest2 = null;
+ AtlasEntityDef typeTest3 = null;
+ AtlasEntityDef typeWithSubTypes = null;
+
+ private final int maxTypesCountInIdxQuery = 10;
+
+
+ @BeforeClass
+ public void init() throws AtlasBaseException {
+ typeTest = new AtlasEntityDef(TEST_TYPE);
+ typeTest1 = new AtlasEntityDef(TEST_TYPE1);
+ typeTest2 = new AtlasEntityDef(TEST_TYPE2);
+ typeTest3 = new AtlasEntityDef(TEST_TYPE3);
+ typeWithSubTypes = new AtlasEntityDef(TEST_TYPE_WITH_SUB_TYPES);
+
+ typeTest1.addSuperType(TEST_TYPE_WITH_SUB_TYPES);
+ typeTest2.addSuperType(TEST_TYPE_WITH_SUB_TYPES);
+ typeTest3.addSuperType(TEST_TYPE_WITH_SUB_TYPES);
+
+ AtlasTypeRegistry.AtlasTransientTypeRegistry ttr = typeRegistry.lockTypeRegistryForUpdate();
+
+ ttr.addType(typeTest);
+ ttr.addType(typeWithSubTypes);
+ ttr.addType(typeTest1);
+ ttr.addType(typeTest2);
+ ttr.addType(typeTest3);
+
+ typeRegistry.releaseTypeRegistryForUpdate(ttr, true);
+ }
+
+ @Test
+ public void getSubTypesForType_NullStringReturnsEmptyString() throws Exception {
+ invokeGetSubTypesForType(null, maxTypesCountInIdxQuery);
+ }
+
+ @Test
+ public void getSubTypesForType_BlankStringReturnsEmptyString() throws Exception {
+ invokeGetSubTypesForType(" ", maxTypesCountInIdxQuery);
+ }
+
+ @Test
+ public void getSubTypesForType_EmptyStringReturnsEmptyString() throws Exception {
+ invokeGetSubTypesForType("", maxTypesCountInIdxQuery);
+ }
+
+ @Test
+ public void getSubTypeForTypeWithNoSubType_ReturnsTypeString() throws Exception {
+ String s = invokeGetSubTypesForType(TEST_TYPE, 10);
+
+ assertEquals(s, "(" + TEST_TYPE + ")");
+ }
+
+ @Test
+ public void getSubTypeForTypeWithSubTypes_ReturnsOrClause() throws Exception {
+ String s = invokeGetSubTypesForType(TEST_TYPE_WITH_SUB_TYPES, maxTypesCountInIdxQuery);
+
+ assertTrue(s.startsWith("(" + TEST_TYPE_WITH_SUB_TYPES));
+ assertTrue(s.contains(" " + TEST_TYPE1));
+ assertTrue(s.contains(" " + TEST_TYPE2));
+ assertTrue(s.contains(" " + TEST_TYPE3));
+ assertTrue(s.endsWith(")"));
+ }
+
+ @Test
+ public void getSubTypeForTypeWithSubTypes_ReturnsEmptyString() throws Exception {
+ String s = invokeGetSubTypesForType(TEST_TYPE_WITH_SUB_TYPES, 2);
+
+ assertTrue(StringUtils.isBlank(s));
+ }
+
+ private String invokeGetSubTypesForType(String inputString, int maxSubTypes) throws Exception {
+ String s = Whitebox.invokeMethod(EntityDiscoveryService.class, "getTypeFilter", typeRegistry, inputString, maxSubTypes);
+
+ assertNotNull(s);
+ return s;
+ }
+}
diff --git a/typesystem/src/test/resources/atlas-application.properties b/typesystem/src/test/resources/atlas-application.properties
index 5ffde5e..c4ce5ea 100644
--- a/typesystem/src/test/resources/atlas-application.properties
+++ b/typesystem/src/test/resources/atlas-application.properties
@@ -72,7 +72,7 @@
# Solr cloud mode properties
atlas.graph.index.search.solr.mode=cloud
atlas.graph.index.search.solr.zookeeper-url=${solr.zk.address}
-
+atlas.graph.index.search.max-result-set-size=150
######### Hive Lineage Configs #########
## Schema
