Merge pull request #83 from shwethags/edge-label

repository/src/main/java/org/apache/hadoop/metadata/repository/graph/GraphBackedMetadataRepository.java

@@ -128,7 +128,16 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
 
     @Override
     public String getEdgeLabel(IDataType<?> dataType, AttributeInfo aInfo) {
-        return EDGE_LABEL_PREFIX + dataType.getName() + "." + aInfo.name;
+        return getEdgeLabel(dataType.getName(), aInfo.name);
+    }
+
+    public String getEdgeLabel(String typeName, String attrName) {
+        return EDGE_LABEL_PREFIX + typeName + "." + attrName;
+    }
+
+    public String getEdgeLabel(ITypedInstance typedInstance, AttributeInfo aInfo) throws MetadataException {
+        IDataType dataType = typeSystem.getDataType(IDataType.class, typedInstance.getTypeName());
+        return getEdgeLabel(dataType, aInfo);
     }
 
     @Override
@@ -275,7 +284,7 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
             }
 
             final String entityTypeName = getTypeName(instanceVertex);
-            String relationshipLabel = entityTypeName + "." + traitNameToBeDeleted;
+            String relationshipLabel = getEdgeLabel(entityTypeName, traitNameToBeDeleted);
             Iterator<Edge> results = instanceVertex.getEdges(
                     Direction.OUT, relationshipLabel).iterator();
             if (results.hasNext()) { // there should only be one edge for this label
@@ -673,6 +682,7 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
             Object attrValue = typedInstance.get(attributeInfo.name);
             LOG.debug("mapping attribute {} = {}", attributeInfo.name, attrValue);
             final String propertyName = getQualifiedName(typedInstance, attributeInfo);
+            String edgeLabel = getEdgeLabel(typedInstance, attributeInfo);
             if (attrValue == null) {
                 return;
             }
@@ -698,11 +708,10 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
 
                 case STRUCT:
                     Vertex structInstanceVertex = mapStructInstanceToVertex(id,
-                            (ITypedStruct) typedInstance.get(attributeInfo.name),
-                            attributeInfo, idToVertexMap);
+                            (ITypedStruct) typedInstance.get(attributeInfo.name), attributeInfo, idToVertexMap);
                     // add an edge to the newly created vertex from the parent
                     GraphHelper.addEdge(
-                            titanGraph, instanceVertex, structInstanceVertex, propertyName);
+                            titanGraph, instanceVertex, structInstanceVertex, edgeLabel);
                     break;
 
                 case TRAIT:
@@ -712,7 +721,7 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
                 case CLASS:
                     Id referenceId = (Id) typedInstance.get(attributeInfo.name);
                     mapClassReferenceAsEdge(
-                            instanceVertex, idToVertexMap, propertyName, referenceId);
+                            instanceVertex, idToVertexMap, edgeLabel, referenceId);
                     break;
 
                 default:
@@ -886,7 +895,7 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
                     traitInstance.fieldMapping().fields, idToVertexMap);
 
             // add an edge to the newly created vertex from the parent
-            String relationshipLabel = typedInstanceTypeName + "." + traitName;
+            String relationshipLabel = getEdgeLabel(typedInstanceTypeName, traitName);
             GraphHelper.addEdge(
                     titanGraph, parentInstanceVertex, traitInstanceVertex, relationshipLabel);
         }
@@ -1017,7 +1026,7 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
                     break;
 
                 case CLASS:
-                    String relationshipLabel = getQualifiedName(typedInstance, attributeInfo);
+                    String relationshipLabel = getEdgeLabel(typedInstance, attributeInfo);
                     Object idOrInstance = mapClassReferenceToVertex(instanceVertex,
                             attributeInfo, relationshipLabel, attributeInfo.dataType());
                     typedInstance.set(attributeInfo.name, idOrInstance);
@@ -1221,7 +1230,7 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
             ITypedStruct structInstance = structType.createInstance();
             typedInstance.set(attributeInfo.name, structInstance);
 
-            String relationshipLabel = getQualifiedName(typedInstance, attributeInfo);
+            String relationshipLabel = getEdgeLabel(typedInstance, attributeInfo);
             LOG.debug("Finding edge for {} -> label {} ", instanceVertex, relationshipLabel);
             for (Edge edge : instanceVertex.getEdges(Direction.OUT, relationshipLabel)) {
                 final Vertex structInstanceVertex = edge.getVertex(Direction.IN);

repository/src/test/java/org/apache/hadoop/metadata/discovery/GraphBackedDiscoveryServiceTest.java

@@ -29,6 +29,7 @@ import org.apache.hadoop.metadata.discovery.graph.GraphBackedDiscoveryService;
 import org.apache.hadoop.metadata.query.HiveTitanSample;
 import org.apache.hadoop.metadata.query.QueryTestsUtils;
 import org.apache.hadoop.metadata.repository.graph.GraphBackedMetadataRepository;
+import org.apache.hadoop.metadata.repository.graph.GraphBackedSearchIndexer;
 import org.apache.hadoop.metadata.repository.graph.GraphHelper;
 import org.apache.hadoop.metadata.repository.graph.GraphProvider;
 import org.apache.hadoop.metadata.typesystem.ITypedReferenceableInstance;
@@ -224,6 +225,7 @@ public class GraphBackedDiscoveryServiceTest {
             {"Table as _loop0 loop (LoadProcess outputTable) withPath"},
             {"Table as src loop (LoadProcess outputTable) as dest select src.name as srcTable, dest.name as destTable withPath"},
             {"Table as t, sd, Column as c where t.name=\"sales_fact\" select c.name as colName, c.dataType as colType"},
+            {"Table where name='sales_fact', db where name='Reporting'"}
         };
     }
 
@@ -268,39 +270,6 @@ public class GraphBackedDiscoveryServiceTest {
     }
 
     @Test
-    public void testSearchByDSLQuery() throws Exception {
-        String dslQuery = "Column as PII";
-        System.out.println("Executing dslQuery = " + dslQuery);
-        String jsonResults = discoveryService.searchByDSL(dslQuery);
-        Assert.assertNotNull(jsonResults);
-
-        JSONObject results = new JSONObject(jsonResults);
-        Assert.assertEquals(results.length(), 3);
-        System.out.println("results = " + results);
-
-        Object query = results.get("query");
-        Assert.assertNotNull(query);
-
-        JSONObject dataType = results.getJSONObject("dataType");
-        Assert.assertNotNull(dataType);
-        String typeName = dataType.getString("typeName");
-        Assert.assertNotNull(typeName);
-
-        JSONArray rows = results.getJSONArray("rows");
-        Assert.assertNotNull(rows);
-        Assert.assertTrue(rows.length() > 0);
-
-        for (int index = 0; index < rows.length(); index++) {
-            JSONObject row = rows.getJSONObject(index);
-            String type = row.getString("$typeName$");
-            Assert.assertEquals(type, "Column");
-
-            String name = row.getString("name");
-            Assert.assertNotEquals(name, "null");
-        }
-    }
-
-    @Test
     public void testSearchForTypeInheritance() throws Exception {
         createTypesWithMultiLevelInheritance();
         createInstances();

repository/src/test/java/org/apache/hadoop/metadata/repository/graph/GraphBackedMetadataRepositoryTest.java

@@ -145,9 +145,8 @@ public class GraphBackedMetadataRepositoryTest {
 
     @Test (dependsOnMethods = "testSubmitEntity")
     public void testGetTraitLabel() throws Exception {
-        Assert.assertEquals(repositoryService.getTraitLabel(
-                        typeSystem.getDataType(ClassType.class, TABLE_TYPE),
-                        CLASSIFICATION), TABLE_TYPE + "." + CLASSIFICATION);
+        Assert.assertEquals(repositoryService.getTraitLabel(typeSystem.getDataType(ClassType.class, TABLE_TYPE),
+                CLASSIFICATION), TABLE_TYPE + "." + CLASSIFICATION);
     }
 
     @Test
@@ -317,6 +316,39 @@ public class GraphBackedMetadataRepositoryTest {
         Assert.assertEquals(repositoryService.getTypeName(tableVertex), TABLE_TYPE);
     }
 
+    @Test(dependsOnMethods = "testCreateEntity")
+    public void testSearchByDSLQuery() throws Exception {
+        String dslQuery = "hive_database as PII";
+        System.out.println("Executing dslQuery = " + dslQuery);
+        String jsonResults = discoveryService.searchByDSL(dslQuery);
+        Assert.assertNotNull(jsonResults);
+
+        JSONObject results = new JSONObject(jsonResults);
+        Assert.assertEquals(results.length(), 3);
+        System.out.println("results = " + results);
+
+        Object query = results.get("query");
+        Assert.assertNotNull(query);
+
+        JSONObject dataType = results.getJSONObject("dataType");
+        Assert.assertNotNull(dataType);
+        String typeName = dataType.getString("typeName");
+        Assert.assertNotNull(typeName);
+
+        JSONArray rows = results.getJSONArray("rows");
+        Assert.assertNotNull(rows);
+        Assert.assertTrue(rows.length() > 0);
+
+        for (int index = 0; index < rows.length(); index++) {
+            JSONObject row = rows.getJSONObject(index);
+            String type = row.getString("$typeName$");
+            Assert.assertEquals(type, "hive_database");
+
+            String name = row.getString("name");
+            Assert.assertEquals(name, DATABASE_NAME);
+        }
+    }
+
     /**
       * Full text search requires GraphBackedSearchIndexer, and GraphBackedSearchIndexer can't be enabled in
       * GraphBackedDiscoveryServiceTest because of its test data. So, test for full text search is in