ATLAS-2433: DSL: Improved support for numeric data type.

repository/src/main/java/org/apache/atlas/query/GremlinQueryComposer.java

@@ -169,8 +169,12 @@ public class GremlinQueryComposer {
             rhs = parseDate(rhs);
         }
 
+        if (lhsI.isNumeric()) {
+            rhs = parseNumber(rhs);
+        }
+
+        rhs = addQuotesIfNecessary(lhsI, rhs);
         SearchParameters.Operator op = SearchParameters.Operator.fromString(operator);
-        rhs = addQuotesIfNecessary(rhs);
         if (op == SearchParameters.Operator.LIKE) {
             add(GremlinClause.TEXT_CONTAINS, lhsI.getQualifiedName(), IdentifierHelper.getFixedRegEx(rhs));
         } else if (op == SearchParameters.Operator.IN) {
@@ -188,6 +192,10 @@ public class GremlinQueryComposer {
         }
     }
 
+    private String parseNumber(String rhs) {
+        return rhs.replace("'", "").replace("\"", "");
+    }
+
     public void addAndClauses(List<String> clauses) {
         add(GremlinClause.AND, String.join(",", clauses));
     }
@@ -464,7 +472,8 @@ public class GremlinQueryComposer {
         add(GremlinClause.INLINE_TRANSFORM_CALL);
     }
 
-    private String addQuotesIfNecessary(String rhs) {
+    private String addQuotesIfNecessary(IdentifierHelper.Info rhsI, String rhs) {
+        if(rhsI.isNumeric()) return rhs;
         if (IdentifierHelper.isTrueOrFalse(rhs)) return rhs;
         if (IdentifierHelper.isQuoted(rhs)) return rhs;
         return IdentifierHelper.getQuoted(rhs);

repository/src/main/java/org/apache/atlas/query/IdentifierHelper.java

@@ -122,6 +122,7 @@ public class IdentifierHelper {
         private boolean  isAttribute;
         private String   qualifiedName;
         private boolean  isDate;
+        private boolean  isNumeric;
 
         public Info(String s) {
             this.raw = removeQuotes(s);
@@ -196,6 +197,7 @@ public class IdentifierHelper {
             isPrimitive = lookup.isPrimitive(context, attributeName);
             setQualifiedName(lookup, context, isAttribute, attributeName);
             setIsDate(lookup, context, isPrimitive, attributeName);
+            setIsNumeric(lookup, context, isPrimitive, attributeName);
         }
 
         private String getDefaultQualifiedNameForSinglePartName(GremlinQueryComposer.Context context, String s) {
@@ -223,6 +225,12 @@ public class IdentifierHelper {
             }
         }
 
+        private void setIsNumeric(Lookup lookup, GremlinQueryComposer.Context context, boolean isPrimitive, String attrName) {
+            if (isPrimitive) {
+                isNumeric = lookup.isNumeric(context, attrName);
+            }
+        }
+
         private void updateParts() {
             parts = StringUtils.split(raw, ".");
         }
@@ -283,5 +291,8 @@ public class IdentifierHelper {
             return raw;
         }
 
+        public boolean isNumeric() {
+            return isNumeric;
+        }
     }
 }

repository/src/main/java/org/apache/atlas/query/Lookup.java

@@ -41,4 +41,6 @@ public interface Lookup {
     String getTypeFromEdge(GremlinQueryComposer.Context context, String item);
 
     boolean isDate(GremlinQueryComposer.Context context, String attributeName);
+
+    boolean isNumeric(GremlinQueryComposer.Context context, String attrName);
 }

repository/src/main/java/org/apache/atlas/query/RegistryBasedLookup.java

@@ -38,6 +38,15 @@ class RegistryBasedLookup implements Lookup {
                     Constants.TIMESTAMP_PROPERTY_KEY,
                     Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY));
 
+    private static final Set<String> NUMERIC_ATTRIBUTES = new HashSet<>(
+            Arrays.asList(AtlasBaseTypeDef.ATLAS_TYPE_SHORT,
+                    AtlasBaseTypeDef.ATLAS_TYPE_INT,
+                    AtlasBaseTypeDef.ATLAS_TYPE_LONG,
+                    AtlasBaseTypeDef.ATLAS_TYPE_FLOAT,
+                    AtlasBaseTypeDef.ATLAS_TYPE_DOUBLE,
+                    AtlasBaseTypeDef.ATLAS_TYPE_BIGINTEGER,
+                    AtlasBaseTypeDef.ATLAS_TYPE_BIGDECIMAL));
+
     private final AtlasTypeRegistry typeRegistry;
 
     public RegistryBasedLookup(AtlasTypeRegistry typeRegistry) {
@@ -201,6 +210,16 @@ class RegistryBasedLookup implements Lookup {
 
         AtlasType attr = et.getAttributeType(attributeName);
         return attr != null && attr.getTypeName().equals(AtlasBaseTypeDef.ATLAS_TYPE_DATE);
+    }
+
+    @Override
+    public boolean isNumeric(GremlinQueryComposer.Context context, String attrName) {
+        AtlasEntityType et = context.getActiveEntityType();
+        if (et == null) {
+            return false;
+        }
 
+        AtlasType attr = et.getAttributeType(attrName);
+        return attr != null && NUMERIC_ATTRIBUTES.contains(attr.getTypeName());
     }
 }

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

@@ -159,6 +159,7 @@ public class DSLQueriesTest extends BasicTestSetup {
                 {"Person where (houseNumber < 153)", 3},
                 {"Person where (houseNumber <= 153)", 4},
                 {"Person where (houseNumber =  17)", 1},
+                {"Person where houseNumber >= 17 or numberOfCars = 2", 2},
                 {"Person where (houseNumber != 17)", 3},
 
                 {"Person where (carMileage > 0)", 2},

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

@@ -306,6 +306,13 @@ public class GremlinQueryComposerTest {
         verify("Table has name and Table has owner and name = 'sales_fact'", "g.V().has('__typeName', 'Table').and(__.has('Table.name'),__.has('Table.owner'),__.has('Table.name', eq('sales_fact'))).limit(local, 25).limit(25).toList()");
     }
 
+
+    @Test
+    public void numericAttributes() {
+        verify("Table where partitionSize = 2048", "g.V().has('__typeName', 'Table').has('Table.partitionSize', eq(2048)).limit(local, 25).limit(25).toList()");
+        verify("Table where partitionSize = 2048 or partitionSize = 10", "g.V().has('__typeName', 'Table').or(__.has('Table.partitionSize', eq(2048)),__.has('Table.partitionSize', eq(10))).limit(local, 25).limit(25).toList()");
+    }
+
     @Test
     public void systemAttributes() {
         verify("Table has __state", "g.V().has('__typeName', 'Table').has('__state').limit(local, 25).limit(25).toList()");
@@ -437,7 +444,8 @@ public class GremlinQueryComposerTest {
                     attributeName.equals("createTime") ||
                     attributeName.equals("clusterName") ||
                     attributeName.equals("__guid") ||
-                    attributeName.equals("__state");
+                    attributeName.equals("__state") ||
+                    attributeName.equals("partitionSize");
         }
 
         @Override
@@ -461,6 +469,7 @@ public class GremlinQueryComposerTest {
                     (context.getActiveTypeName().equals("Table") && attributeName.equals("isFile")) ||
                     (context.getActiveTypeName().equals("Table") && attributeName.equals("__guid")) ||
                     (context.getActiveTypeName().equals("Table") && attributeName.equals("__state")) ||
+                    (context.getActiveTypeName().equals("Table") && attributeName.equals("partitionSize")) ||
                     (context.getActiveTypeName().equals("hive_db") && attributeName.equals("name")) ||
                     (context.getActiveTypeName().equals("hive_db") && attributeName.equals("owner")) ||
                     (context.getActiveTypeName().equals("hive_db") && attributeName.equals("createTime")) ||
@@ -491,13 +500,21 @@ public class GremlinQueryComposerTest {
         public String getTypeFromEdge(GremlinQueryComposer.Context context, String item) {
             if(context.getActiveTypeName().equals("DB") && item.equals("Table")) {
                 return "Table";
-            } else if(context.getActiveTypeName().equals("Table") && item.equals("Column")) {
+            }
+
+            if(context.getActiveTypeName().equals("Table") && item.equals("Column")) {
                 return "Column";
-            } else if(context.getActiveTypeName().equals("Table") && item.equals("db")) {
+            }
+
+            if(context.getActiveTypeName().equals("Table") && item.equals("db")) {
                 return "DB";
-            } else if(context.getActiveTypeName().equals("Table") && item.equals("columns")) {
+            }
+
+            if(context.getActiveTypeName().equals("Table") && item.equals("columns")) {
                 return "Column";
-            } else if(context.getActiveTypeName().equals(item)) {
+            }
+
+            if(context.getActiveTypeName().equals(item)) {
                 return null;
             }
             return context.getActiveTypeName();
@@ -507,5 +524,10 @@ public class GremlinQueryComposerTest {
         public boolean isDate(GremlinQueryComposer.Context context, String attributeName) {
             return attributeName.equals("createTime");
         }
+
+        @Override
+        public boolean isNumeric(GremlinQueryComposer.Context context, String attrName) {
+            return attrName.equals("partitionSize");
+        }
     }
 }