ATLAS-1950: Import API: Improvement: Specify Supertypes in Import Transforms

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

@@ -17,6 +17,7 @@
  */
 package org.apache.atlas.repository.impexp;
 
+import com.google.common.annotations.VisibleForTesting;
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.impexp.AtlasImportRequest;
@@ -24,6 +25,7 @@ import org.apache.atlas.model.impexp.AtlasImportResult;
 import org.apache.atlas.model.typedef.AtlasTypesDef;
 import org.apache.atlas.repository.store.graph.BulkImporter;
 import org.apache.atlas.store.AtlasTypeDefStore;
+import org.apache.atlas.type.AtlasEntityType;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.commons.collections.MapUtils;
 import org.apache.commons.io.FileUtils;
@@ -75,7 +77,7 @@ public class ImportService {
 
             String transforms = MapUtils.isNotEmpty(request.getOptions()) ? request.getOptions().get(AtlasImportRequest.TRANSFORMS_KEY) : null;
 
-            source.setImportTransform(ImportTransforms.fromJson(transforms));
+            setImportTransform(source, transforms);
             startTimestamp = System.currentTimeMillis();
             processTypes(source.getTypesDef(), result);
             setStartPosition(request, source);
@@ -99,6 +101,30 @@ public class ImportService {
         return result;
     }
 
+    @VisibleForTesting
+    void setImportTransform(ZipSource source, String transforms) throws AtlasBaseException {
+        ImportTransforms importTransform = ImportTransforms.fromJson(transforms);
+        if (importTransform == null) {
+            return;
+        }
+
+        updateTransformsWithSubTypes(importTransform);
+        source.setImportTransform(importTransform);
+    }
+
+    private void updateTransformsWithSubTypes(ImportTransforms importTransforms) throws AtlasBaseException {
+        String[] transformTypes = importTransforms.getTypes().toArray(new String[importTransforms.getTypes().size()]);
+        for (int i = 0; i < transformTypes.length; i++) {
+            String typeName = transformTypes[i];
+            AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName);
+            if (entityType == null) {
+                continue;
+            }
+
+            importTransforms.addParentTransformsToSubTypes(typeName, entityType.getAllSubTypes());
+        }
+    }
+
     private void setStartPosition(AtlasImportRequest request, ZipSource source) throws AtlasBaseException {
         if (request.getStartGuid() != null) {
             source.setPositionUsingEntityGuid(request.getStartGuid());

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

@@ -30,6 +30,7 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 public class ImportTransforms {
     private static final Logger LOG = LoggerFactory.getLogger(ImportTransforms.class);
@@ -53,6 +54,23 @@ public class ImportTransforms {
 
     public Map<String, List<ImportTransformer>> getTransforms(String typeName) { return transforms.get(typeName); }
 
+    public Set<String> getTypes() {
+        return getTransforms().keySet();
+    }
+
+    public void addParentTransformsToSubTypes(String parentType, Set<String> subTypes) {
+        Map<String, List<ImportTransformer>> attribtueTransformMap = getTransforms().get(parentType);
+        for (String subType : subTypes) {
+            if(!getTransforms().containsKey(subType)) {
+                getTransforms().put(subType, attribtueTransformMap);
+            } else {
+                for (Map.Entry<String, List<ImportTransformer>> entry : attribtueTransformMap.entrySet()) {
+                    getTransforms().get(subType).get(entry.getKey()).addAll(entry.getValue());
+                }
+            }
+        }
+    }
+
     public AtlasEntity.AtlasEntityWithExtInfo apply(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo) throws AtlasBaseException {
         if (entityWithExtInfo != null) {
             apply(entityWithExtInfo.getEntity());

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

@@ -54,7 +54,10 @@ import static org.apache.atlas.graph.GraphSandboxUtil.useLocalSolr;
 import static org.apache.atlas.repository.impexp.ZipFileResourceTestUtils.*;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
-import static org.testng.Assert.*;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.assertFalse;
 
 @Guice(modules = TestModules.TestOnlyModule.class)
 public class ImportServiceTest {
@@ -374,4 +377,33 @@ public class ImportServiceTest {
     private AtlasEntity.AtlasEntityWithExtInfo getEntity(AtlasEntityHeader header) throws AtlasBaseException {
         return entityStore.getById(header.getGuid());
     }
+
+    @Test(dataProvider = "salesNewTypeAttrs-next")
+    public void transformUpdatesForSubTypes(ZipSource zipSource) throws IOException, AtlasBaseException {
+        loadBaseModel();
+        loadHiveModel();
+
+        String transformJSON = "{ \"Asset\": { \"qualifiedName\":[ \"lowercase\", \"replace:@cl1:@cl2\" ] } }";
+        importService.setImportTransform(zipSource, transformJSON);
+        ImportTransforms importTransforms = zipSource.getImportTransform();
+
+        assertTrue(importTransforms.getTransforms().containsKey("Asset"));
+        assertTrue(importTransforms.getTransforms().containsKey("hive_table"));
+        assertTrue(importTransforms.getTransforms().containsKey("hive_column"));
+    }
+
+    @Test(dataProvider = "salesNewTypeAttrs-next")
+    public void transformUpdatesForSubTypesAddsToExistingTransforms(ZipSource zipSource) throws IOException, AtlasBaseException {
+        loadBaseModel();
+        loadHiveModel();
+
+        String transformJSON = "{ \"Asset\": { \"qualifiedName\":[ \"replace:@cl1:@cl2\" ] }, \"hive_table\": { \"qualifiedName\":[ \"lowercase\" ] } }";
+        importService.setImportTransform(zipSource, transformJSON);
+        ImportTransforms importTransforms = zipSource.getImportTransform();
+
+        assertTrue(importTransforms.getTransforms().containsKey("Asset"));
+        assertTrue(importTransforms.getTransforms().containsKey("hive_table"));
+        assertTrue(importTransforms.getTransforms().containsKey("hive_column"));
+        assertEquals(importTransforms.getTransforms().get("hive_table").get("qualifiedName").size(), 2);
+    }
 }

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

@@ -23,19 +23,23 @@ import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.testng.annotations.BeforeTest;
 import org.testng.annotations.Test;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
+import static org.apache.atlas.repository.impexp.ZipFileResourceTestUtils.loadModelFromJson;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
 
 public class ImportTransformsTest {
     private final String qualifiedName  = "qualifiedName";
     private final String lowerCaseCL1   = "@cl1";
     private final String lowerCaseCL2   = "@cl2";
     private final String jsonTransforms = "{ \"hive_table\": { \"qualifiedName\":[ \"lowercase\", \"replace:@cl1:@cl2\" ] } }";
+    private final String jsonTransforms2 = "{ \"Asset\": { \"qualifiedName\":[ \"replace:@cl1:@cl2\" ] }, \"hive_table\": { \"qualifiedName\":[ \"lowercase\", \"replace:@cl1:@cl2\" ] } }";
 
     private ImportTransforms transform;
 
@@ -86,6 +90,14 @@ public class ImportTransformsTest {
         assertEquals(entityWithExtInfo.getEntity().getGuid(), transformedEntityWithExtInfo.getEntity().getGuid());
     }
 
+    @Test
+    public void transformFromJsonWithMultipleEntries() {
+        ImportTransforms t = ImportTransforms.fromJson(jsonTransforms2);
+
+        assertNotNull(t);
+        assertEquals(t.getTransforms().size(), 2);
+    }
+
     private String[] getExtEntityExpectedValues(AtlasEntityWithExtInfo entityWithExtInfo) {
         String[] ret = new String[entityWithExtInfo.getReferredEntities().size()];