ATLAS-1589 : DSL queries returns wrong object when filter traverses edges

release-log.txt

@@ -9,6 +9,7 @@ ATLAS-1060 Add composite indexes for exact match performance improvements for al
 ATLAS-1127 Modify creation and modification timestamps to Date instead of Long(sumasai)
 
 ALL CHANGES:
+ATLAS_1589 DSL queries return wrong object when filter traverses an edge (jnhagelberg)
 ATLAS-1590 UI : Edit Button is enabled for Deleted entities. (kevalbhatt)
 ATLAS-1487 Create Entity in UI : types list doesn't list fs_permissions (struct type) hence no entity could be created for it. (kevalbhatt)
 ATLAS-1573 Full text mapping for Entity store V2

repository/src/main/java/org/apache/atlas/gremlin/Gremlin2ExpressionFactory.java

@@ -337,5 +337,13 @@ public class Gremlin2ExpressionFactory extends GremlinExpressionFactory {
        LiteralExpression aliasName =  (LiteralExpression)fc.getArguments().get(0);
        return Collections.singletonList(aliasName.getValue().toString());
     }
+
+    @Override
+    public boolean isRepeatExpression(GroovyExpression expr) {
+        if(!(expr instanceof FunctionCallExpression)) {
+            return false;
+        }
+        return ((FunctionCallExpression)expr).getFunctionName().equals(LOOP_METHOD);
+    }
 }
 

repository/src/main/java/org/apache/atlas/gremlin/Gremlin3ExpressionFactory.java

@@ -443,4 +443,12 @@ public class Gremlin3ExpressionFactory extends GremlinExpressionFactory {
     public List<String> getAliasesRequiredByExpression(GroovyExpression expr) {
         return Collections.emptyList();
     }
+
+    @Override
+    public boolean isRepeatExpression(GroovyExpression expr) {
+        if(!(expr instanceof FunctionCallExpression)) {
+            return false;
+        }
+        return ((FunctionCallExpression)expr).getFunctionName().equals(REPEAT_METHOD);
+    }
 }

repository/src/main/java/org/apache/atlas/gremlin/GremlinExpressionFactory.java

@@ -649,4 +649,6 @@ public abstract class GremlinExpressionFactory {
        return aliasName.getValue().toString();
 
     }
+
+    public abstract boolean isRepeatExpression(GroovyExpression expr);
 }
\ No newline at end of file

repository/src/main/java/org/apache/atlas/gremlin/optimizer/ExpandOrsOptimization.java

@@ -288,10 +288,14 @@ public class ExpandOrsOptimization implements GremlinOptimization {
 
     private GroovyExpression removeMapFromPathsIfNeeded(GroovyExpression expr, List<LiteralExpression> aliases) {
         if(aliases.size() > 0 && factory.isSelectGeneratesMap(aliases.size())) {
+            RepeatExpressionFinder repeatExprFinder = new RepeatExpressionFinder(factory);
+            GremlinQueryOptimizer.visitCallHierarchy(expr, repeatExprFinder);
+            boolean hasRepeat = repeatExprFinder.isRepeatExpressionFound();
+
             PathExpressionFinder pathExprFinder = new PathExpressionFinder();
             GremlinQueryOptimizer.visitCallHierarchy(expr, pathExprFinder);
             boolean hasPath = pathExprFinder.isPathExpressionFound();
-            if(hasPath) {
+            if(! hasRepeat && hasPath) {
                 //the path will now start with the map that we added.  That is an artifact
                 //of the optimization process and must be removed.
                 if(expr.getType() != TraversalStepType.END && expr.getType() != TraversalStepType.NONE) {

repository/src/main/java/org/apache/atlas/gremlin/optimizer/RepeatExpressionFinder.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
+ *
+ *     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.gremlin.optimizer;
+
+import org.apache.atlas.gremlin.GremlinExpressionFactory;
+import org.apache.atlas.groovy.AbstractFunctionExpression;
+import org.apache.atlas.groovy.GroovyExpression;
+
+/**
+ * Determines whether an expression contains a repeat/loop function.
+ */
+public class RepeatExpressionFinder implements CallHierarchyVisitor {
+
+    private boolean found = false;
+    private GremlinExpressionFactory factory;
+
+    public RepeatExpressionFinder(GremlinExpressionFactory factory) {
+        this.factory = factory;
+    }
+
+    @Override
+    public boolean preVisitFunctionCaller(AbstractFunctionExpression expr) {
+
+        found = factory.isRepeatExpression(expr);
+        if(found) {
+            return false;
+        }
+        return true;
+    }
+
+    @Override
+    public void visitNonFunctionCaller(GroovyExpression expr) {
+
+    }
+
+    @Override
+    public void visitNullCaller() {
+
+    }
+
+    public boolean isRepeatExpressionFound() {
+        return found;
+    }
+
+    @Override
+    public boolean postVisitFunctionCaller(AbstractFunctionExpression functionCall) {
+
+        return false;
+    }
+}

repository/src/main/scala/org/apache/atlas/query/GremlinQuery.scala

@@ -440,8 +440,13 @@ class GremlinTranslator(expr: Expression,
            return GremlinExpressionFactory.INSTANCE.generateHasExpression(gPersistenceBehavior, childExpr, qualifiedPropertyName, c.symbol, persistentExprValue, fInfo);
         }
         case fil@FilterExpression(child, condExpr) => {
-            val newParent = genQuery(parent, child, inClosure);
-            return genQuery(newParent, condExpr, inClosure);
+            var newParent = genQuery(parent, child, inClosure);
+            val alias = "a" + counter.next;
+            newParent = GremlinExpressionFactory.INSTANCE.generateAliasExpression(newParent, alias);
+            val translated =  genQuery(newParent, condExpr, inClosure);
+            //we want the query to return instances of the class whose instances we are filtering out
+            //The act of filtering may traverse edges and have other side effects.
+            GremlinExpressionFactory.INSTANCE.generateBackReferenceExpression(translated, false, alias);
         }
         case l@LogicalExpression(symb, children) => {
             val translatedChildren : java.util.List[GroovyExpression] = translateList(children, false);

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

@@ -1228,6 +1228,23 @@ public class GraphBackedDiscoveryServiceTest extends BaseRepositoryTest {
         assertEquals(rows.length(), 0);
     }
 
+    @Test
+    public void testTypePreservedWhenFilterTraversesEdges() throws DiscoveryException, JSONException {
+
+        String dsl = "hive_table db.name=\"Reporting\" limit 10";
+        ImmutableSet<String> expectedTableNames = ImmutableSet.of("table1", "table2", "sales_fact_monthly_mv", "sales_fact_daily_mv");
+        String jsonResults = discoveryService.searchByDSL(dsl, null);
+        assertNotNull(jsonResults);
+
+        JSONObject results = new JSONObject(jsonResults);
+        JSONArray rows = results.getJSONArray("rows");
+        assertEquals(rows.length(), expectedTableNames.size());
+        for(int i = 0; i < rows.length(); i++) {
+            JSONObject row = rows.getJSONObject(i);
+            Assert.assertTrue(expectedTableNames.contains(row.get("name")));
+        }
+    }
+
     private FieldValueValidator makeCountValidator(int count) {
         return new FieldValueValidator().withFieldNames("count()").withExpectedValues(count);
     }