Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,12 @@ public List<String> get(AtlasExportRequest exportRequest, AtlasObjectId item) {

return ret;
}

if (StringUtils.isEmpty(item.getTypeName()) && MapUtils.isNotEmpty(item.getUniqueAttributes())) {
LOG.info("Request missing typeName but has uniqueAttributes. Attempting generic type search.");
return getEntitiesForMatchTypeType(item, MATCH_TYPE_FOR_TYPE);
}

} catch (AtlasBaseException ex) {
LOG.error("Error fetching starting entity for: {}", item, ex);
} finally {
Expand Down Expand Up @@ -145,16 +151,22 @@ List<String> executeGremlinQuery(String query, Map<String, Object> bindings) {
}

private List<String> getEntitiesForMatchTypeUsingUniqueAttributes(AtlasObjectId item, String matchType) throws AtlasBaseException {
final String queryTemplate = getQueryTemplateForMatchType(matchType);
final String typeName = item.getTypeName();
final AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName);

Set<String> ret = new HashSet<>();
final String typeName = item.getTypeName();
final AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName);

if (entityType == null) {
throw new AtlasBaseException(AtlasErrorCode.UNKNOWN_TYPENAME, typeName);
}

final String queryTemplate = getQueryTemplateForMatchType(matchType);
Set<String> ret = new HashSet<>();

Set<String> typeNamesToQuery = new HashSet<>();
typeNamesToQuery.add(typeName);
if (CollectionUtils.isNotEmpty(entityType.getAllSubTypes())) {
typeNamesToQuery.addAll(entityType.getAllSubTypes());
}

for (Map.Entry<String, Object> e : item.getUniqueAttributes().entrySet()) {
String attrName = e.getKey();
Object attrValue = e.getValue();
Expand All @@ -165,27 +177,50 @@ private List<String> getEntitiesForMatchTypeUsingUniqueAttributes(AtlasObjectId
continue;
}

List<String> guids = executeGremlinQuery(queryTemplate, getBindingsForObjectId(typeName, attribute.getQualifiedName(), e.getValue()));
for (String typeToSearch : typeNamesToQuery) {
List<String> guids = executeGremlinQuery(queryTemplate, getBindingsForObjectId(typeToSearch, attribute.getQualifiedName(), e.getValue()));

if (!CollectionUtils.isNotEmpty(guids)) {
continue;
}
if (CollectionUtils.isNotEmpty(guids)) {
ret.addAll(guids);

ret.addAll(guids);
}
}
}

return new ArrayList<>(ret);
}

private List<String> getEntitiesForMatchTypeType(AtlasObjectId item, String matchType) {
private List<String> getEntitiesForMatchTypeType(AtlasObjectId item, String matchType) throws AtlasBaseException {
return executeGremlinQuery(getQueryTemplateForMatchType(matchType), getBindingsForTypeName(item.getTypeName()));
}

private HashMap<String, Object> getBindingsForTypeName(String typeName) {
HashMap<String, Object> ret = new HashMap<>();

ret.put(BINDING_PARAMETER_TYPENAME, new HashSet<>(Arrays.asList(StringUtils.split(typeName, ","))));
private HashMap<String, Object> getBindingsForTypeName(String typeName) throws AtlasBaseException {
HashMap<String, Object> ret = new HashMap<>();
Set<String> typeNamesToQuery = new HashSet<>();

if (StringUtils.isBlank(typeName)) {
typeNamesToQuery.addAll(typeRegistry.getAllEntityDefNames());
} else {
List<String> providedTypeNames = Arrays.asList(StringUtils.split(typeName, ","));

for (String name : providedTypeNames) {
AtlasType type = typeRegistry.getType(name);

if (type instanceof AtlasEntityType) {
AtlasEntityType entityType = (AtlasEntityType) type;
typeNamesToQuery.add(entityType.getTypeName());

Set<String> subTypes = entityType.getAllSubTypes();
if (CollectionUtils.isNotEmpty(subTypes)) {
typeNamesToQuery.addAll(subTypes);
}
} else {
typeNamesToQuery.add(name);
}
}
}

ret.put(BINDING_PARAMETER_TYPENAME, typeNamesToQuery);
return ret;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,14 +35,18 @@
import javax.inject.Inject;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Set;

import static org.apache.atlas.repository.impexp.StartEntityFetchByExportRequest.BINDING_PARAMETER_ATTR_NAME;
import static org.apache.atlas.repository.impexp.StartEntityFetchByExportRequest.BINDING_PARAMETER_TYPENAME;
import static org.apache.atlas.repository.impexp.StartEntityFetchByExportRequest.BINDING_PARAMTER_ATTR_VALUE;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertTrue;

@Guice(modules = TestModules.TestOnlyModule.class)
public class StartEntityFetchByExportRequestTest extends AtlasTestBase {
Expand All @@ -58,6 +62,14 @@ public class StartEntityFetchByExportRequestTest extends AtlasTestBase {
private AtlasGremlin3QueryProvider atlasGremlin3QueryProvider;
private StartEntityFetchByExportRequestSpy startEntityFetchByExportRequestSpy;

@BeforeClass
void setup() throws IOException, AtlasBaseException {
super.basicSetup(typeDefStore, typeRegistry);

atlasGremlin3QueryProvider = new AtlasGremlin3QueryProvider();
startEntityFetchByExportRequestSpy = new StartEntityFetchByExportRequestSpy(atlasGraph, typeRegistry);
}

@Test
public void fetchTypeGuid() {
String exportRequestJson = "{ \"itemsToExport\": [ { \"typeName\": \"hive_db\", \"guid\": \"111-222-333\" } ]}";
Expand All @@ -72,44 +84,117 @@ public void fetchTypeUniqueAttributes() {
String exportRequestJson = "{ \"itemsToExport\": [ { \"typeName\": \"hive_db\", \"uniqueAttributes\": {\"qualifiedName\": \"stocks@cl1\"} } ]}";
AtlasExportRequest exportRequest = AtlasType.fromJson(exportRequestJson, AtlasExportRequest.class);

startEntityFetchByExportRequestSpy.clearRecordedCalls();
startEntityFetchByExportRequestSpy.get(exportRequest);

assertEquals(startEntityFetchByExportRequestSpy.getGeneratedQuery(), startEntityFetchByExportRequestSpy.getQueryTemplateForMatchType(exportRequest.getMatchTypeOptionValue()));
assertEquals(startEntityFetchByExportRequestSpy.getSuppliedBindingsMap().get(BINDING_PARAMETER_TYPENAME), "hive_db");
assertEquals(startEntityFetchByExportRequestSpy.getSuppliedBindingsMap().get(BINDING_PARAMETER_ATTR_NAME), "Referenceable.qualifiedName");
assertEquals(startEntityFetchByExportRequestSpy.getSuppliedBindingsMap().get(BINDING_PARAMTER_ATTR_VALUE), "stocks@cl1");
assertFalse(startEntityFetchByExportRequestSpy.getRecordedBindings().isEmpty());

boolean foundHiveDbSearch = startEntityFetchByExportRequestSpy.getRecordedBindings().stream()
.anyMatch(b -> b.get(BINDING_PARAMETER_TYPENAME).equals("hive_db") &&
b.get(BINDING_PARAMETER_ATTR_NAME).equals("Referenceable.qualifiedName") &&
b.get(BINDING_PARAMTER_ATTR_VALUE).equals("stocks@cl1"));

assertTrue(foundHiveDbSearch, "Should have searched for hive_db specifically");
}

@BeforeClass
void setup() throws IOException, AtlasBaseException {
super.basicSetup(typeDefStore, typeRegistry);
@Test
public void fetchReferenceableUniqueAttributes() {
String exportRequestJson = "{ \"itemsToExport\": [ { \"typeName\": \"Referenceable\", \"uniqueAttributes\": {\"qualifiedName\": \"stocks@cl1\"} } ]}";
AtlasExportRequest exportRequest = AtlasType.fromJson(exportRequestJson, AtlasExportRequest.class);

atlasGremlin3QueryProvider = new AtlasGremlin3QueryProvider();
startEntityFetchByExportRequestSpy = new StartEntityFetchByExportRequestSpy(atlasGraph, typeRegistry);
startEntityFetchByExportRequestSpy.clearRecordedCalls();
startEntityFetchByExportRequestSpy.get(exportRequest);

List<Map<String, Object>> allBindings = startEntityFetchByExportRequestSpy.getRecordedBindings();

assertTrue(allBindings.size() > 1, "Should have triggered multiple queries for Referenceable subtypes");

boolean foundHiveTable = allBindings.stream()
.anyMatch(b -> b.get(BINDING_PARAMETER_TYPENAME).equals("hive_table"));

assertTrue(foundHiveTable, "Search should have included subtype: hive_table");
}

@Test
public void fetchTypeExpansion() {
String exportRequestJson = "{ \"itemsToExport\": [ { \"typeName\": \"Asset\" } ], \"options\": {\"matchType\": \"forType\"} }";
AtlasExportRequest exportRequest = AtlasType.fromJson(exportRequestJson, AtlasExportRequest.class);

startEntityFetchByExportRequestSpy.clearRecordedCalls();
startEntityFetchByExportRequestSpy.get(exportRequest);

Map<String, Object> lastBindings = startEntityFetchByExportRequestSpy.getSuppliedBindingsMap();
Set<String> boundTypes = (Set<String>) lastBindings.get(BINDING_PARAMETER_TYPENAME);

assertTrue(boundTypes.contains("Asset"));
assertTrue(boundTypes.contains("hive_db")); // Asset is a supertype of hive_db
}

@Test
public void fetchEmptyTypeUniqueAttributes() throws AtlasBaseException {
String exportRequestJson = "{ \"itemsToExport\": [ { \"typeName\": \"\", \"uniqueAttributes\": {\"qualifiedName\": \"stocks@cl1\"} } ]}";
AtlasExportRequest exportRequest = AtlasType.fromJson(exportRequestJson, AtlasExportRequest.class);

startEntityFetchByExportRequestSpy.clearRecordedCalls();
startEntityFetchByExportRequestSpy.get(exportRequest);

Map<String, Object> lastBindings = startEntityFetchByExportRequestSpy.getSuppliedBindingsMap();

assertTrue(lastBindings != null && lastBindings.containsKey(BINDING_PARAMETER_TYPENAME), "Bindings should contain typeName");

Set<String> boundTypes = (Set<String>) lastBindings.get(BINDING_PARAMETER_TYPENAME);
int expectedSize = typeRegistry.getAllEntityDefNames().size();

assertEquals(boundTypes.size(), expectedSize, "Should contain all entity types from registry");
}

@Test
public void fetchUnknownType() {
String exportRequestJson = "{ \"itemsToExport\": [ { \"typeName\": \"InvalidType\" } ]}";
AtlasExportRequest exportRequest = AtlasType.fromJson(exportRequestJson, AtlasExportRequest.class);

List<AtlasObjectId> result = startEntityFetchByExportRequestSpy.get(exportRequest);
assertTrue(result.isEmpty());
}

/**
* Spy class to capture Gremlin queries and bindings.
* to record multiple calls because the new implementation uses loops.
*/
private class StartEntityFetchByExportRequestSpy extends StartEntityFetchByExportRequest {
String generatedQuery;
Map<String, Object> suppliedBindingsMap;
private String lastQuery;
private Map<String, Object> lastBindings;
private List<Map<String, Object>> recordedBindings = new ArrayList<>();

public StartEntityFetchByExportRequestSpy(AtlasGraph atlasGraph, AtlasTypeRegistry typeRegistry) {
super(atlasGraph, typeRegistry, atlasGremlin3QueryProvider);
}

public String getGeneratedQuery() {
return generatedQuery;
return lastQuery;
}

public Map<String, Object> getSuppliedBindingsMap() {
return suppliedBindingsMap;
return lastBindings;
}

public List<Map<String, Object>> getRecordedBindings() {
return recordedBindings;
}

public void clearRecordedCalls() {
recordedBindings.clear();
lastQuery = null;
lastBindings = null;
}

@Override
List<String> executeGremlinQuery(String query, Map<String, Object> bindings) {
this.generatedQuery = query;
this.suppliedBindingsMap = bindings;
this.lastQuery = query;
this.lastBindings = bindings;
this.recordedBindings.add(new java.util.HashMap<>(bindings));

return Collections.emptyList();
}
}
}
}