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.discovery;
  
  import org.apache.atlas.SortOrder;
  import org.apache.atlas.model.discovery.SearchParameters.FilterCriteria;
  import org.apache.atlas.repository.Constants;
  import org.apache.atlas.repository.graphdb.AtlasGraphQuery;
  import org.apache.atlas.repository.graphdb.AtlasIndexQuery;
  import org.apache.atlas.repository.graphdb.AtlasVertex;
  import org.apache.atlas.type.AtlasClassificationType;
  import org.apache.atlas.type.AtlasEntityType;
  import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
  import org.apache.atlas.util.SearchPredicateUtil;
  import org.apache.atlas.utils.AtlasPerfTracer;
  import org.apache.commons.collections.CollectionUtils;
  import org.apache.commons.collections.Predicate;
  import org.apache.commons.collections.PredicateUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.ArrayList;
  import java.util.HashSet;
  import java.util.Iterator;
  import java.util.LinkedHashMap;
  import java.util.LinkedList;
  import java.util.List;
  import java.util.Set;
  import java.util.stream.StreamSupport;
  
  import static org.apache.atlas.discovery.SearchContext.MATCH_ALL_CLASSIFICATION_TYPES;
  import static org.apache.atlas.discovery.SearchContext.MATCH_ALL_CLASSIFIED;
  import static org.apache.atlas.discovery.SearchContext.MATCH_ALL_NOT_CLASSIFIED;
  import static org.apache.atlas.discovery.SearchContext.MATCH_ALL_WILDCARD_CLASSIFICATION;
  import static org.apache.atlas.discovery.SearchContext.MarkerUtil;
  import static org.apache.atlas.repository.Constants.PROPAGATED_TRAIT_NAMES_PROPERTY_KEY;
  import static org.apache.atlas.repository.Constants.TRAIT_NAMES_PROPERTY_KEY;
  import static org.apache.atlas.repository.Constants.TYPE_NAME_PROPERTY_KEY;
  import static org.apache.atlas.repository.graphdb.AtlasGraphQuery.ComparisionOperator.EQUAL;
  import static org.apache.atlas.repository.graphdb.AtlasGraphQuery.ComparisionOperator.NOT_EQUAL;
  import static org.apache.atlas.repository.graphdb.AtlasGraphQuery.SortOrder.ASC;
  import static org.apache.atlas.repository.graphdb.AtlasGraphQuery.SortOrder.DESC;
  
  public class EntitySearchProcessor extends SearchProcessor {
      private static final Logger LOG      = LoggerFactory.getLogger(EntitySearchProcessor.class);
      private static final Logger PERF_LOG = AtlasPerfTracer.getPerfLogger("EntitySearchProcessor");
  
      private final AtlasIndexQuery indexQuery;
      private final AtlasGraphQuery graphQuery;
      private       Predicate       graphQueryPredicate;
      private       Predicate       filterGraphQueryPredicate;
  
      public EntitySearchProcessor(SearchContext context) {
          super(context);
  
          final Set<AtlasEntityType> entityTypes           = context.getEntityTypes();
          final FilterCriteria       filterCriteria        = context.getSearchParameters().getEntityFilters();
          final Set<String>          indexAttributes       = new HashSet<>();
          final Set<String>          graphAttributes       = new HashSet<>();
          final Set<String>          allAttributes         = new HashSet<>();
          final Set<String>          typeAndSubTypes       = context.getEntityTypeNames();
          final String               typeAndSubTypesQryStr = context.getEntityTypesQryStr();
          final String               sortBy                = context.getSearchParameters().getSortBy();
          final SortOrder            sortOrder             = context.getSearchParameters().getSortOrder();
  
          final Set<AtlasClassificationType> classificationTypes           = context.getClassificationTypes();
          final Set<String>                  classificationTypeAndSubTypes = context.getClassificationTypeNames();
          final boolean                      filterClassification;
  
          if (CollectionUtils.isNotEmpty(classificationTypes)) {
              filterClassification = !context.needClassificationProcessor();
          } else {
              filterClassification = false;
          }
  
          final Predicate typeNamePredicate;
          final Predicate traitPredicate  = buildTraitPredict(classificationTypes);
          final Predicate activePredicate = SearchPredicateUtil.getEQPredicateGenerator().generatePredicate(Constants.STATE_PROPERTY_KEY, "ACTIVE", String.class);
  
          if (!isEntityRootType()) {
              typeNamePredicate = SearchPredicateUtil.getINPredicateGenerator().generatePredicate(TYPE_NAME_PROPERTY_KEY, typeAndSubTypes, String.class);
          } else {
              typeNamePredicate = SearchPredicateUtil.generateIsEntityVertexPredicate(context.getTypeRegistry());
          }
  
          processSearchAttributes(entityTypes, filterCriteria, indexAttributes, graphAttributes, allAttributes);
  
          final boolean typeSearchByIndex = !filterClassification && typeAndSubTypesQryStr.length() <= MAX_QUERY_STR_LENGTH_TYPES;
          final boolean attrSearchByIndex = !filterClassification && CollectionUtils.isNotEmpty(indexAttributes) && canApplyIndexFilter(entityTypes, filterCriteria, false);
  
          StringBuilder indexQuery = new StringBuilder();
  
          // TypeName check to be done in-memory as well to address ATLAS-2121 (case sensitivity)
          inMemoryPredicate = typeNamePredicate;
  
          if (typeSearchByIndex) {
              graphIndexQueryBuilder.addTypeAndSubTypesQueryFilter(indexQuery, typeAndSubTypesQryStr);
          }
  
          if (attrSearchByIndex) {
              constructFilterQuery(indexQuery, entityTypes, filterCriteria, indexAttributes);
  
              Predicate attributePredicate = constructInMemoryPredicate(entityTypes, filterCriteria, indexAttributes);
  
              if (attributePredicate != null) {
                  inMemoryPredicate = PredicateUtils.andPredicate(inMemoryPredicate, attributePredicate);
              }
          } else {
              graphAttributes.addAll(indexAttributes);
          }
  
          if (indexQuery.length() > 0) {
              graphIndexQueryBuilder.addActiveStateQueryFilter(indexQuery);
  
              String indexQueryString = STRAY_AND_PATTERN.matcher(indexQuery).replaceAll(")");
  
              indexQueryString = STRAY_OR_PATTERN.matcher(indexQueryString).replaceAll(")");
              indexQueryString = STRAY_ELIPSIS_PATTERN.matcher(indexQueryString).replaceAll("");
  
              this.indexQuery = context.getGraph().indexQuery(Constants.VERTEX_INDEX, indexQueryString);
          } else {
              this.indexQuery = null;
          }
  
          if (CollectionUtils.isNotEmpty(graphAttributes) || !typeSearchByIndex) {
              AtlasGraphQuery query = context.getGraph().query();
  
              if (!typeSearchByIndex && !isEntityRootType()) {
                  query.in(TYPE_NAME_PROPERTY_KEY, typeAndSubTypes);
              }
  
              // If we need to filter on the trait names then we need to build the query and equivalent in-memory predicate
              if (filterClassification) {
                  AtlasClassificationType classificationType = classificationTypes.iterator().next();
                  List<AtlasGraphQuery>   orConditions       = new LinkedList<>();
  
                  if (classificationType == MATCH_ALL_WILDCARD_CLASSIFICATION || classificationType == MATCH_ALL_CLASSIFIED || classificationType == MATCH_ALL_CLASSIFICATION_TYPES) {
                      orConditions.add(query.createChildQuery().has(TRAIT_NAMES_PROPERTY_KEY, NOT_EQUAL, null));
                      orConditions.add(query.createChildQuery().has(PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, NOT_EQUAL, null));
                  } else if (classificationType == MATCH_ALL_NOT_CLASSIFIED) {
                      orConditions.add(query.createChildQuery().has(TRAIT_NAMES_PROPERTY_KEY, EQUAL, null)
                              .has(PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, EQUAL, null));
                  } else {
                      orConditions.add(query.createChildQuery().in(TRAIT_NAMES_PROPERTY_KEY, classificationTypeAndSubTypes));
                      orConditions.add(query.createChildQuery().in(PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, classificationTypeAndSubTypes));
                  }
  
                  query.or(orConditions);
  
                  // Construct a parallel in-memory predicate
                  if (isEntityRootType()) {
                      inMemoryPredicate = typeNamePredicate;
                  }
  
                  if (graphQueryPredicate != null) {
                      graphQueryPredicate = PredicateUtils.andPredicate(graphQueryPredicate, traitPredicate);
                  } else {
                      graphQueryPredicate = traitPredicate;
                  }
              }
  
              graphQuery = toGraphFilterQuery(entityTypes, filterCriteria, graphAttributes, query);
  
              // Prepare in-memory predicate for attribute filtering
              Predicate attributePredicate = constructInMemoryPredicate(entityTypes, filterCriteria, graphAttributes);
  
              if (attributePredicate != null) {
                  if (graphQueryPredicate != null) {
                      graphQueryPredicate = PredicateUtils.andPredicate(graphQueryPredicate, attributePredicate);
                  } else {
                      graphQueryPredicate = attributePredicate;
                  }
              }
  
              // Filter condition for the STATUS
              if (context.getSearchParameters().getExcludeDeletedEntities() && this.indexQuery == null) {
                  graphQuery.has(Constants.STATE_PROPERTY_KEY, "ACTIVE");
  
                  if (graphQueryPredicate != null) {
                      graphQueryPredicate = PredicateUtils.andPredicate(graphQueryPredicate, activePredicate);
                  } else {
                      graphQueryPredicate = activePredicate;
                  }
              }
              if (sortBy != null && !sortBy.isEmpty()) {
                  final AtlasEntityType entityType      = context.getEntityTypes().iterator().next();
                  AtlasAttribute        sortByAttribute = entityType.getAttribute(sortBy);
  
                  if (sortByAttribute != null) {
                      AtlasGraphQuery.SortOrder qrySortOrder = sortOrder == SortOrder.ASCENDING ? ASC : DESC;
  
                      graphQuery.orderBy(sortByAttribute.getVertexPropertyName(), qrySortOrder);
                  }
              }
          } else {
              graphQuery          = null;
              graphQueryPredicate = null;
          }
  
          // Prepare the graph query and in-memory filter for the filtering phase
          filterGraphQueryPredicate = typeNamePredicate;
  
          Predicate attributesPredicate = constructInMemoryPredicate(entityTypes, filterCriteria, allAttributes);
  
          if (attributesPredicate != null) {
              filterGraphQueryPredicate = filterGraphQueryPredicate == null ? attributesPredicate : PredicateUtils.andPredicate(filterGraphQueryPredicate, attributesPredicate);
          }
  
          if (filterClassification) {
              filterGraphQueryPredicate = filterGraphQueryPredicate == null ? traitPredicate : PredicateUtils.andPredicate(filterGraphQueryPredicate, traitPredicate);
          }
  
          // Filter condition for the STATUS
          if (context.getSearchParameters().getExcludeDeletedEntities()) {
              filterGraphQueryPredicate = filterGraphQueryPredicate == null ? activePredicate : PredicateUtils.andPredicate(filterGraphQueryPredicate, activePredicate);
          }
      }
  
      @Override
      public List<AtlasVertex> execute() {
          LOG.debug("==> EntitySearchProcessor.execute({})", context);
  
          List<AtlasVertex> ret  = new ArrayList<>();
          AtlasPerfTracer   perf = null;
  
          if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
              perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntitySearchProcessor.execute(" + context + ")");
          }
  
          try {
              final int     limit    = context.getSearchParameters().getLimit();
              final Integer marker   = context.getMarker();
              final int     startIdx = marker != null ? marker : context.getSearchParameters().getOffset();
  
              // when subsequent filtering stages are involved, query should start at 0 even though startIdx can be higher
              //
              // first 'startIdx' number of entries will be ignored
              // if marker is provided, start query with marker offset
              int qryOffset;
  
              if (marker != null) {
                  qryOffset = marker;
              } else {
                  qryOffset = (nextProcessor != null || (graphQuery != null && indexQuery != null) || isEntityRootType()) ? 0 : startIdx;
              }
  
              int                                 resultIdx             = qryOffset;
              LinkedHashMap<Integer, AtlasVertex> offsetEntityVertexMap = new LinkedHashMap<>();
  
              for (; ret.size() < limit; qryOffset += limit) {
                  offsetEntityVertexMap.clear();
  
                  if (context.terminateSearch()) {
                      LOG.warn("query terminated: {}", context.getSearchParameters());
  
                      break;
                  }
  
                  final boolean isLastResultPage;
  
                  if (indexQuery != null) {
                      Iterator<AtlasIndexQuery.Result> idxQueryResult = executeIndexQuery(context, indexQuery, qryOffset, limit);
  
                      offsetEntityVertexMap = getVerticesFromIndexQueryResult(idxQueryResult, offsetEntityVertexMap, qryOffset);
                  } else {
                      Iterator<AtlasVertex> queryResult = graphQuery.vertices(qryOffset, limit).iterator();
  
                      offsetEntityVertexMap = getVertices(queryResult, offsetEntityVertexMap, qryOffset);
                  }
  
                  isLastResultPage = offsetEntityVertexMap.size() < limit;
                  // Do in-memory filtering
                  offsetEntityVertexMap = super.filter(offsetEntityVertexMap, inMemoryPredicate);
  
                  //incase when operator is NEQ in pipeSeperatedSystemAttributes
                  if (graphQueryPredicate != null) {
                      offsetEntityVertexMap = super.filter(offsetEntityVertexMap, graphQueryPredicate);
                  }
  
                  offsetEntityVertexMap = super.filter(offsetEntityVertexMap);
  
                  resultIdx = collectResultVertices(ret, startIdx, limit, resultIdx, offsetEntityVertexMap, marker);
  
                  if (isLastResultPage) {
                      resultIdx = MarkerUtil.MARKER_END - 1;
                      break;
                  }
              }
  
              if (marker != null) {
                  nextOffset = resultIdx + 1;
              }
          } finally {
              AtlasPerfTracer.log(perf);
          }
  
          LOG.debug("<== EntitySearchProcessor.execute({}): ret.size()={}", context, ret.size());
  
          return ret;
      }
  
      @Override
      public long getResultCount() {
          if (indexQuery != null) {
              return indexQuery.vertexTotals();
          } else if (graphQuery != null) {
              return StreamSupport.stream(graphQuery.vertexIds().spliterator(), false).count();
          } else {
              return -1L;
          }
      }
  
      @Override
      public LinkedHashMap<Integer, AtlasVertex> filter(LinkedHashMap<Integer, AtlasVertex> offsetEntityVertexMap) {
          LOG.debug("==> EntitySearchProcessor.filter({})", offsetEntityVertexMap.size());
  
          // Since we already have the entity vertices, a in-memory filter will be faster than fetching the same
          // vertices again with the required filtering
          if (filterGraphQueryPredicate != null) {
              LOG.debug("Filtering in-memory");
  
              offsetEntityVertexMap = super.filter(offsetEntityVertexMap, filterGraphQueryPredicate);
          }
  
          offsetEntityVertexMap = super.filter(offsetEntityVertexMap);
  
          LOG.debug("<== EntitySearchProcessor.filter(): ret.size()={}", offsetEntityVertexMap.size());
  
          return offsetEntityVertexMap;
      }
  }
  