diff --git a/repository/src/main/java/org/apache/atlas/repository/impexp/StartEntityFetchByExportRequest.java b/repository/src/main/java/org/apache/atlas/repository/impexp/StartEntityFetchByExportRequest.java index da19f19747d..13bf2f1467b 100644 --- a/repository/src/main/java/org/apache/atlas/repository/impexp/StartEntityFetchByExportRequest.java +++ b/repository/src/main/java/org/apache/atlas/repository/impexp/StartEntityFetchByExportRequest.java @@ -109,6 +109,12 @@ public List 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 { @@ -145,16 +151,22 @@ List executeGremlinQuery(String query, Map bindings) { } private List getEntitiesForMatchTypeUsingUniqueAttributes(AtlasObjectId item, String matchType) throws AtlasBaseException { - final String queryTemplate = getQueryTemplateForMatchType(matchType); - final String typeName = item.getTypeName(); - final AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName); - - Set 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 ret = new HashSet<>(); + + Set typeNamesToQuery = new HashSet<>(); + typeNamesToQuery.add(typeName); + if (CollectionUtils.isNotEmpty(entityType.getAllSubTypes())) { + typeNamesToQuery.addAll(entityType.getAllSubTypes()); + } + for (Map.Entry e : item.getUniqueAttributes().entrySet()) { String attrName = e.getKey(); Object attrValue = e.getValue(); @@ -165,27 +177,50 @@ private List getEntitiesForMatchTypeUsingUniqueAttributes(AtlasObjectId continue; } - List guids = executeGremlinQuery(queryTemplate, getBindingsForObjectId(typeName, attribute.getQualifiedName(), e.getValue())); + for (String typeToSearch : typeNamesToQuery) { + List 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 getEntitiesForMatchTypeType(AtlasObjectId item, String matchType) { + private List getEntitiesForMatchTypeType(AtlasObjectId item, String matchType) throws AtlasBaseException { return executeGremlinQuery(getQueryTemplateForMatchType(matchType), getBindingsForTypeName(item.getTypeName())); } - private HashMap getBindingsForTypeName(String typeName) { - HashMap ret = new HashMap<>(); - - ret.put(BINDING_PARAMETER_TYPENAME, new HashSet<>(Arrays.asList(StringUtils.split(typeName, ",")))); + private HashMap getBindingsForTypeName(String typeName) throws AtlasBaseException { + HashMap ret = new HashMap<>(); + Set typeNamesToQuery = new HashSet<>(); + + if (StringUtils.isBlank(typeName)) { + typeNamesToQuery.addAll(typeRegistry.getAllEntityDefNames()); + } else { + List 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 subTypes = entityType.getAllSubTypes(); + if (CollectionUtils.isNotEmpty(subTypes)) { + typeNamesToQuery.addAll(subTypes); + } + } else { + typeNamesToQuery.add(name); + } + } + } + ret.put(BINDING_PARAMETER_TYPENAME, typeNamesToQuery); return ret; } diff --git a/repository/src/test/java/org/apache/atlas/repository/impexp/StartEntityFetchByExportRequestTest.java b/repository/src/test/java/org/apache/atlas/repository/impexp/StartEntityFetchByExportRequestTest.java index ddff0f088cc..85fb315ce13 100644 --- a/repository/src/test/java/org/apache/atlas/repository/impexp/StartEntityFetchByExportRequestTest.java +++ b/repository/src/test/java/org/apache/atlas/repository/impexp/StartEntityFetchByExportRequestTest.java @@ -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 { @@ -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\" } ]}"; @@ -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> 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 lastBindings = startEntityFetchByExportRequestSpy.getSuppliedBindingsMap(); + Set boundTypes = (Set) 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 lastBindings = startEntityFetchByExportRequestSpy.getSuppliedBindingsMap(); + + assertTrue(lastBindings != null && lastBindings.containsKey(BINDING_PARAMETER_TYPENAME), "Bindings should contain typeName"); + + Set boundTypes = (Set) 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 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 suppliedBindingsMap; + private String lastQuery; + private Map lastBindings; + private List> recordedBindings = new ArrayList<>(); public StartEntityFetchByExportRequestSpy(AtlasGraph atlasGraph, AtlasTypeRegistry typeRegistry) { super(atlasGraph, typeRegistry, atlasGremlin3QueryProvider); } public String getGeneratedQuery() { - return generatedQuery; + return lastQuery; } public Map getSuppliedBindingsMap() { - return suppliedBindingsMap; + return lastBindings; + } + + public List> getRecordedBindings() { + return recordedBindings; + } + + public void clearRecordedCalls() { + recordedBindings.clear(); + lastQuery = null; + lastBindings = null; } @Override List executeGremlinQuery(String query, Map bindings) { - this.generatedQuery = query; - this.suppliedBindingsMap = bindings; + this.lastQuery = query; + this.lastBindings = bindings; + this.recordedBindings.add(new java.util.HashMap<>(bindings)); return Collections.emptyList(); } } -} +} \ No newline at end of file