This is an automated email from the ASF dual-hosted git repository. asf-gitbox-commits pushed a commit to branch atlas-2.6 in repository https://gitbox.apache.org/repos/asf/atlas.git
commit f5d223316e2837fb30de1faebaf956bae47d96c4 Author: sheetalshah1007 <[email protected]> AuthorDate: Fri Sep 25 17:55:36 2026 +0530 ATLAS-5413: Initialize PatchContext before typedef JSON patches register Java handlers (#766) Co-authored-by: Sheetal Shah <[email protected]> --- .../repository/patches/AtlasPatchManager.java | 32 +++++++++++++++++++--- .../bootstrap/AtlasTypeDefStoreInitializer.java | 4 +-- .../repository/patches/AtlasPatchManagerTest.java | 19 +++++++++++++ .../AtlasTypeDefStoreInitializerTest.java | 8 ++++-- 4 files changed, 55 insertions(+), 8 deletions(-) diff --git a/repository/src/main/java/org/apache/atlas/repository/patches/AtlasPatchManager.java b/repository/src/main/java/org/apache/atlas/repository/patches/AtlasPatchManager.java index 30d3b894cf..10a8321435 100644 --- a/repository/src/main/java/org/apache/atlas/repository/patches/AtlasPatchManager.java +++ b/repository/src/main/java/org/apache/atlas/repository/patches/AtlasPatchManager.java @@ -47,6 +47,7 @@ public class AtlasPatchManager { private final GraphBackedSearchIndexer indexer; private final EntityGraphMapper entityGraphMapper; private PatchContext context; + private boolean defaultHandlersRegistered; @Inject public AtlasPatchManager(AtlasGraph atlasGraph, AtlasTypeRegistry typeRegistry, GraphBackedSearchIndexer indexer, EntityGraphMapper entityGraphMapper) { @@ -57,7 +58,7 @@ public class AtlasPatchManager { } public AtlasPatches getAllPatches() { - return context.getPatchRegistry().getAllPatches(); + return getOrCreatePatchContext().getPatchRegistry().getAllPatches(); } public void applyAll() { @@ -92,13 +93,36 @@ public class AtlasPatchManager { } public PatchContext getContext() { - return this.context; + return getOrCreatePatchContext(); } private void init() { LOG.info("==> AtlasPatchManager.init()"); - this.context = new PatchContext(atlasGraph, typeRegistry, indexer, entityGraphMapper); + getOrCreatePatchContext(); + registerDefaultHandlersIfNeeded(); + + LOG.info("<== AtlasPatchManager.init()"); + } + + /** + * Ensures {@link PatchContext} exists before JSON typedef patches register Java patch handlers + * (AtlasTypeDefStoreInitializer runs before {@link AtlasPatchService#applyAll()}). + */ + public PatchContext getOrCreatePatchContext() { + if (this.context == null) { + LOG.info("AtlasPatchManager.getOrCreatePatchContext(): creating PatchContext"); + + this.context = new PatchContext(atlasGraph, typeRegistry, indexer, entityGraphMapper); + } + + return this.context; + } + + private void registerDefaultHandlersIfNeeded() { + if (defaultHandlersRegistered) { + return; + } // register all java patches here handlers.add(new UniqueAttributePatch(context)); @@ -113,6 +137,6 @@ public class AtlasPatchManager { handlers.add(new ProcessImpalaNamePatch(context)); handlers.add(new ReplaceHugeSparkProcessAttributesPatch(context)); - LOG.info("<== AtlasPatchManager.init()"); + defaultHandlersRegistered = true; } } diff --git a/repository/src/main/java/org/apache/atlas/repository/store/bootstrap/AtlasTypeDefStoreInitializer.java b/repository/src/main/java/org/apache/atlas/repository/store/bootstrap/AtlasTypeDefStoreInitializer.java index 31e748096c..18de9a0e9b 100644 --- a/repository/src/main/java/org/apache/atlas/repository/store/bootstrap/AtlasTypeDefStoreInitializer.java +++ b/repository/src/main/java/org/apache/atlas/repository/store/bootstrap/AtlasTypeDefStoreInitializer.java @@ -870,7 +870,7 @@ public class AtlasTypeDefStoreInitializer implements ActiveStateChangeHandler { LOG.info("adding a Java patch to update entities of {} with new mandatory attributes", typeName); // Java patch handler to add mandatory attributes - patchManager.addPatchHandler(new AddMandatoryAttributesPatch(patchManager.getContext(), patch.getId(), typeName, attributesToAdd)); + patchManager.addPatchHandler(new AddMandatoryAttributesPatch(patchManager.getOrCreatePatchContext(), patch.getId(), typeName, attributesToAdd)); ret = APPLIED; } finally { @@ -1534,7 +1534,7 @@ public class AtlasTypeDefStoreInitializer implements ActiveStateChangeHandler { LOG.info("Update entities of {} with new supertypes", typeName); // add to java patch handlers to update entity supertypes - patchManager.addPatchHandler(new SuperTypesUpdatePatch(patchManager.getContext(), patch.getId(), typeName)); + patchManager.addPatchHandler(new SuperTypesUpdatePatch(patchManager.getOrCreatePatchContext(), patch.getId(), typeName)); ret = APPLIED; } else { diff --git a/repository/src/test/java/org/apache/atlas/repository/patches/AtlasPatchManagerTest.java b/repository/src/test/java/org/apache/atlas/repository/patches/AtlasPatchManagerTest.java index 066c74c895..3d5932e232 100644 --- a/repository/src/test/java/org/apache/atlas/repository/patches/AtlasPatchManagerTest.java +++ b/repository/src/test/java/org/apache/atlas/repository/patches/AtlasPatchManagerTest.java @@ -91,6 +91,7 @@ public class AtlasPatchManagerTest { // Mock vertex creation for createOrUpdatePatchVertex AtlasVertex mockVertex = mock(AtlasVertex.class); when(atlasGraph.addVertex()).thenReturn(mockVertex); + doNothing().when(atlasGraph).commit(); patchManager = new AtlasPatchManager(atlasGraph, typeRegistry, indexer, entityGraphMapper); } @@ -156,6 +157,24 @@ public class AtlasPatchManagerTest { assertEquals(handlers.get(0), mockHandler); } + @Test + public void testGetOrCreatePatchContextCreatesContextOnce() { + assertNotNull(patchManager.getOrCreatePatchContext()); + PatchContext first = patchManager.getOrCreatePatchContext(); + PatchContext second = patchManager.getOrCreatePatchContext(); + assertEquals(first, second); + } + + @Test + public void testInitDoesNotDuplicateDefaultHandlers() throws Exception { + patchManager.getOrCreatePatchContext(); + invokeInitMethod(); + int handlerCountAfterFirstInit = getHandlersFromManager().size(); + + invokeInitMethod(); + assertEquals(getHandlersFromManager().size(), handlerCountAfterFirstInit); + } + @Test public void testGetContext() throws Exception { invokeInitMethod(); diff --git a/repository/src/test/java/org/apache/atlas/repository/store/bootstrap/AtlasTypeDefStoreInitializerTest.java b/repository/src/test/java/org/apache/atlas/repository/store/bootstrap/AtlasTypeDefStoreInitializerTest.java index 6c1ce5fcfd..320d8f42f4 100644 --- a/repository/src/test/java/org/apache/atlas/repository/store/bootstrap/AtlasTypeDefStoreInitializerTest.java +++ b/repository/src/test/java/org/apache/atlas/repository/store/bootstrap/AtlasTypeDefStoreInitializerTest.java @@ -42,6 +42,7 @@ import org.apache.atlas.repository.graphdb.AtlasGraphQuery; import org.apache.atlas.repository.graphdb.AtlasVertex; import org.apache.atlas.repository.patches.AtlasPatchManager; import org.apache.atlas.repository.patches.AtlasPatchRegistry; +import org.apache.atlas.repository.patches.PatchContext; import org.apache.atlas.store.AtlasTypeDefStore; import org.apache.atlas.type.AtlasEntityType; import org.apache.atlas.type.AtlasStructType.AtlasAttribute; @@ -134,8 +135,11 @@ public class AtlasTypeDefStoreInitializerTest { lenient().when(typeDefStore.updateClassificationDefByName(anyString(), any())).thenReturn(new AtlasClassificationDef()); lenient().when(typeDefStore.updateRelationshipDefByName(anyString(), any())).thenReturn(new AtlasRelationshipDef()); - // Mock patch manager - lenient().when(patchManager.getContext()).thenReturn(null); + // Mock patch manager — production uses getOrCreatePatchContext before Java patch registration + PatchContext mockPatchContext = mock(PatchContext.class); + lenient().when(mockPatchContext.getPatchRegistry()).thenReturn(patchRegistry); + lenient().when(patchManager.getOrCreatePatchContext()).thenReturn(mockPatchContext); + lenient().when(patchManager.getContext()).thenReturn(mockPatchContext); initializer = new AtlasTypeDefStoreInitializer(typeDefStore, typeRegistry, graph, conf, patchManager);
