ATLAS-3776: fixed basic-search handling of sortBy attribute while using graphQuery

graphdb/common/src/main/java/org/apache/atlas/repository/graphdb/tinkerpop/query/TinkerpopGraphQuery.java

@@ -212,7 +212,7 @@ public abstract class TinkerpopGraphQuery<V, E> implements AtlasGraphQuery<V, E>
         Preconditions.checkArgument(limit >= 0, "Limit must be non-negative");
 
         // Compute the overall result by combining the results of all the AndConditions (nested within OR) together.
-        Set<AtlasVertex<V, E>> result = new HashSet<>();
+        Set<AtlasVertex<V, E>> result = new LinkedHashSet<>();
         long resultIdx = 0;
         for(AndCondition andExpr : queryCondition.getAndTerms()) {
             if (result.size() == limit) {

repository/src/main/java/org/apache/atlas/discovery/EntitySearchProcessor.java

@@ -25,7 +25,7 @@ import org.apache.atlas.repository.graphdb.AtlasIndexQuery;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.type.AtlasClassificationType;
 import org.apache.atlas.type.AtlasEntityType;
-import org.apache.atlas.type.AtlasStructType;
+import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
 import org.apache.atlas.util.SearchPredicateUtil;
 import org.apache.atlas.utils.AtlasPerfTracer;
 import org.apache.commons.collections.CollectionUtils;
@@ -42,6 +42,7 @@ import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Set;
+import java.util.stream.StreamSupport;
 
 import static org.apache.atlas.SortOrder.ASCENDING;
 import static org.apache.atlas.discovery.SearchContext.MATCH_ALL_CLASSIFICATION_TYPES;
@@ -197,13 +198,15 @@ public class EntitySearchProcessor extends SearchProcessor {
                     graphQueryPredicate = activePredicate;
                 }
             }
-
             if (sortBy != null && !sortBy.isEmpty()) {
-                AtlasGraphQuery.SortOrder qrySortOrder = sortOrder == SortOrder.ASCENDING ? ASC : DESC;
-                graphQuery.orderBy(sortBy, qrySortOrder);
-            }
+                AtlasAttribute sortByAttribute = context.getEntityType().getAttribute(sortBy);
 
+                if (sortByAttribute != null) {
+                    AtlasGraphQuery.SortOrder qrySortOrder = sortOrder == SortOrder.ASCENDING ? ASC : DESC;
 
+                    graphQuery.orderBy(sortByAttribute.getVertexPropertyName(), qrySortOrder);
+                }
+            }
         } else {
             graphQuery = null;
             graphQueryPredicate = null;
@@ -263,7 +266,7 @@ public class EntitySearchProcessor extends SearchProcessor {
             String sortBy = context.getSearchParameters().getSortBy();
 
             final AtlasEntityType entityType = context.getEntityType();
-            AtlasStructType.AtlasAttribute sortByAttribute = entityType.getAttribute(sortBy);
+            AtlasAttribute sortByAttribute = entityType.getAttribute(sortBy);
             if (sortByAttribute == null) {
                 sortBy = null;
             } else {
@@ -360,6 +363,12 @@ public class EntitySearchProcessor extends SearchProcessor {
 
     @Override
     public long getResultCount() {
-        return (indexQuery != null) ? indexQuery.vertexTotals() : -1;
+        if (indexQuery != null) {
+            return indexQuery.vertexTotals();
+        } else if (graphQuery != null) {
+            return StreamSupport.stream(graphQuery.vertexIds().spliterator(), false).count();
+        } else {
+            return -1L;
+        }
     }
 }

repository/src/test/java/org/apache/atlas/query/BasicTestSetup.java → repository/src/test/java/org/apache/atlas/BasicTestSetup.java

@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.atlas.query;
+package org.apache.atlas;
 
 import com.google.common.collect.ImmutableList;
 import org.apache.atlas.AtlasClient;
@@ -32,6 +32,8 @@ import org.apache.atlas.type.AtlasTypeRegistry;
 
 import javax.inject.Inject;
 import java.io.IOException;
+import java.time.LocalDate;
+import java.time.ZoneId;
 import java.util.*;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -179,6 +181,7 @@ public abstract class BasicTestSetup {
         AtlasEntity salesFactDaily =
                 table("sales_fact_daily_mv", "sales fact daily materialized view", reportingDB, sd, "Joe BI", "Managed",
                       salesFactColumns, "Metric");
+        salesFactDaily.setAttribute("createTime", Date.from(LocalDate.of(2016, 8, 19).atStartOfDay(ZoneId.systemDefault()).toInstant()));
         entities.add(salesFactDaily);
 
         sd = storageDescriptor("hdfs://host:8000/apps/warehouse/sales", "TextInputFormat", "TextOutputFormat", true, ImmutableList.of(column("time_id", "int", "time id")));

repository/src/test/java/org/apache/atlas/discovery/EntitySearchProcessorTest.java

+/**
+ * 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
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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.discovery;
+
+import com.google.common.collect.Sets;
+import org.apache.atlas.BasicTestSetup;
+import org.apache.atlas.SortOrder;
+import org.apache.atlas.TestModules;
+import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.model.discovery.SearchParameters;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
+import org.apache.atlas.repository.graphdb.AtlasVertex;
+import org.apache.atlas.repository.store.graph.v2.EntityGraphRetriever;
+import org.apache.atlas.type.AtlasTypeRegistry;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Guice;
+import org.testng.annotations.Test;
+
+import javax.inject.Inject;
+import java.util.Collections;
+import java.util.Date;
+import java.util.List;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+@Guice(modules = TestModules.TestOnlyModule.class)
+public class EntitySearchProcessorTest extends BasicTestSetup {
+
+    @Inject
+    private AtlasGraph graph;
+
+    @Inject
+    private AtlasTypeRegistry typeRegistry;
+
+    @Inject
+    private EntityGraphRetriever entityRetriever;
+
+    @BeforeClass
+    public void setup() {
+        setupTestData();
+    }
+
+
+    @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);
+    }
+
+    @Test
+    public void searchByClassificationSortBy() throws AtlasBaseException {
+        SearchParameters params = new SearchParameters();
+        params.setTypeName("hive_table");
+        params.setClassification("Metric");
+        params.setLimit(10);
+        params.setSortBy("createTime");
+        params.setSortOrder(SortOrder.ASCENDING);
+
+        SearchContext context = new SearchContext(params, typeRegistry, graph, Collections.<String>emptySet());
+        EntitySearchProcessor processor = new EntitySearchProcessor(context);
+
+        List<AtlasVertex> vertices = processor.execute();
+
+        assertEquals(processor.getResultCount(), 4);
+        assertEquals(vertices.size(), 4);
+
+
+        AtlasVertex firstVertex = vertices.get(0);
+
+        Date firstDate = (Date) entityRetriever.toAtlasEntityHeader(firstVertex, Sets.newHashSet(params.getSortBy())).getAttribute(params.getSortBy());
+
+        AtlasVertex secondVertex = vertices.get(1);
+        Date secondDate = (Date) entityRetriever.toAtlasEntityHeader(secondVertex, Sets.newHashSet(params.getSortBy())).getAttribute(params.getSortBy());
+
+        assertTrue(firstDate.before(secondDate));
+    }
+
+    @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);
+    }
+
+    @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);
+    }
+
+}

repository/src/test/java/org/apache/atlas/query/DSLQueriesTest.java

@@ -17,6 +17,7 @@
  */
 package org.apache.atlas.query;
 
+import org.apache.atlas.BasicTestSetup;
 import org.apache.atlas.TestModules;
 import org.apache.atlas.discovery.EntityDiscoveryService;
 import org.apache.atlas.exception.AtlasBaseException;