content: /**
   * 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.store.graph.v2;
  
  import org.apache.atlas.AtlasErrorCode;
  import org.apache.atlas.RequestContext;
  import org.apache.atlas.exception.AtlasBaseException;
  import org.apache.atlas.model.instance.AtlasCheckStateRequest;
  import org.apache.atlas.model.instance.AtlasCheckStateResult;
  import org.apache.atlas.model.instance.AtlasCheckStateResult.AtlasEntityState;
  import org.apache.atlas.repository.Constants;
  import org.apache.atlas.repository.graph.AtlasGraphProvider;
  import org.apache.atlas.repository.graph.GraphHelper;
  import org.apache.atlas.repository.graphdb.AtlasEdge;
  import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
  import org.apache.atlas.repository.graphdb.AtlasGraph;
  import org.apache.atlas.repository.graphdb.AtlasGraphQuery;
  import org.apache.atlas.repository.graphdb.AtlasVertex;
  import org.apache.atlas.type.AtlasClassificationType;
  import org.apache.atlas.type.AtlasEntityType;
  import org.apache.atlas.type.AtlasTypeRegistry;
  import org.apache.commons.collections.CollectionUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import org.springframework.stereotype.Component;
  
  import javax.inject.Inject;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  
  import static org.apache.atlas.repository.Constants.CLASSIFICATION_EDGE_IS_PROPAGATED_PROPERTY_KEY;
  import static org.apache.atlas.repository.graph.GraphHelper.getDelimitedClassificationNames;
  
  @Component
  public final class EntityStateChecker {
      private static final Logger LOG = LoggerFactory.getLogger(EntityStateChecker.class);
  
      private final AtlasGraph           graph;
      private final AtlasTypeRegistry    typeRegistry;
      private final EntityGraphRetriever entityRetriever;
  
      @Inject
      public EntityStateChecker(AtlasGraph graph, AtlasTypeRegistry typeRegistry) {
          this.graph           = graph;
          this.typeRegistry    = typeRegistry;
          this.entityRetriever = new EntityGraphRetriever(graph, typeRegistry);
      }
  
      public AtlasCheckStateResult checkState(AtlasCheckStateRequest request) throws AtlasBaseException {
          LOG.debug("==> checkState({})", request);
  
          AtlasCheckStateResult ret = new AtlasCheckStateResult();
  
          if (request != null) {
              if (CollectionUtils.isNotEmpty(request.getEntityGuids())) {
                  for (String guid : request.getEntityGuids()) {
                      checkEntityState(guid, request.getFixIssues(), ret);
                  }
              } else if (CollectionUtils.isNotEmpty(request.getEntityTypes())) {
                  final Collection<String> entityTypes;
  
                  if (request.getEntityTypes().contains("*")) {
                      entityTypes = typeRegistry.getAllEntityDefNames();
                  } else {
                      entityTypes = request.getEntityTypes();
                  }
  
                  LOG.info("checkState(): scanning for entities of {} types", entityTypes.size());
  
                  for (String typeName : entityTypes) {
                      AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName);
  
                      if (entityType == null) {
                          LOG.warn("checkState(): {} - entity-type not found", typeName);
  
                          continue;
                      }
  
                      LOG.info("checkState(): scanning for {} entities", typeName);
  
                      AtlasGraphQuery query = AtlasGraphProvider.getGraphInstance().query().has(Constants.ENTITY_TYPE_PROPERTY_KEY, typeName);
  
                      int count = 0;
  
                      for (Iterator<AtlasVertex> iter = query.vertices().iterator(); iter.hasNext(); count++) {
                          checkEntityState(iter.next(), request.getFixIssues(), ret);
                      }
  
                      LOG.info("checkState(): scanned {} {} entities", count, typeName);
                  }
              }
  
              int incorrectFixed          = ret.getEntitiesFixed();
              int incorrectPartiallyFixed = ret.getEntitiesPartiallyFixed();
              int incorrectNotFixed       = ret.getEntitiesNotFixed();
  
              if (incorrectFixed == 0 && incorrectPartiallyFixed == 0 && incorrectNotFixed == 0) {
                  ret.setState(AtlasCheckStateResult.State.OK);
              } else if (incorrectPartiallyFixed != 0) {
                  ret.setState(AtlasCheckStateResult.State.PARTIALLY_FIXED);
              } else if (incorrectNotFixed != 0) {
                  ret.setState(incorrectFixed > 0 ? AtlasCheckStateResult.State.PARTIALLY_FIXED : AtlasCheckStateResult.State.NOT_FIXED);
              } else {
                  ret.setState(AtlasCheckStateResult.State.FIXED);
              }
          }
  
          LOG.debug("<== checkState({}, {})", request, ret);
  
          return ret;
      }
  
      /**
       * Check an entity state given its GUID
       * @param guid
       * @return
       * @throws AtlasBaseException
       */
      public AtlasEntityState checkEntityState(String guid, boolean fixIssues, AtlasCheckStateResult result) throws AtlasBaseException {
          AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(this.graph, guid);
  
          if (entityVertex == null) {
              throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
          }
  
          return checkEntityState(entityVertex, fixIssues, result);
      }
  
      /**
       * Check an entity state given its vertex
       * @param entityVertex
       * @return
       * @throws AtlasBaseException
       */
      public AtlasEntityState checkEntityState(AtlasVertex entityVertex, boolean fixIssues, AtlasCheckStateResult result) throws AtlasBaseException {
          LOG.debug("==> checkEntityState(guid={})", AtlasGraphUtilsV2.getIdFromVertex(entityVertex));
  
          AtlasEntityState ret = new AtlasEntityState();
  
          ret.setGuid(AtlasGraphUtilsV2.getIdFromVertex(entityVertex));
          ret.setTypeName(AtlasGraphUtilsV2.getTypeName(entityVertex));
          ret.setName(getEntityName(entityVertex));
          ret.setStatus(AtlasGraphUtilsV2.getState(entityVertex));
          ret.setState(AtlasCheckStateResult.State.OK);
  
          checkEntityState_Classifications(entityVertex, ret, fixIssues);
  
          if (ret.getState() != AtlasCheckStateResult.State.OK) { // don't include clean entities in the response
              if (result.getEntities() == null) {
                  result.setEntities(new HashMap<>());
              }
  
              result.getEntities().put(ret.getGuid(), ret);
          }
  
          result.incrEntitiesScanned();
  
          switch (ret.getState()) {
              case FIXED:
                  result.incrEntitiesFixed();
                  break;
  
              case PARTIALLY_FIXED:
                  result.incrEntitiesPartiallyFixed();
                  break;
  
              case NOT_FIXED:
                  result.incrEntitiesNotFixed();
                  break;
  
              case OK:
                  result.incrEntitiesOk();
                  break;
          }
  
          LOG.info("checkEntityState(guid={}; type={}; name={}): {}", ret.getGuid(), ret.getTypeName(), ret.getName(), ret.getState());
  
          LOG.debug("<== checkEntityState({}): {}", ret.getGuid(), ret);
  
          return ret;
      }
  
      private void checkEntityState_Classifications(AtlasVertex entityVertex, AtlasEntityState result, boolean fixIssues) {
          LOG.debug("==> checkEntityState_Classifications({}, {})", result, fixIssues);
  
          Collection<String>  traitNames                 = entityVertex.getPropertyValues(Constants.TRAIT_NAMES_PROPERTY_KEY, String.class);
          Collection<String>  propagatedTraitNames       = entityVertex.getPropertyValues(Constants.PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, String.class);
          Collection<String>  traitVertexNames           = null;
          Collection<String>  propagatedTraitVertexNames = null;
          Iterable<AtlasEdge> edges                      = entityVertex.getEdges(AtlasEdgeDirection.OUT, Constants.CLASSIFICATION_LABEL);
  
          if (edges != null) {
              for (AtlasEdge edge : edges) {
                  Boolean                 isPropagated       = AtlasGraphUtilsV2.getEncodedProperty(edge, CLASSIFICATION_EDGE_IS_PROPAGATED_PROPERTY_KEY, Boolean.class);
                  String                  classificationName = GraphHelper.getTypeName(edge.getInVertex());
                  AtlasClassificationType classification     = typeRegistry.getClassificationTypeByName(classificationName);
  
                  if (classification != null) {
                      if (isPropagated != null && isPropagated) {
                          propagatedTraitVertexNames = addToCollection(propagatedTraitVertexNames, classificationName);
                      } else {
                          traitVertexNames = addToCollection(traitVertexNames, classificationName);
                      }
                  }
              }
  
              if (traitVertexNames == null) {
                  traitVertexNames = Collections.emptyList();
              }
  
              if (propagatedTraitVertexNames == null) {
                  propagatedTraitVertexNames = Collections.emptyList();
              }
          }
  
          Collection<String> traitNamesToAdd              = subtract(traitVertexNames, traitNames);
          Collection<String> traitNamesToRemove           = subtract(traitNames, traitVertexNames);
          Collection<String> propagatedTraitNamesToAdd    = subtract(propagatedTraitVertexNames, propagatedTraitNames);
          Collection<String> propagatedTraitNamesToRemove = subtract(propagatedTraitNames, propagatedTraitVertexNames);
  
          if (traitNamesToAdd != null || traitNamesToRemove != null || propagatedTraitNamesToAdd != null || propagatedTraitNamesToRemove != null) {
              List<String> issues = result.getIssues();
  
              if (issues == null) {
                  issues = new ArrayList<>();
  
                  result.setIssues(issues);
              }
  
              if (fixIssues) {
                  if (traitNamesToAdd != null || traitNamesToRemove != null) {
                      if (traitNamesToAdd != null) {
                          issues.add("incorrect property: __traitNames has missing classifications: " + traitNamesToAdd);
                      }
  
                      if (traitNamesToRemove != null) {
                          issues.add("incorrect property: __traitNames has unassigned classifications: " + traitNamesToRemove);
                      }
  
                      entityVertex.removeProperty(Constants.TRAIT_NAMES_PROPERTY_KEY);
                      entityVertex.removeProperty(Constants.CLASSIFICATION_NAMES_KEY);
  
                      if (traitVertexNames != null) {
                          for (String classificationName : traitVertexNames) {
                              AtlasGraphUtilsV2.addEncodedProperty(entityVertex, Constants.TRAIT_NAMES_PROPERTY_KEY, classificationName);
                          }
                      }
  
                      entityVertex.setProperty(Constants.CLASSIFICATION_NAMES_KEY, getDelimitedClassificationNames(traitVertexNames));
                  }
  
                  if (propagatedTraitNamesToAdd != null || propagatedTraitNamesToRemove != null) {
                      if (propagatedTraitNamesToAdd != null) {
                          issues.add("incorrect property: __propagatedTraitNames has missing classifications: " + propagatedTraitNamesToAdd);
                      }
  
                      if (propagatedTraitNamesToRemove != null) {
                          issues.add("incorrect property: __propagatedTraitNames has unassigned classifications: " + propagatedTraitNamesToRemove);
                      }
  
                      entityVertex.removeProperty(Constants.PROPAGATED_TRAIT_NAMES_PROPERTY_KEY);
                      entityVertex.removeProperty(Constants.PROPAGATED_CLASSIFICATION_NAMES_KEY);
  
                      if (propagatedTraitVertexNames != null) {
                          for (String classificationName : propagatedTraitVertexNames) {
                              AtlasGraphUtilsV2.addEncodedProperty(entityVertex, Constants.PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, classificationName);
                          }
                      }
  
                      entityVertex.setProperty(Constants.PROPAGATED_CLASSIFICATION_NAMES_KEY, getDelimitedClassificationNames(propagatedTraitVertexNames));
                  }
  
                  AtlasGraphUtilsV2.setEncodedProperty(entityVertex, Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY, RequestContext.get().getRequestTime());
  
                  result.setState(AtlasCheckStateResult.State.FIXED);
              } else {
                  result.setState(AtlasCheckStateResult.State.NOT_FIXED);
              }
          }
  
          LOG.debug("<== checkEntityState_Classifications({}, {})", result, fixIssues);
      }
  
      private String getEntityName(AtlasVertex entityVertex) throws AtlasBaseException {
          String              ret              = null;
          Map<String, Object> uniqueAttributes = entityRetriever.getEntityUniqueAttribute(entityVertex);
  
          if (uniqueAttributes != null) {
              Object val = uniqueAttributes.get("qualifiedName");
  
              if (val == null) {
                  for (Object attrVal : uniqueAttributes.values()) {
                      if (attrVal != null) {
                          ret = attrVal.toString();
  
                          break;
                      }
                  }
              } else {
                  ret = val.toString();
              }
          }
  
          return ret;
      }
  
      private Collection<String> addToCollection(Collection<String> list, String str) {
          if (list == null) {
              list = new ArrayList<>();
          }
  
          list.add(str);
  
          return list;
      }
  
      // return elements in 'col1' that are not in 'col2'
      private Collection<String> subtract(Collection<String> col1, Collection<String> col2) {
          Collection<String> ret = null;
  
          if (col2 == null) {
              ret = col1;
          } else if (col1 != null) {
              for (String elem : col1) {
                  if (!col2.contains(elem)) {
                      if (ret == null) {
                          ret = new ArrayList<>();
                      }
  
                      ret.add(elem);
                  }
              }
          }
  
          return ret;
      }
  }
  