ATLAS-1618: updated export to support scope option - full/connected

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

@@ -18,12 +18,6 @@
 package org.apache.atlas.model.impexp;
 
 
-import java.io.Serializable;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
 import org.codehaus.jackson.annotate.JsonAutoDetect;
@@ -33,6 +27,11 @@ import org.codehaus.jackson.map.annotate.JsonSerialize;
 import javax.xml.bind.annotation.XmlAccessType;
 import javax.xml.bind.annotation.XmlAccessorType;
 import javax.xml.bind.annotation.XmlRootElement;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 
 import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.NONE;
 import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.PUBLIC_ONLY;
@@ -44,12 +43,17 @@ import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.PUBLIC_ONL
 @XmlRootElement
 @XmlAccessorType(XmlAccessType.PROPERTY)
 public class AtlasExportRequest implements Serializable {
-    private static final long serialVersionUID = 1L;
 
-    public static final String EXPORT_PARAM_OPTION             = "FORMAT";
-    public static final String EXPORT_PARAM_OPTION_FORMAT_JSON = "JSON";
-    public static final String EXPORT_PARAM_OPTION_FORMAT_ZIP  = "ZIP";
+    private static final long serialVersionUID = 1L;
 
+    public static final String OPTION_FETCH_TYPE      = "fetchType";
+    public static final String OPTION_ATTR_MATCH_TYPE = "matchType";
+    public static final String FETCH_TYPE_FULL        = "full";
+    public static final String FETCH_TYPE_CONNECTED   = "connected";
+    public static final String MATCH_TYPE_STARTS_WITH = "startsWith";
+    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";
 
     private List<AtlasObjectId> itemsToExport = new ArrayList<>();
     private Map<String, Object> options       = new HashMap<>();

repository/src/main/java/org/apache/atlas/repository/store/graph/AtlasEntityStore.java

@@ -23,9 +23,10 @@ import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.impexp.AtlasImportResult;
 import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasEntity;
-import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.instance.EntityMutationResponse;
+import org.apache.atlas.repository.store.graph.v1.EntityImportStream;
 import org.apache.atlas.repository.store.graph.v1.EntityStream;
 import org.apache.atlas.type.AtlasEntityType;
 
@@ -76,7 +77,7 @@ public interface AtlasEntityStore {
      * @return EntityMutationResponse Entity mutations operations with the corresponding set of entities on which these operations were performed
      * @throws AtlasBaseException
      */
-    EntityMutationResponse bulkImport(EntityStream entityStream, AtlasImportResult importResult) throws AtlasBaseException;
+    EntityMutationResponse bulkImport(EntityImportStream entityStream, AtlasImportResult importResult) throws AtlasBaseException;
 
     /**
      * Update a single entity

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1.java

@@ -142,7 +142,7 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
     }
 
     @Override
-    public EntityMutationResponse bulkImport(EntityStream entityStream, AtlasImportResult importResult) throws AtlasBaseException {
+    public EntityMutationResponse bulkImport(EntityImportStream entityStream, AtlasImportResult importResult) throws AtlasBaseException {
         if (LOG.isDebugEnabled()) {
             LOG.debug("==> bulkImport()");
         }
@@ -168,6 +168,10 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
 
             EntityMutationResponse resp = createOrUpdate(oneEntityStream, false, true);
 
+            if(CollectionUtils.isNotEmpty(entity.getClassifications())) {
+                addClassifications(entity.getGuid(), entity.getClassifications());
+            }
+
             updateImportMetrics("entity:%s:created", resp.getCreatedEntities(), processedGuids, importResult);
             updateImportMetrics("entity:%s:updated", resp.getUpdatedEntities(), processedGuids, importResult);
             updateImportMetrics("entity:%s:deleted", resp.getDeletedEntities(), processedGuids, importResult);
@@ -181,6 +185,8 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
             if (resp.getGuidAssignments() != null) {
                 ret.getGuidAssignments().putAll(resp.getGuidAssignments());
             }
+
+            entityStream.onImportComplete(entity.getGuid());
         }
 
         importResult.getProcessedEntities().addAll(processedGuids);

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStreamForImport.java

@@ -18,6 +18,9 @@
 package org.apache.atlas.repository.store.graph.v1;
 
 import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.instance.AtlasEntityHeader;
+
+import java.util.List;
 
 public class AtlasEntityStreamForImport extends AtlasEntityStream implements EntityImportStream {
     public AtlasEntityStreamForImport(AtlasEntity entity) {
@@ -27,4 +30,9 @@ public class AtlasEntityStreamForImport extends AtlasEntityStream implements Ent
     public AtlasEntityStreamForImport(AtlasEntity entity, EntityStream entityStream) {
         super(entity, entityStream);
     }
+
+    @Override
+    public void onImportComplete(String guid) {
+
+    }
 }

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityImportStream.java

@@ -17,5 +17,8 @@
  */
 package org.apache.atlas.repository.store.graph.v1;
 
+
 public interface EntityImportStream extends EntityStream {
+
+    void onImportComplete(String guid);
 }

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

@@ -37,8 +37,12 @@ public class AtlasGremlin2QueryProvider extends AtlasGremlinQueryProvider {
                 return "g.V().has('__superTypeNames', T.in, ['Referenceable']).has('__traitNames').count()";
             case ENTITIES_FOR_TAG_METRIC:
                 return "g.V().has('__typeName', T.in, g.V().has('__type', 'typeSystem').filter{it.'__type.category'.name() == 'TRAIT'}.'__type.name'.toSet()).groupCount{it.'__typeName'}.cap.toList()";
-            case EXPORT_BY_GUID:
+            case EXPORT_BY_GUID_FULL:
                 return "g.V('__guid', startGuid).bothE().bothV().has('__guid').__guid.dedup().toList()";
+            case EXPORT_BY_GUID_CONNECTED_IN_EDGE:
+                return "g.V('__guid', startGuid).inE().outV().has('__guid').__guid.dedup().toList()";
+            case EXPORT_BY_GUID_CONNECTED_OUT_EDGE:
+                return "g.V('__guid', startGuid).outE().inV().has('__guid').__guid.dedup().toList()";
             case EXPORT_TYPE_STARTS_WITH:
                 return "g.V().has('__typeName','%s').filter({it.'%s'.startsWith(attrValue)}).has('__guid').__guid.toList()";
             case EXPORT_TYPE_ENDS_WITH:

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

@@ -42,7 +42,9 @@ public abstract class AtlasGremlinQueryProvider {
         ENTITIES_FOR_TAG_METRIC,
 
         // Import Export related Queries
-        EXPORT_BY_GUID,
+        EXPORT_BY_GUID_FULL,
+        EXPORT_BY_GUID_CONNECTED_IN_EDGE,
+        EXPORT_BY_GUID_CONNECTED_OUT_EDGE,
         EXPORT_TYPE_STARTS_WITH,
         EXPORT_TYPE_ENDS_WITH,
         EXPORT_TYPE_CONTAINS,

typesystem/src/main/java/org/apache/atlas/typesystem/types/TypeUtils.java

@@ -22,12 +22,7 @@ import com.google.common.collect.ImmutableList;
 import org.apache.atlas.AtlasException;
 
 import java.io.IOException;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -82,6 +77,22 @@ public class TypeUtils {
         public static <L, R> Pair<L, R> of(L left, R right) {
             return new Pair<>(left, right);
         }
+
+        public boolean equals(Object o) {
+            if (this == o) {
+                return true;
+            }
+
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+
+            Pair p = (Pair)o;
+
+            return Objects.equals(left, p.left) && Objects.equals(right, p.right);
+        }
+
+        public int hashCode() { return Objects.hash(left, right); }
     }
 
     /**

webapp/src/main/java/org/apache/atlas/web/resources/AdminResource.java

@@ -22,7 +22,9 @@ import com.google.inject.Inject;
 import org.apache.atlas.ApplicationProperties;
 import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasErrorCode;
-import org.apache.atlas.AtlasException;
+import org.apache.atlas.authorize.AtlasActionTypes;
+import org.apache.atlas.authorize.AtlasResourceTypes;
+import org.apache.atlas.authorize.simple.AtlasAuthorizationUtils;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.impexp.AtlasExportRequest;
 import org.apache.atlas.model.impexp.AtlasExportResult;
@@ -31,14 +33,12 @@ import org.apache.atlas.model.impexp.AtlasImportResult;
 import org.apache.atlas.model.metrics.AtlasMetrics;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
 import org.apache.atlas.services.MetricsService;
-import org.apache.atlas.authorize.AtlasActionTypes;
-import org.apache.atlas.authorize.AtlasResourceTypes;
-import org.apache.atlas.authorize.simple.AtlasAuthorizationUtils;
 import org.apache.atlas.store.AtlasTypeDefStore;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.atlas.web.filters.AtlasCSRFPreventionFilter;
 import org.apache.atlas.web.service.ServiceState;
 import org.apache.atlas.web.util.Servlets;
+import org.apache.commons.configuration.Configuration;
 import org.apache.commons.configuration.ConfigurationException;
 import org.apache.commons.configuration.PropertiesConfiguration;
 import org.apache.commons.lang.StringUtils;
@@ -54,13 +54,7 @@ import javax.inject.Singleton;
 import javax.servlet.ServletOutputStream;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
-import javax.ws.rs.Consumes;
-import javax.ws.rs.GET;
-import javax.ws.rs.POST;
-import javax.ws.rs.Path;
-import javax.ws.rs.Produces;
-import javax.ws.rs.QueryParam;
-import javax.ws.rs.WebApplicationException;
+import javax.ws.rs.*;
 import javax.ws.rs.core.Context;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
@@ -70,9 +64,6 @@ import java.util.Collection;
 import java.util.HashSet;
 import java.util.Set;
 import java.util.concurrent.locks.ReentrantLock;
-import org.apache.commons.configuration.Configuration;
-
-import static org.apache.atlas.repository.converters.AtlasInstanceConverter.toAtlasBaseException;
 
 
 /**
@@ -369,7 +360,7 @@ public class AdminResource {
         try {
             AtlasImportRequest   request       = new AtlasImportRequest(Servlets.getParameterMap(httpServletRequest));
             ByteArrayInputStream inputStream   = new ByteArrayInputStream(bytes);
-            ImportService        importService = new ImportService(this.typesDefStore, this.entityStore);
+            ImportService        importService = new ImportService(this.typesDefStore, this.entityStore, this.typeRegistry);
 
             ZipSource zipSource = new ZipSource(inputStream);
 
@@ -405,7 +396,7 @@ public class AdminResource {
 
         try {
             AtlasImportRequest request       = new AtlasImportRequest(Servlets.getParameterMap(httpServletRequest));
-            ImportService      importService = new ImportService(this.typesDefStore, this.entityStore);
+            ImportService      importService = new ImportService(this.typesDefStore, this.entityStore, this.typeRegistry);
 
             result = importService.run(request, Servlets.getUserName(httpServletRequest),
                                        Servlets.getHostName(httpServletRequest),

webapp/src/main/java/org/apache/atlas/web/resources/ImportService.java

@@ -18,21 +18,22 @@
 package org.apache.atlas.web.resources;
 
 import org.apache.atlas.AtlasErrorCode;
+import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.impexp.AtlasImportRequest;
 import org.apache.atlas.model.impexp.AtlasImportResult;
-import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.typedef.*;
+import org.apache.atlas.repository.store.bootstrap.AtlasTypeDefStoreInitializer;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
-import org.apache.commons.io.FileUtils;
-import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.store.AtlasTypeDefStore;
+import org.apache.atlas.type.AtlasTypeRegistry;
+import org.apache.commons.io.FileUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.*;
-import java.util.List;
-import java.util.concurrent.TimeUnit;
+import java.io.ByteArrayInputStream;
+import java.io.File;
+import java.io.FileNotFoundException;
 
 
 public class ImportService {
@@ -40,14 +41,16 @@ public class ImportService {
 
     private final AtlasTypeDefStore typeDefStore;
     private final AtlasEntityStore  entityStore;
+    private final AtlasTypeRegistry typeRegistry;
 
     private long startTimestamp;
     private long endTimestamp;
 
 
-    public ImportService(final AtlasTypeDefStore typeDefStore, final AtlasEntityStore entityStore) {
+    public ImportService(final AtlasTypeDefStore typeDefStore, final AtlasEntityStore entityStore, AtlasTypeRegistry typeRegistry) {
         this.typeDefStore = typeDefStore;
         this.entityStore  = entityStore;
+        this.typeRegistry = typeRegistry;
     }
 
     public AtlasImportResult run(ZipSource source, AtlasImportRequest request, String userName,
@@ -116,9 +119,13 @@ public class ImportService {
     }
 
     private void processTypes(AtlasTypesDef typeDefinitionMap, AtlasImportResult result) throws AtlasBaseException {
-        setGuidToEmpty(typeDefinitionMap.getEntityDefs());
-        typeDefStore.updateTypesDef(typeDefinitionMap);
+        setGuidToEmpty(typeDefinitionMap);
+        AtlasTypesDef typesToCreate = AtlasTypeDefStoreInitializer.getTypesToCreate(typeDefinitionMap, this.typeRegistry);
+        if (!typesToCreate.isEmpty()) {
+            typeDefStore.createTypesDef(typesToCreate);
+        }
 
+        typeDefStore.updateTypesDef(typeDefinitionMap);
         updateMetricsForTypesDef(typeDefinitionMap, result);
     }
 
@@ -129,9 +136,21 @@ public class ImportService {
         result.incrementMeticsCounter("typedef:struct", typeDefinitionMap.getStructDefs().size());
     }
 
-    private void setGuidToEmpty(List<AtlasEntityDef> entityDefList) {
-        for (AtlasEntityDef edf: entityDefList) {
-            edf.setGuid("");
+    private void setGuidToEmpty(AtlasTypesDef typesDef) {
+        for (AtlasEntityDef def: typesDef.getEntityDefs()) {
+            def.setGuid(null);
+        }
+
+        for (AtlasClassificationDef def: typesDef.getClassificationDefs()) {
+            def.setGuid(null);
+        }
+
+        for (AtlasEnumDef def: typesDef.getEnumDefs()) {
+            def.setGuid(null);
+        }
+
+        for (AtlasStructDef def: typesDef.getStructDefs()) {
+            def.setGuid(null);
         }
     }
 

webapp/src/main/java/org/apache/atlas/web/resources/ZipSource.java

@@ -17,6 +17,7 @@
  */
 package org.apache.atlas.web.resources;
 
+import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.codehaus.jackson.type.TypeReference;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.instance.AtlasEntity;
@@ -184,4 +185,11 @@ public class ZipSource implements EntityImportStream {
             return null;
         }
     }
+
+    @Override
+    public void onImportComplete(String guid) {
+        if(guid != null) {
+            guidEntityJsonMap.remove(guid);
+        }
+    }
 }