ATLAS-3077: Handle java patches in patch framework

intg/src/main/java/org/apache/atlas/model/patches/AtlasPatch.java

@@ -50,12 +50,12 @@ public class AtlasPatch implements Serializable {
     private long        updatedTime;
     private PatchStatus status;
 
-    public enum PatchStatus { APPLIED, SKIPPED, FAILED, UNKNOWN }
+    public enum PatchStatus { UNKNOWN, APPLIED, SKIPPED, FAILED }
 
     public AtlasPatch() { }
 
-    public AtlasPatch(String id, String patchName, String type, String action, PatchStatus status, String updatedBy,
-                      String createdBy, long createdTime, long updatedTime) {
+    public AtlasPatch(String id, String patchName, String type, String action, PatchStatus status,
+                      String updatedBy, String createdBy, long createdTime, long updatedTime) {
         this.id          = id;
         this.description = patchName;
         this.type        = type;

repository/src/main/java/org/apache/atlas/repository/graph/GraphBackedSearchIndexer.java

@@ -108,7 +108,7 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
     private boolean     recomputeIndexedKeys = true;
     private Set<String> vertexIndexKeys      = new HashSet<>();
 
-    private enum UniqueKind { NONE, GLOBAL_UNIQUE, PER_TYPE_UNIQUE }
+    public enum UniqueKind { NONE, GLOBAL_UNIQUE, PER_TYPE_UNIQUE }
 
     @Inject
     public GraphBackedSearchIndexer(AtlasTypeRegistry typeRegistry) throws AtlasException {
@@ -431,7 +431,7 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
         return type instanceof AtlasRelationshipType;
     }
 
-    private Class getPrimitiveClass(String attribTypeName) {
+    public Class getPrimitiveClass(String attribTypeName) {
         String attributeTypeName = attribTypeName.toLowerCase();
 
         switch (attributeTypeName) {
@@ -461,7 +461,7 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
         throw new IllegalArgumentException(String.format("Unknown primitive typename %s", attribTypeName));
     }
 
-    private AtlasCardinality toAtlasCardinality(AtlasAttributeDef.Cardinality cardinality) {
+    public AtlasCardinality toAtlasCardinality(AtlasAttributeDef.Cardinality cardinality) {
         switch (cardinality) {
             case SINGLE:
                 return SINGLE;
@@ -500,7 +500,7 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
         return propertyKey;
     }
 
-    private void createVertexIndex(AtlasGraphManagement management, String propertyName, UniqueKind uniqueKind, Class propertyClass,
+    public void createVertexIndex(AtlasGraphManagement management, String propertyName, UniqueKind uniqueKind, Class propertyClass,
                                   AtlasCardinality cardinality, boolean createCompositeIndex, boolean createCompositeIndexWithTypeAndSuperTypes) {
         if (propertyName != null) {
             AtlasPropertyKey propertyKey = management.getPropertyKey(propertyName);
@@ -704,7 +704,7 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
         return !(INDEX_EXCLUSION_CLASSES.contains(propertyClass) || cardinality.isMany());
     }
     
-    private void commit(AtlasGraphManagement management) throws IndexException {
+    public void commit(AtlasGraphManagement management) throws IndexException {
         try {
             management.commit();
 
@@ -715,7 +715,7 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
         }
     }
 
-    private void rollback(AtlasGraphManagement management) throws IndexException {
+    public void rollback(AtlasGraphManagement management) throws IndexException {
         try {
             management.rollback();
 

repository/src/main/java/org/apache/atlas/repository/patches/AtlasJavaPatchHandler.java

+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.atlas.repository.patches;
+
+import org.apache.atlas.RequestContext;
+import org.apache.atlas.model.patches.AtlasPatch.PatchStatus;
+import org.apache.atlas.repository.graph.GraphBackedSearchIndexer;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
+import org.apache.atlas.repository.graphdb.AtlasVertex;
+import org.apache.atlas.repository.store.graph.v2.EntityGraphRetriever;
+import org.apache.atlas.type.AtlasTypeRegistry;
+import org.apache.commons.collections.MapUtils;
+
+import java.util.Map;
+
+import static org.apache.atlas.model.patches.AtlasPatch.PatchStatus.UNKNOWN;
+import static org.apache.atlas.repository.Constants.CREATED_BY_KEY;
+import static org.apache.atlas.repository.Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY;
+import static org.apache.atlas.repository.Constants.MODIFIED_BY_KEY;
+import static org.apache.atlas.repository.Constants.PATCH_DESCRIPTION_PROPERTY_KEY;
+import static org.apache.atlas.repository.Constants.PATCH_ID_PROPERTY_KEY;
+import static org.apache.atlas.repository.Constants.PATCH_STATE_PROPERTY_KEY;
+import static org.apache.atlas.repository.Constants.PATCH_TYPE_PROPERTY_KEY;
+import static org.apache.atlas.repository.Constants.TIMESTAMP_PROPERTY_KEY;
+import static org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2.findByPatchId;
+import static org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2.setEncodedProperty;
+import static org.apache.atlas.repository.store.graph.v2.AtlasTypeDefGraphStoreV2.getCurrentUser;
+
+public abstract class AtlasJavaPatchHandler {
+    public final AtlasGraph               graph;
+    public final AtlasTypeRegistry        typeRegistry;
+    public final Map<String, PatchStatus> patchesRegistry;
+    public final EntityGraphRetriever     entityRetriever;
+    public final GraphBackedSearchIndexer indexer;
+    public final PatchContext             context;
+    public final String                   patchId;
+    public final String                   patchDescription;
+
+    private PatchStatus patchStatus;
+
+    public static final String JAVA_PATCH_TYPE = "JAVA_PATCH";
+
+    public AtlasJavaPatchHandler(PatchContext context, String patchId, String patchDescription) {
+        this.context          = context;
+        this.graph            = context.getGraph();
+        this.typeRegistry     = context.getTypeRegistry();
+        this.indexer          = context.getIndexer();
+        this.patchesRegistry  = context.getPatchesRegistry();
+        this.patchId          = patchId;
+        this.patchDescription = patchDescription;
+        this.patchStatus      = getPatchStatus(patchesRegistry);
+        this.entityRetriever  = new EntityGraphRetriever(typeRegistry);
+
+        init();
+    }
+
+    private void init() {
+        PatchStatus patchStatus = getPatchStatus();
+
+        if (patchStatus == UNKNOWN) {
+            AtlasVertex patchVertex = graph.addVertex();
+
+            setEncodedProperty(patchVertex, PATCH_ID_PROPERTY_KEY, patchId);
+            setEncodedProperty(patchVertex, PATCH_DESCRIPTION_PROPERTY_KEY, patchDescription);
+            setEncodedProperty(patchVertex, PATCH_TYPE_PROPERTY_KEY, JAVA_PATCH_TYPE);
+            setEncodedProperty(patchVertex, PATCH_STATE_PROPERTY_KEY, getPatchStatus().toString());
+            setEncodedProperty(patchVertex, TIMESTAMP_PROPERTY_KEY, RequestContext.get().getRequestTime());
+            setEncodedProperty(patchVertex, MODIFICATION_TIMESTAMP_PROPERTY_KEY, RequestContext.get().getRequestTime());
+            setEncodedProperty(patchVertex, CREATED_BY_KEY, getCurrentUser());
+            setEncodedProperty(patchVertex, MODIFIED_BY_KEY, getCurrentUser());
+
+            graph.commit();
+
+            addToPatchesRegistry(patchId, getPatchStatus());
+        }
+    }
+
+    private PatchStatus getPatchStatus(Map<String, PatchStatus> patchesRegistry) {
+        PatchStatus ret = UNKNOWN;
+
+        if (MapUtils.isNotEmpty(patchesRegistry) && patchesRegistry.containsKey(patchId)) {
+            ret = patchesRegistry.get(patchId);
+        }
+
+        return ret;
+    }
+
+    public void updatePatchVertex(PatchStatus patchStatus) {
+        AtlasVertex patchVertex = findByPatchId(patchId);
+
+        if (patchVertex != null) {
+            setEncodedProperty(patchVertex, PATCH_STATE_PROPERTY_KEY, patchStatus.toString());
+            setEncodedProperty(patchVertex, MODIFICATION_TIMESTAMP_PROPERTY_KEY, RequestContext.get().getRequestTime());
+            setEncodedProperty(patchVertex, MODIFIED_BY_KEY, getCurrentUser());
+
+            graph.commit();
+
+            addToPatchesRegistry(getPatchId(), getPatchStatus());
+        }
+    }
+
+    public PatchStatus getPatchStatus() {
+        return patchStatus;
+    }
+
+    public void addToPatchesRegistry(String patchId, PatchStatus status) {
+        getPatchesRegistry().put(patchId, status);
+    }
+
+    public void setPatchStatus(PatchStatus patchStatus) {
+        this.patchStatus = patchStatus;
+    }
+
+    public String getPatchId() {
+        return patchId;
+    }
+
+    public Map<String, PatchStatus> getPatchesRegistry() {
+        return patchesRegistry;
+    }
+
+    public abstract void applyPatch();
+}
\ No newline at end of file

repository/src/main/java/org/apache/atlas/repository/patches/PatchContext.java

+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.atlas.repository.patches;
+
+import org.apache.atlas.model.patches.AtlasPatch.PatchStatus;
+import org.apache.atlas.repository.graph.GraphBackedSearchIndexer;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
+import org.apache.atlas.type.AtlasTypeRegistry;
+
+import java.util.Map;
+
+/**
+ * Patch context for typedef and java patches.
+ */
+public class PatchContext {
+    private final AtlasGraph               graph;
+    private final AtlasTypeRegistry        typeRegistry;
+    private final GraphBackedSearchIndexer indexer;
+    private final Map<String, PatchStatus> patchesRegistry;
+
+    public PatchContext(AtlasGraph graph, AtlasTypeRegistry typeRegistry, GraphBackedSearchIndexer indexer,
+                        Map<String, PatchStatus> patchesRegistry) {
+        this.graph           = graph;
+        this.typeRegistry    = typeRegistry;
+        this.indexer         = indexer;
+        this.patchesRegistry = patchesRegistry;
+    }
+
+    public AtlasGraph getGraph() {
+        return graph;
+    }
+
+    public AtlasTypeRegistry getTypeRegistry() {
+        return typeRegistry;
+    }
+
+    public GraphBackedSearchIndexer getIndexer() {
+        return indexer;
+    }
+
+    public Map<String, PatchStatus> getPatchesRegistry() {
+        return patchesRegistry;
+    }
+}
\ No newline at end of file

repository/src/main/java/org/apache/atlas/repository/patches/UniqueAttributePatchHandler.java

+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.atlas.repository.patches;
+
+import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
+import org.apache.atlas.repository.IndexException;
+import org.apache.atlas.repository.graph.GraphBackedSearchIndexer.UniqueKind;
+import org.apache.atlas.repository.graphdb.AtlasCardinality;
+import org.apache.atlas.repository.graphdb.AtlasGraphManagement;
+import org.apache.atlas.repository.graphdb.AtlasVertex;
+import org.apache.atlas.type.AtlasEntityType;
+import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
+import org.apache.commons.collections.MapUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.Map;
+
+import static org.apache.atlas.model.patches.AtlasPatch.PatchStatus.APPLIED;
+import static org.apache.atlas.model.patches.AtlasPatch.PatchStatus.FAILED;
+import static org.apache.atlas.repository.graph.GraphHelper.getGuid;
+import static org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2.findActiveEntityVerticesByType;
+import static org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2.setEncodedProperty;
+
+public class UniqueAttributePatchHandler extends AtlasJavaPatchHandler {
+    private static final String PATCH_ID          = "JAVA_PATCH_0000_001";
+    private static final String PATCH_DESCRIPTION = "Add new vertex property for each unique attribute of active entities";
+    private static final Logger LOG               = LoggerFactory.getLogger(UniqueAttributePatchHandler.class);
+
+    public UniqueAttributePatchHandler(PatchContext context) {
+        super(context, PATCH_ID, PATCH_DESCRIPTION);
+    }
+
+    @Override
+    public void applyPatch() {
+        Collection<AtlasEntityType> allEntityTypes = typeRegistry.getAllEntityTypes();
+        boolean                     patchFailed    = false;
+
+        for (AtlasEntityType entityType : allEntityTypes) {
+            String                      typeName          = entityType.getTypeName();
+            Map<String, AtlasAttribute> uniqAttributes    = entityType.getUniqAttributes();
+            int                         entitiesProcessed = 0;
+
+            LOG.info("Applying java patch: {} for type: {}", getPatchId(), typeName);
+
+            if (MapUtils.isNotEmpty(uniqAttributes)) {
+                Collection<AtlasAttribute> attributes = uniqAttributes.values();
+
+                try {
+                    // register unique attribute property keys in graph
+                    registerUniqueAttrPropertyKeys(attributes);
+
+                    Iterator<AtlasVertex> iterator = findActiveEntityVerticesByType(typeName);
+
+                    while (iterator.hasNext()) {
+                        AtlasVertex entityVertex = iterator.next();
+                        boolean     patchApplied = false;
+
+                        for (AtlasAttribute attribute : attributes) {
+                            String                       uniquePropertyKey = attribute.getVertexUniquePropertyName();
+                            Collection<? extends String> propertyKeys      = entityVertex.getPropertyKeys();
+
+                            if (!propertyKeys.contains(uniquePropertyKey)) {
+                                String            propertyKey   = attribute.getVertexPropertyName();
+                                AtlasAttributeDef attributeDef  = attribute.getAttributeDef();
+                                Object            uniqAttrValue = entityRetriever.mapVertexToPrimitive(entityVertex, propertyKey, attributeDef);
+
+                                // add the unique attribute property to vertex
+                                setEncodedProperty(entityVertex, uniquePropertyKey, uniqAttrValue);
+
+                                try {
+                                    graph.commit();
+
+                                    patchApplied = true;
+
+                                    if (LOG.isDebugEnabled()) {
+                                        LOG.debug("{}: Added unique attribute property: {} to entity: {} ({})",
+                                                   PATCH_ID, uniquePropertyKey, getGuid(entityVertex), typeName);
+                                    }
+                                } catch (Throwable t) {
+                                    LOG.warn("Java patch ({}): failed to update entity guid: {}; typeName: {}; attrName: {}; attrValue: {}",
+                                              getPatchId(), getGuid(entityVertex), typeName, attribute.getName(), uniqAttrValue);
+
+                                    continue;
+                                }
+                            }
+                        }
+
+                        if (patchApplied) {
+                            entitiesProcessed++;
+                        }
+
+                        if (entitiesProcessed % 1000 == 0) {
+                            LOG.info("Java patch: {} : processed {} {} entities.", getPatchId(), entitiesProcessed, typeName);
+                        }
+                    }
+                } catch (IndexException e) {
+                    LOG.error("Java patch: {} failed! error: {}", getPatchId(), e);
+
+                    patchFailed = true;
+
+                    break;
+                }
+            }
+
+            LOG.info("Applied java patch ({}) for type: {}; Total processed: {}", getPatchId(), typeName, entitiesProcessed);
+        }
+
+        if (patchFailed) {
+            setPatchStatus(FAILED);
+        } else {
+            setPatchStatus(APPLIED);
+        }
+
+        LOG.info("Applied java patch: {}; status: {}", getPatchId(), getPatchStatus());
+
+        updatePatchVertex(getPatchStatus());
+    }
+
+    private void registerUniqueAttrPropertyKeys(Collection<AtlasAttribute> attributes) throws IndexException {
+        AtlasGraphManagement management = graph.getManagementSystem();
+        boolean              idxCreated = false;
+
+        for (AtlasAttribute attribute : attributes) {
+            String  uniquePropertyName       = attribute.getVertexUniquePropertyName();
+            boolean uniquePropertyNameExists = management.getPropertyKey(uniquePropertyName) != null;
+
+            if (!uniquePropertyNameExists) {
+                AtlasAttributeDef attributeDef   = attribute.getAttributeDef();
+                boolean           isIndexable    = attributeDef.getIsIndexable();
+                String            attribTypeName = attributeDef.getTypeName();
+                Class             propertyClass  = indexer.getPrimitiveClass(attribTypeName);
+                AtlasCardinality  cardinality    = indexer.toAtlasCardinality(attributeDef.getCardinality());
+
+                indexer.createVertexIndex(management, uniquePropertyName, UniqueKind.NONE, propertyClass, cardinality, isIndexable, true);
+
+                idxCreated = true;
+            }
+        }
+
+        //Commit indexes
+        if (idxCreated) {
+            indexer.commit(management);
+        }
+    }
+}
\ No newline at end of file

webapp/src/main/java/org/apache/atlas/web/resources/AdminResource.java

@@ -564,7 +564,7 @@ public class AdminResource {
             LOG.debug("==> AdminResource.getAtlasPatches()");
         }
 
-        AtlasPatches ret = AtlasGraphUtilsV2.getPatches();
+        AtlasPatches ret = AtlasGraphUtilsV2.getAllPatches();
 
         if (LOG.isDebugEnabled()) {
             LOG.debug("<== AdminResource.getAtlasPatches()");