From 6d0c5c8c7162f720aabe4c57e1d4bbaf8ddb4d9a Mon Sep 17 00:00:00 2001
From: Ashutosh Mestry <amestry@hortonworks.com>
Date: Tue, 10 Jul 2018 15:37:16 -0700
Subject: [PATCH] ATLAS-1950: Import API: Improvement: Specify Supertypes in Import Transforms
Signed-off-by: Sarath Subramanian <ssubramanian@hortonworks.com>
---
repository/src/main/java/org/apache/atlas/repository/impexp/ImportService.java | 28 +++++++++++++++++++++++++++-
repository/src/main/java/org/apache/atlas/repository/impexp/ImportTransforms.java | 18 ++++++++++++++++++
repository/src/test/java/org/apache/atlas/repository/impexp/ImportServiceTest.java | 34 +++++++++++++++++++++++++++++++++-
repository/src/test/java/org/apache/atlas/repository/impexp/ImportTransformsTest.java | 12 ++++++++++++
4 files changed, 90 insertions(+), 2 deletions(-)
diff --git a/repository/src/main/java/org/apache/atlas/repository/impexp/ImportService.java b/repository/src/main/java/org/apache/atlas/repository/impexp/ImportService.java
index c976c59..972e1e3 100644
--- a/repository/src/main/java/org/apache/atlas/repository/impexp/ImportService.java
+++ b/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());
diff --git a/repository/src/main/java/org/apache/atlas/repository/impexp/ImportTransforms.java b/repository/src/main/java/org/apache/atlas/repository/impexp/ImportTransforms.java
index 63e53c5..2f27448 100644
--- a/repository/src/main/java/org/apache/atlas/repository/impexp/ImportTransforms.java
+++ b/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());
diff --git a/repository/src/test/java/org/apache/atlas/repository/impexp/ImportServiceTest.java b/repository/src/test/java/org/apache/atlas/repository/impexp/ImportServiceTest.java
index de56d3f..a910d08 100644
--- a/repository/src/test/java/org/apache/atlas/repository/impexp/ImportServiceTest.java
+++ b/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);
+ }
}
diff --git a/repository/src/test/java/org/apache/atlas/repository/impexp/ImportTransformsTest.java b/repository/src/test/java/org/apache/atlas/repository/impexp/ImportTransformsTest.java
index ccedeb1..a73abcd 100644
--- a/repository/src/test/java/org/apache/atlas/repository/impexp/ImportTransformsTest.java
+++ b/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()];
--
libgit2 0.27.1