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.TypeCategory;
  import org.apache.atlas.model.instance.AtlasEntity;
  import org.apache.atlas.model.instance.AtlasObjectId;
  import org.apache.atlas.model.instance.AtlasStruct;
  import org.apache.atlas.repository.graphdb.AtlasGraph;
  import org.apache.atlas.repository.store.graph.EntityGraphDiscovery;
  import org.apache.atlas.repository.store.graph.EntityGraphDiscoveryContext;
  import org.apache.atlas.repository.store.graph.EntityResolver;
  import org.apache.atlas.type.AtlasArrayType;
  import org.apache.atlas.type.AtlasBuiltInTypes.AtlasObjectIdType;
  import org.apache.atlas.type.AtlasEntityType;
  import org.apache.atlas.type.AtlasMapType;
  import org.apache.atlas.type.AtlasStructType;
  import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
  import org.apache.atlas.type.AtlasType;
  import org.apache.atlas.type.AtlasTypeRegistry;
  import org.apache.atlas.type.AtlasTypeUtil;
  import org.apache.atlas.type.TemplateToken;
  import org.apache.atlas.utils.AtlasEntityUtil;
  import org.apache.atlas.utils.AtlasPerfMetrics.MetricRecorder;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.HashSet;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  
  import static org.apache.atlas.repository.store.graph.v2.EntityGraphMapper.validateCustomAttributes;
  import static org.apache.atlas.repository.store.graph.v2.EntityGraphMapper.validateLabels;
  
  public class AtlasEntityGraphDiscoveryV2 implements EntityGraphDiscovery {
      private static final Logger LOG = LoggerFactory.getLogger(AtlasEntityGraphDiscoveryV2.class);
  
      private final AtlasGraph                  graph;
      private final AtlasTypeRegistry           typeRegistry;
      private final EntityGraphDiscoveryContext discoveryContext;
      private final EntityGraphMapper           entityGraphMapper;
  
      public AtlasEntityGraphDiscoveryV2(AtlasGraph graph, AtlasTypeRegistry typeRegistry, EntityStream entityStream, EntityGraphMapper entityGraphMapper) {
          this.graph             = graph;
          this.typeRegistry      = typeRegistry;
          this.discoveryContext  = new EntityGraphDiscoveryContext(typeRegistry, entityStream);
          this.entityGraphMapper = entityGraphMapper;
      }
  
      @Override
      public void init() throws AtlasBaseException {
          //Nothing to do
      }
  
      @Override
      public EntityGraphDiscoveryContext discoverEntities() throws AtlasBaseException {
          // walk through entities in stream and validate them; record entity references
          discover();
  
          // resolve entity references discovered in previous step
          resolveReferences();
  
          return discoveryContext;
      }
  
      @Override
      public void validateAndNormalize(AtlasEntity entity) throws AtlasBaseException {
          List<String> messages = new ArrayList<>();
  
          if (!AtlasTypeUtil.isValidGuid(entity.getGuid())) {
              throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, "invalid guid " + entity.getGuid());
          }
  
          AtlasEntityType type = typeRegistry.getEntityTypeByName(entity.getTypeName());
  
          if (type == null) {
              throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), entity.getTypeName());
          }
  
          validateCustomAttributes(entity);
  
          validateLabels(entity.getLabels());
  
          type.validateValue(entity, entity.getTypeName(), messages);
  
          if (!messages.isEmpty()) {
              throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS, messages);
          }
  
          type.getNormalizedValue(entity);
      }
  
      @Override
      public void validateAndNormalizeForUpdate(AtlasEntity entity) throws AtlasBaseException {
          List<String> messages = new ArrayList<>();
  
          if (!AtlasTypeUtil.isValidGuid(entity.getGuid())) {
              throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, "invalid guid " + entity.getGuid());
          }
  
          AtlasEntityType type = typeRegistry.getEntityTypeByName(entity.getTypeName());
  
          if (type == null) {
              throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), entity.getTypeName());
          }
  
          validateCustomAttributes(entity);
  
          validateLabels(entity.getLabels());
  
          type.validateValueForUpdate(entity, entity.getTypeName(), messages);
  
          if (!messages.isEmpty()) {
              throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS, messages);
          }
  
          type.getNormalizedValueForUpdate(entity);
      }
  
      @Override
      public void cleanUp() {
          discoveryContext.cleanUp();
      }
  
      protected void discover() throws AtlasBaseException {
          MetricRecorder metric = RequestContext.get().startMetricRecord("walkEntityGraph");
  
          EntityStream entityStream = discoveryContext.getEntityStream();
  
          Set<String> walkedEntities = new HashSet<>();
  
          // walk through top-level entities and find entity references
          while (entityStream.hasNext()) {
              AtlasEntity entity = entityStream.next();
  
              if (entity == null) {
                  throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "found null entity");
              }
  
              processDynamicAttributes(entity);
  
              walkEntityGraph(entity);
  
              walkedEntities.add(entity.getGuid());
          }
  
          // walk through entities referenced by other entities
          // referencedGuids will be updated within this for() loop; avoid use of iterators
          List<String> referencedGuids = discoveryContext.getReferencedGuids();
  
          for (int i = 0; i < referencedGuids.size(); i++) {
              String guid = referencedGuids.get(i);
  
              if (walkedEntities.contains(guid)) {
                  continue;
              }
  
              AtlasEntity entity = entityStream.getByGuid(guid);
  
              if (entity != null) {
                  walkEntityGraph(entity);
  
                  walkedEntities.add(entity.getGuid());
              }
          }
  
          RequestContext.get().endMetricRecord(metric);
      }
  
      protected void resolveReferences() throws AtlasBaseException {
          MetricRecorder metric = RequestContext.get().startMetricRecord("resolveReferences");
  
          EntityResolver[] entityResolvers = new EntityResolver[] {new IDBasedEntityResolver(this.graph, typeRegistry, entityGraphMapper),
                  new UniqAttrBasedEntityResolver(this.graph, typeRegistry, entityGraphMapper)
          };
  
          for (EntityResolver resolver : entityResolvers) {
              resolver.resolveEntityReferences(discoveryContext);
          }
  
          RequestContext.get().endMetricRecord(metric);
      }
  
      void visitAttribute(AtlasType attrType, Object val) throws AtlasBaseException {
          if (attrType == null || val == null) {
              return;
          }
  
          switch (attrType.getTypeCategory()) {
              case PRIMITIVE:
              case ENUM:
                  return;
  
              case ARRAY: {
                  AtlasArrayType arrayType = (AtlasArrayType) attrType;
                  AtlasType      elemType  = arrayType.getElementType();
  
                  visitCollectionReferences(elemType, val);
              }
              break;
  
              case MAP: {
                  AtlasType keyType   = ((AtlasMapType) attrType).getKeyType();
                  AtlasType valueType = ((AtlasMapType) attrType).getValueType();
  
                  visitMapReferences(keyType, valueType, val);
              }
              break;
  
              case STRUCT:
                  visitStruct((AtlasStructType) attrType, val);
                  break;
  
              case OBJECT_ID_TYPE:
                  visitReference((AtlasObjectIdType) attrType, val);
                  break;
  
              default:
                  throw new AtlasBaseException(AtlasErrorCode.TYPE_CATEGORY_INVALID, attrType.getTypeCategory().name());
          }
      }
  
      void visitMapReferences(AtlasType keyType, AtlasType valueType, Object val) throws AtlasBaseException {
          if (keyType == null || valueType == null || val == null) {
              return;
          }
  
          if (isPrimitive(keyType.getTypeCategory()) && isPrimitive(valueType.getTypeCategory())) {
              return;
          }
  
          if (Map.class.isAssignableFrom(val.getClass())) {
              for (Map.Entry<?, ?> e : ((Map<?, ?>) val).entrySet()) {
                  visitAttribute(keyType, e.getKey());
                  visitAttribute(valueType, e.getValue());
              }
          }
      }
  
      void visitCollectionReferences(AtlasType elemType, Object val) throws AtlasBaseException {
          if (elemType == null || val == null || isPrimitive(elemType.getTypeCategory())) {
              return;
          }
  
          Iterator<?> it = null;
  
          if (val instanceof Collection) {
              it = ((Collection<?>) val).iterator();
          } else if (val instanceof Iterable) {
              it = ((Iterable<?>) val).iterator();
          } else if (val instanceof Iterator) {
              it = (Iterator<?>) val;
          }
  
          if (it != null) {
              while (it.hasNext()) {
                  Object elem = it.next();
                  visitAttribute(elemType, elem);
              }
          }
      }
  
      void visitStruct(AtlasStructType structType, Object val) throws AtlasBaseException {
          if (structType == null || val == null) {
              return;
          }
  
          final AtlasStruct struct;
  
          if (val instanceof AtlasStruct) {
              struct = (AtlasStruct) val;
          } else if (val instanceof Map) {
              Map<String, Object> attributes = AtlasTypeUtil.toStructAttributes((Map) val);
  
              struct = new AtlasStruct(structType.getTypeName(), attributes);
          } else {
              throw new AtlasBaseException(AtlasErrorCode.INVALID_STRUCT_VALUE, val.toString());
          }
  
          visitStruct(structType, struct);
      }
  
      void visitEntity(AtlasEntityType entityType, AtlasEntity entity) throws AtlasBaseException {
          List<String> visitedAttributes = new ArrayList<>();
  
          // visit relationship attributes
          visitRelationships(entityType, entity, visitedAttributes);
  
          // visit struct attributes
          for (AtlasAttribute attribute : entityType.getAllAttributes().values()) {
              AtlasType attrType = attribute.getAttributeType();
              String    attrName = attribute.getName();
              Object    attrVal  = entity.getAttribute(attrName);
  
              if (entity.hasAttribute(attrName) && !visitedAttributes.contains(attrName)) {
                  visitAttribute(attrType, attrVal);
              }
          }
      }
  
      void visitStruct(AtlasStructType structType, AtlasStruct struct) throws AtlasBaseException {
          for (AtlasAttribute attribute : structType.getAllAttributes().values()) {
              AtlasType attrType = attribute.getAttributeType();
              Object    attrVal  = struct.getAttribute(attribute.getName());
  
              visitAttribute(attrType, attrVal);
          }
      }
  
      void walkEntityGraph(AtlasEntity entity) throws AtlasBaseException {
          if (entity == null) {
              return;
          }
  
          AtlasEntityType type = typeRegistry.getEntityTypeByName(entity.getTypeName());
  
          if (type == null) {
              throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), entity.getTypeName());
          }
  
          recordObjectReference(entity.getGuid());
  
          visitEntity(type, entity);
      }
  
      boolean isPrimitive(TypeCategory typeCategory) {
          return typeCategory == TypeCategory.PRIMITIVE || typeCategory == TypeCategory.ENUM;
      }
  
      private void visitReference(AtlasObjectIdType type, Object val) throws AtlasBaseException {
          if (type == null || val == null) {
              return;
          }
  
          if (val instanceof AtlasObjectId) {
              AtlasObjectId objId = (AtlasObjectId) val;
  
              if (!AtlasTypeUtil.isValid(objId)) {
                  throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, objId.toString());
              }
  
              recordObjectReference(objId);
          } else if (val instanceof Map) {
              AtlasObjectId objId = new AtlasObjectId((Map) val);
  
              if (!AtlasTypeUtil.isValid(objId)) {
                  throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, objId.toString());
              }
  
              recordObjectReference(objId);
          } else {
              throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, val.toString());
          }
      }
  
      private void visitRelationships(AtlasEntityType entityType, AtlasEntity entity, List<String> visitedAttributes) throws AtlasBaseException {
          for (String attrName : entityType.getRelationshipAttributes().keySet()) {
              // if attribute is not in 'relationshipAttributes', try 'attributes'
              if (entity.hasRelationshipAttribute(attrName)) {
                  Object         attrVal          = entity.getRelationshipAttribute(attrName);
                  String         relationshipType = AtlasEntityUtil.getRelationshipType(attrVal);
                  AtlasAttribute attribute        = entityType.getRelationshipAttribute(attrName, relationshipType);
  
                  visitAttribute(attribute.getAttributeType(), attrVal);
  
                  visitedAttributes.add(attrName);
              } else if (entity.hasAttribute(attrName)) {
                  Object         attrVal          = entity.getAttribute(attrName);
                  String         relationshipType = AtlasEntityUtil.getRelationshipType(attrVal);
                  AtlasAttribute attribute        = entityType.getRelationshipAttribute(attrName, relationshipType);
  
                  visitAttribute(attribute.getAttributeType(), attrVal);
  
                  visitedAttributes.add(attrName);
              }
          }
      }
  
      private void recordObjectReference(String guid) {
          discoveryContext.addReferencedGuid(guid);
      }
  
      private void recordObjectReference(AtlasObjectId objId) {
          if (AtlasTypeUtil.isValidGuid(objId)) {
              discoveryContext.addReferencedGuid(objId.getGuid());
          } else {
              discoveryContext.addReferencedByUniqAttribs(objId);
          }
      }
  
      private void processDynamicAttributes(AtlasEntity entity) throws AtlasBaseException {
          AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entity.getTypeName());
  
          if (entityType == null) {
              throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), entity.getTypeName());
          }
  
          for (AtlasAttribute attribute : entityType.getDynEvalAttributes()) {
              String              attributeName = attribute.getName();
              List<TemplateToken> tokens        = entityType.getParsedTemplates().get(attributeName);
  
              if (tokens == null) {
                  continue;
              }
  
              StringBuilder dynAttributeValue = new StringBuilder();
  
              boolean set = true;
  
              for (TemplateToken token : tokens) {
                  String evaluated = token.eval(entity);
  
                  if (evaluated != null) {
                      dynAttributeValue.append(evaluated);
                  } else {
                      set = false;
  
                      LOG.warn("Attribute {} for {} unable to be generated because of dynamic attribute token {}", attributeName, entityType, token.getValue());
  
                      break;
                  }
              }
  
              if (set) {
                  entity.setAttribute(attributeName, dynAttributeValue.toString());
              }
          }
      }
  }
  