From 3869ffcb8faf903a86809480cbee2a769984cf02 Mon Sep 17 00:00:00 2001 From: Umesh Patil Date: Fri, 12 Dec 2025 18:37:23 +0530 Subject: [PATCH 1/2] Export APIs: Exporting parent type doesn't export child type --- .../StartEntityFetchByExportRequest.java | 72 +++++++++++++++++-- 1 file changed, 65 insertions(+), 7 deletions(-) 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..e4592ae9428 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 @@ -40,6 +40,7 @@ import java.util.ArrayList; import java.util.Arrays; +import java.util.Collection; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -104,6 +105,10 @@ public List get(AtlasExportRequest exportRequest, AtlasObjectId item) { return ret; } + if (StringUtils.isEmpty(item.getTypeName()) && MapUtils.isNotEmpty(item.getUniqueAttributes())) { + return getEntitiesForMatchTypeType(item, MATCH_TYPE_FOR_TYPE); + } + if (StringUtils.isNotEmpty(item.getTypeName()) && MapUtils.isNotEmpty(item.getUniqueAttributes())) { ret = getEntitiesForMatchTypeUsingUniqueAttributes(item, matchType); @@ -145,9 +150,16 @@ List executeGremlinQuery(String query, Map bindings) { } private List getEntitiesForMatchTypeUsingUniqueAttributes(AtlasObjectId item, String matchType) throws AtlasBaseException { + final String typeName = item.getTypeName(); + final AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName); + + if (entityType == null || (CollectionUtils.isNotEmpty(entityType.getAllSubTypes()) && entityType.getSuperTypes().isEmpty())) { + LOG.info("Unique attribute lookup requested for abstract/generic type '{}'. Rerouting to general search.", typeName); + + return getEntitiesForMatchTypeType(item, AtlasExportRequest.MATCH_TYPE_FOR_TYPE); + } + final String queryTemplate = getQueryTemplateForMatchType(matchType); - final String typeName = item.getTypeName(); - final AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName); Set ret = new HashSet<>(); @@ -177,15 +189,61 @@ private List getEntitiesForMatchTypeUsingUniqueAttributes(AtlasObjectId return new ArrayList<>(ret); } - private List getEntitiesForMatchTypeType(AtlasObjectId item, String matchType) { - return executeGremlinQuery(getQueryTemplateForMatchType(matchType), getBindingsForTypeName(item.getTypeName())); + private List getEntitiesForMatchTypeType(AtlasObjectId item, String matchType) throws AtlasBaseException { + String typeName = item.getTypeName(); + + AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName); + + if (StringUtils.isNotEmpty(typeName) && entityType == null) { + throw new AtlasBaseException(AtlasErrorCode.UNKNOWN_TYPENAME, typeName); + } + + boolean isRootOrGenericSearch = StringUtils.isEmpty(typeName) || (CollectionUtils.isNotEmpty(entityType.getAllSubTypes()) && entityType.getSuperTypes().isEmpty()); + + if (isRootOrGenericSearch) { + LOG.info("Handling export for root or generic type: Executing generic query for all concrete entities."); + + Collection allConcreteEntityTypes = typeRegistry.getAllEntityDefNames(); + + if (CollectionUtils.isEmpty(allConcreteEntityTypes)) { + return new ArrayList<>(); + } + + String allForTypeQuery = getQueryTemplateForMatchType(AtlasExportRequest.MATCH_TYPE_FOR_TYPE); + + HashMap bindings = new HashMap<>(); + bindings.put(BINDING_PARAMETER_TYPENAME, new HashSet<>(allConcreteEntityTypes)); + + return executeGremlinQuery(allForTypeQuery, bindings); + } + + return executeGremlinQuery(getQueryTemplateForMatchType(matchType), getBindingsForTypeName(typeName)); } - private HashMap getBindingsForTypeName(String typeName) { + private HashMap getBindingsForTypeName(String typeName) throws AtlasBaseException { HashMap ret = new HashMap<>(); + Set typeNamesToQuery = new HashSet<>(); - ret.put(BINDING_PARAMETER_TYPENAME, new HashSet<>(Arrays.asList(StringUtils.split(typeName, ",")))); + 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; } @@ -211,4 +269,4 @@ private static Map initMatchTypeQueryMap(AtlasGremlinQueryProvid return ret; } -} +} \ No newline at end of file From 8dad881b1d09210b32ed06f0bcb32a53f4c83304 Mon Sep 17 00:00:00 2001 From: Umesh Patil Date: Mon, 29 Dec 2025 16:55:01 +0530 Subject: [PATCH 2/2] ATLAS-5159 : Added UTs for export entity fetch logic. --- .../StartEntityFetchByExportRequestTest.java | 54 ++++++++++++++++++- 1 file changed, 53 insertions(+), 1 deletion(-) 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..0ec3ac914d6 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 @@ -38,11 +38,14 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Set; +import static org.apache.atlas.model.impexp.AtlasExportRequest.MATCH_TYPE_FOR_TYPE; 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.assertTrue; @Guice(modules = TestModules.TestOnlyModule.class) public class StartEntityFetchByExportRequestTest extends AtlasTestBase { @@ -80,6 +83,55 @@ public void fetchTypeUniqueAttributes() { assertEquals(startEntityFetchByExportRequestSpy.getSuppliedBindingsMap().get(BINDING_PARAMTER_ATTR_VALUE), "stocks@cl1"); } + @Test + public void fetchReferenceableUniqueAttributes() { + String exportRequestJson = "{ \"itemsToExport\": [ { \"typeName\": \"Referenceable\", \"uniqueAttributes\": {\"qualifiedName\": \"stocks@cl1\"} } ]}"; + AtlasExportRequest exportRequest = AtlasType.fromJson(exportRequestJson, AtlasExportRequest.class); + + startEntityFetchByExportRequestSpy.get(exportRequest); + + assertEquals(startEntityFetchByExportRequestSpy.getGeneratedQuery(), startEntityFetchByExportRequestSpy.getQueryTemplateForMatchType(MATCH_TYPE_FOR_TYPE)); + + Object typeNameBinding = startEntityFetchByExportRequestSpy.getSuppliedBindingsMap().get(BINDING_PARAMETER_TYPENAME); + assertTrue(typeNameBinding instanceof Set); + + Set boundTypes = (Set) typeNameBinding; + assertTrue(boundTypes.size() > 1); + assertTrue(boundTypes.contains("hive_table")); + } + + @Test + public void fetchTypeExpansion() { + String exportRequestJson = "{ \"itemsToExport\": [ { \"typeName\": \"Asset\" } ], \"options\": {\"matchType\": \"forType\"} }"; + AtlasExportRequest exportRequest = AtlasType.fromJson(exportRequestJson, AtlasExportRequest.class); + + startEntityFetchByExportRequestSpy.get(exportRequest); + + Set boundTypes = (Set) startEntityFetchByExportRequestSpy.getSuppliedBindingsMap().get(BINDING_PARAMETER_TYPENAME); + assertTrue(boundTypes.contains("Asset")); + assertTrue(boundTypes.contains("hive_db")); + } + + @Test + public void fetchEmptyTypeUniqueAttributes() { + String exportRequestJson = "{ \"itemsToExport\": [ { \"typeName\": \"\", \"uniqueAttributes\": {\"qualifiedName\": \"stocks@cl1\"} } ]}"; + AtlasExportRequest exportRequest = AtlasType.fromJson(exportRequestJson, AtlasExportRequest.class); + + startEntityFetchByExportRequestSpy.get(exportRequest); + + Set boundTypes = (Set) startEntityFetchByExportRequestSpy.getSuppliedBindingsMap().get(BINDING_PARAMETER_TYPENAME); + assertEquals(boundTypes.size(), typeRegistry.getAllEntityDefNames().size()); + } + + @Test + public void fetchUnknownType() { + String exportRequestJson = "{ \"itemsToExport\": [ { \"typeName\": \"InvalidType\" } ]}"; + AtlasExportRequest exportRequest = AtlasType.fromJson(exportRequestJson, AtlasExportRequest.class); + + List result = startEntityFetchByExportRequestSpy.get(exportRequest); + assertTrue(result.isEmpty()); + } + @BeforeClass void setup() throws IOException, AtlasBaseException { super.basicSetup(typeDefStore, typeRegistry); @@ -112,4 +164,4 @@ List executeGremlinQuery(String query, Map bindings) { return Collections.emptyList(); } } -} +} \ No newline at end of file