ATLAS-2422: Added to master from branch-0.8.

intg/src/main/java/org/apache/atlas/model/impexp/AtlasExportRequest.java

@@ -54,6 +54,7 @@ public class AtlasExportRequest implements Serializable {
     public static final String MATCH_TYPE_ENDS_WITH   = "endsWith";
     public static final String MATCH_TYPE_CONTAINS    = "contains";
     public static final String MATCH_TYPE_MATCHES     = "matches";
+    public static final String MATCH_TYPE_FOR_TYPE     = "forType";
 
     private List<AtlasObjectId> itemsToExport = new ArrayList<>();
     private Map<String, Object> options       = new HashMap<>();

repository/src/main/java/org/apache/atlas/repository/impexp/ExportService.java

@@ -60,13 +60,7 @@ import org.springframework.stereotype.Component;
 import javax.inject.Inject;
 import javax.script.ScriptEngine;
 import javax.script.ScriptException;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 
 import static org.apache.atlas.model.impexp.AtlasExportRequest.*;
 
@@ -182,19 +176,19 @@ public class ExportService {
         return overall;
     }
 
-    private AtlasExportResult.OperationStatus processObjectId(AtlasObjectId item, ExportContext context) throws AtlasServiceException, AtlasException, AtlasBaseException {
+    private AtlasExportResult.OperationStatus processObjectId(AtlasObjectId item, ExportContext context) {
         if (LOG.isDebugEnabled()) {
             LOG.debug("==> processObjectId({})", item);
         }
 
         try {
-            List<AtlasEntityWithExtInfo> entities = getStartingEntity(item, context);
-            if(entities.size() == 0) {
+            List<String> entityGuids = getStartingEntity(item, context);
+            if(entityGuids.size() == 0) {
                 return AtlasExportResult.OperationStatus.FAIL;
             }
 
-            for (AtlasEntityWithExtInfo entityWithExtInfo : entities) {
-                processEntity(entityWithExtInfo.getEntity().getGuid(), context);
+            for (String guid : entityGuids) {
+                processEntity(guid, context);
             }
 
             while (!context.guidsToProcess.isEmpty()) {
@@ -221,72 +215,100 @@ public class ExportService {
         return AtlasExportResult.OperationStatus.SUCCESS;
     }
 
-    private List<AtlasEntityWithExtInfo> getStartingEntity(AtlasObjectId item, ExportContext context) throws AtlasBaseException {
-        List<AtlasEntityWithExtInfo> ret = new ArrayList<>();
+    private List<String> getStartingEntity(AtlasObjectId item, ExportContext context) throws AtlasBaseException {
+        List<String> ret = null;
 
         if (StringUtils.isNotEmpty(item.getGuid())) {
-            AtlasEntityWithExtInfo entity = entityGraphRetriever.toAtlasEntityWithExtInfo(item);
+            ret = Collections.singletonList(item.getGuid());
+        } else if (StringUtils.equalsIgnoreCase(context.matchType, MATCH_TYPE_FOR_TYPE) && StringUtils.isNotEmpty(item.getTypeName())) {
+            final String queryTemplate = getQueryTemplateForMatchType(context);
 
-            if (entity != null) {
-                ret = Collections.singletonList(entity);
-            }
+            setupBindingsForTypeName(context, item.getTypeName());
+
+            ret = executeGremlinQueryForGuids(queryTemplate, context);
         } else if (StringUtils.isNotEmpty(item.getTypeName()) && MapUtils.isNotEmpty(item.getUniqueAttributes())) {
-            String          typeName   = item.getTypeName();
-            AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName);
+            final String          queryTemplate = getQueryTemplateForMatchType(context);
+            final String          typeName      = item.getTypeName();
+            final AtlasEntityType entityType    = typeRegistry.getEntityTypeByName(typeName);
 
             if (entityType == null) {
                 throw new AtlasBaseException(AtlasErrorCode.UNKNOWN_TYPENAME, typeName);
             }
 
-            final String queryTemplate;
-            if (StringUtils.equalsIgnoreCase(context.matchType, MATCH_TYPE_STARTS_WITH)) {
-                queryTemplate = gremlinQueryProvider.getQuery(AtlasGremlinQuery.EXPORT_TYPE_STARTS_WITH);
-            } else if (StringUtils.equalsIgnoreCase(context.matchType, MATCH_TYPE_ENDS_WITH)) {
-                queryTemplate = gremlinQueryProvider.getQuery(AtlasGremlinQuery.EXPORT_TYPE_ENDS_WITH);
-            } else if (StringUtils.equalsIgnoreCase(context.matchType, MATCH_TYPE_CONTAINS)) {
-                queryTemplate = gremlinQueryProvider.getQuery(AtlasGremlinQuery.EXPORT_TYPE_CONTAINS);
-            } else if (StringUtils.equalsIgnoreCase(context.matchType, MATCH_TYPE_MATCHES)) {
-                queryTemplate = gremlinQueryProvider.getQuery(AtlasGremlinQuery.EXPORT_TYPE_MATCHES);
-            } else { // default
-                queryTemplate = gremlinQueryProvider.getQuery(AtlasGremlinQuery.EXPORT_TYPE_DEFAULT);
-            }
-
             for (Map.Entry<String, Object> e : item.getUniqueAttributes().entrySet()) {
                 String attrName  = e.getKey();
                 Object attrValue = e.getValue();
 
                 AtlasAttribute attribute = entityType.getAttribute(attrName);
-
                 if (attribute == null || attrValue == null) {
                     continue;
                 }
 
-                context.bindings.clear();
-                context.bindings.put("typeName", typeName);
-                context.bindings.put("attrName", attribute.getQualifiedName());
-                context.bindings.put("attrValue", attrValue);
+                setupBindingsForTypeNameAttrNameAttrValue(context, typeName, attrValue, attribute);
 
                 List<String> guids = executeGremlinQueryForGuids(queryTemplate, context);
 
                 if (CollectionUtils.isNotEmpty(guids)) {
+                    if (ret == null) {
+                        ret = new ArrayList<>();
+                    }
+
                     for (String guid : guids) {
-                        AtlasEntityWithExtInfo entityWithExtInfo = entityGraphRetriever.toAtlasEntityWithExtInfo(guid);
+                        if (!ret.contains(guid)) {
+                            ret.add(guid);
+                        }
+                    }
+                }
+            }
+        }
 
-                        if (entityWithExtInfo == null) {
-                            continue;
+        if (ret == null) {
+            ret = Collections.emptyList();
         }
 
-                        ret.add(entityWithExtInfo);
+        logInfoStartingEntitiesFound(item, context, ret);
+        return ret;
     }
+
+    private void logInfoStartingEntitiesFound(AtlasObjectId item, ExportContext context, List<String> ret) {
+        LOG.info("export(item={}; matchType={}, fetchType={}): found {} entities", item, context.matchType, context.fetchType, ret.size());
     }
 
-                break;
+    private void setupBindingsForTypeName(ExportContext context, String typeName) {
+        context.bindings.clear();
+        context.bindings.put("typeName", new HashSet<String>(Arrays.asList(StringUtils.split(typeName,","))));
     }
 
-            LOG.info("export(item={}; matchType={}, fetchType={}): found {} entities", item, context.matchType, context.fetchType, ret.size());
+    private void setupBindingsForTypeNameAttrNameAttrValue(ExportContext context,
+                                                           String typeName, Object attrValue, AtlasAttribute attribute) {
+        context.bindings.clear();
+        context.bindings.put("typeName", typeName);
+        context.bindings.put("attrName", attribute.getQualifiedName());
+        context.bindings.put("attrValue", attrValue);
     }
 
-        return ret;
+    private String getQueryTemplateForMatchType(ExportContext context) {
+        if (StringUtils.equalsIgnoreCase(context.matchType, MATCH_TYPE_STARTS_WITH)) {
+            return gremlinQueryProvider.getQuery(AtlasGremlinQuery.EXPORT_TYPE_STARTS_WITH);
+        }
+
+        if (StringUtils.equalsIgnoreCase(context.matchType, MATCH_TYPE_ENDS_WITH)) {
+            return gremlinQueryProvider.getQuery(AtlasGremlinQuery.EXPORT_TYPE_ENDS_WITH);
+        }
+
+        if (StringUtils.equalsIgnoreCase(context.matchType, MATCH_TYPE_CONTAINS)) {
+            return gremlinQueryProvider.getQuery(AtlasGremlinQuery.EXPORT_TYPE_CONTAINS);
+        }
+
+        if (StringUtils.equalsIgnoreCase(context.matchType, MATCH_TYPE_MATCHES)) {
+            return gremlinQueryProvider.getQuery(AtlasGremlinQuery.EXPORT_TYPE_MATCHES);
+        }
+
+        if (StringUtils.equalsIgnoreCase(context.matchType, MATCH_TYPE_FOR_TYPE)) {
+            return gremlinQueryProvider.getQuery(AtlasGremlinQuery.EXPORT_TYPE_ALL_FOR_TYPE);
+        }
+
+        return gremlinQueryProvider.getQuery(AtlasGremlinQuery.EXPORT_TYPE_DEFAULT);
     }
 
     private void processEntity(String guid, ExportContext context) throws AtlasBaseException {

repository/src/main/java/org/apache/atlas/util/AtlasGremlin2QueryProvider.java

@@ -43,6 +43,8 @@ public class AtlasGremlin2QueryProvider extends AtlasGremlinQueryProvider {
                 return "g.V('__guid', startGuid).inE().outV().has('__guid').transform{[__guid:it.__guid,isProcess:(it.__superTypeNames != null) ? it.__superTypeNames.contains('Process') : false ]}.dedup().toList()";
             case EXPORT_BY_GUID_CONNECTED_OUT_EDGE:
                 return "g.V('__guid', startGuid).outE().inV().has('__guid').transform{[__guid:it.__guid,isProcess:(it.__superTypeNames != null) ? it.__superTypeNames.contains('Process') : false ]}.dedup().toList()";
+            case EXPORT_TYPE_ALL_FOR_TYPE:
+                return "g.V().has('__typeName',T.in,typeName).has('__guid').__guid.toList()";
             case EXPORT_TYPE_STARTS_WITH:
                 return "g.V().has('__typeName',typeName).filter({it.getProperty(attrName).startsWith(attrValue)}).has('__guid').__guid.toList()";
             case EXPORT_TYPE_ENDS_WITH:

repository/src/main/java/org/apache/atlas/util/AtlasGremlin3QueryProvider.java

@@ -43,6 +43,8 @@ public class AtlasGremlin3QueryProvider extends AtlasGremlin2QueryProvider {
                 return "g.V().has('__guid', startGuid).inE().outV().has('__guid').project('__guid', 'isProcess').by('__guid').by(map {it.get().values('__superTypeNames').toSet().contains('Process')}).dedup().toList()";
             case EXPORT_BY_GUID_CONNECTED_OUT_EDGE:
                 return "g.V().has('__guid', startGuid).outE().inV().has('__guid').project('__guid', 'isProcess').by('__guid').by(map {it.get().values('__superTypeNames').toSet().contains('Process')}).dedup().toList()";
+            case EXPORT_TYPE_ALL_FOR_TYPE:
+                return "g.V().has('__typeName', within(typeName)).has('__guid').values('__guid').toList()";
             case FULL_LINEAGE:
                 return "g.V().has('__guid', '%s').repeat(__.inE('%s').as('e1').outV().outE('%s').as('e2').inV()).emit().select('e1', 'e2').toList()";
             case PARTIAL_LINEAGE:

repository/src/main/java/org/apache/atlas/util/AtlasGremlinQueryProvider.java

@@ -45,6 +45,7 @@ public abstract class AtlasGremlinQueryProvider {
         EXPORT_BY_GUID_FULL,
         EXPORT_BY_GUID_CONNECTED_IN_EDGE,
         EXPORT_BY_GUID_CONNECTED_OUT_EDGE,
+        EXPORT_TYPE_ALL_FOR_TYPE,
         EXPORT_TYPE_STARTS_WITH,
         EXPORT_TYPE_ENDS_WITH,
         EXPORT_TYPE_CONTAINS,

repository/src/test/java/org/apache/atlas/repository/impexp/ExportServiceTest.java

@@ -37,6 +37,7 @@ import org.apache.atlas.repository.store.graph.v1.EntityGraphMapper;
 import org.apache.atlas.repository.store.graph.v1.SoftDeleteHandlerV1;
 import org.apache.atlas.runner.LocalSolrRunner;
 import org.apache.atlas.store.AtlasTypeDefStore;
+import org.apache.atlas.type.AtlasType;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -189,7 +190,7 @@ public class ExportServiceTest {
     }
 
     @Test
-    public void exportType_Succeeds() throws AtlasBaseException, FileNotFoundException {
+    public void exportType() throws AtlasBaseException {
         String requestingIP = "1.0.0.0";
         String hostName = "root";
 
@@ -266,9 +267,7 @@ public class ExportServiceTest {
     }
 
     @Test
-    public void verifyOverallStatus() throws Exception {
-
-//        ExportService service = new ExportService(typeRegistry);
+    public void verifyOverallStatus() {
         assertEquals(AtlasExportResult.OperationStatus.FAIL, exportService.getOverallOperationStatus());
 
         assertEquals(AtlasExportResult.OperationStatus.SUCCESS, exportService.getOverallOperationStatus(AtlasExportResult.OperationStatus.SUCCESS));
@@ -305,6 +304,26 @@ public class ExportServiceTest {
         assertEquals(AtlasExportResult.OperationStatus.FAIL, zipSource.getExportResult().getOperationStatus());
     }
 
+    @Test
+    public void requestForTypeFull() {
+        AtlasExportRequest req = getRequestForTypeFull("Department,Employee");
+
+        assertNotNull(req);
+        assertEquals(req.getItemsToExport().size(), 1);
+        assertEquals(req.getOptions().get(AtlasExportRequest.OPTION_ATTR_MATCH_TYPE), "forType");
+    }
+
+    @Test
+    public void verifyTypeFull() throws AtlasBaseException, IOException {
+        ZipSource zipSource = runExportWithParameters(getRequestForTypeFull("Department,Employee,Manager"));
+        verifyExportForFullEmployeeData(zipSource);
+    }
+
+    private AtlasExportRequest getRequestForTypeFull(String type) {
+        String jsonRequest = "{ \"itemsToExport\": [ { \"typeName\": \"%s\" } ], \"options\": {  \"fetchType\": \"FULL\", \"matchType\": \"forType\"} }";
+        return AtlasType.fromJson(String.format(jsonRequest, type), AtlasExportRequest.class);
+    }
+
     private void tamperEmployeeRequest(AtlasExportRequest request) {
         AtlasObjectId objectId = request.getItemsToExport().get(0);
         objectId.getUniqueAttributes().remove("name");
@@ -329,6 +348,23 @@ public class ExportServiceTest {
         verifyTypeDefs(zipSource);
     }
 
+    private void verifyExportForFullEmployeeData(ZipSource zipSource) throws AtlasBaseException {
+        final List<String> expectedEntityTypes = Arrays.asList(new String[]{"Manager", "Employee", "Department"});
+
+        assertNotNull(zipSource.getCreationOrder());
+        assertTrue(zipSource.hasNext());
+
+        while (zipSource.hasNext()) {
+            AtlasEntity entity = zipSource.next();
+
+            assertNotNull(entity);
+            assertEquals(AtlasEntity.Status.ACTIVE, entity.getStatus());
+            assertTrue(expectedEntityTypes.contains(entity.getTypeName()));
+        }
+
+        verifyTypeDefs(zipSource);
+    }
+
     private void verifyExportForHrData(ZipSource zipSource) throws IOException, AtlasBaseException {
         assertNotNull(zipSource.getCreationOrder());
         assertTrue(zipSource.getCreationOrder().size() == 1);