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.AtlasConfiguration;
  import org.apache.atlas.AtlasErrorCode;
  import org.apache.atlas.annotation.GraphTransaction;
  import org.apache.atlas.authorize.AtlasAuthorizationUtils;
  import org.apache.atlas.authorize.AtlasEntityAccessRequest;
  import org.apache.atlas.authorize.AtlasPrivilege;
  import org.apache.atlas.exception.AtlasBaseException;
  import org.apache.atlas.model.instance.AtlasEntity;
  import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
  import org.apache.atlas.model.instance.AtlasEntityHeader;
  import org.apache.atlas.model.instance.AtlasObjectId;
  import org.apache.atlas.model.lineage.AtlasLineageInfo;
  import org.apache.atlas.model.lineage.AtlasLineageInfo.LineageDirection;
  import org.apache.atlas.model.lineage.AtlasLineageInfo.LineageInfoOnDemand;
  import org.apache.atlas.model.lineage.AtlasLineageInfo.LineageRelation;
  import org.apache.atlas.model.lineage.LineageOnDemandConstraints;
  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.AtlasVertex;
  import org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2;
  import org.apache.atlas.repository.store.graph.v2.EntityGraphRetriever;
  import org.apache.atlas.type.AtlasEntityType;
  import org.apache.atlas.type.AtlasTypeRegistry;
  import org.apache.atlas.type.AtlasTypeUtil;
  import org.apache.atlas.util.AtlasGremlinQueryProvider;
  import org.apache.atlas.v1.model.lineage.SchemaResponse.SchemaDetails;
  import org.apache.commons.collections.CollectionUtils;
  import org.apache.commons.collections.MapUtils;
  import org.apache.commons.lang.StringUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import org.springframework.stereotype.Service;
  
  import javax.inject.Inject;
  import javax.script.ScriptEngine;
  import javax.script.ScriptException;
  
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Set;
  import java.util.stream.Collectors;
  
  import static org.apache.atlas.AtlasClient.DATA_SET_SUPER_TYPE;
  import static org.apache.atlas.AtlasClient.PROCESS_SUPER_TYPE;
  import static org.apache.atlas.AtlasErrorCode.INSTANCE_LINEAGE_QUERY_FAILED;
  import static org.apache.atlas.model.lineage.AtlasLineageInfo.LineageDirection.BOTH;
  import static org.apache.atlas.model.lineage.AtlasLineageInfo.LineageDirection.INPUT;
  import static org.apache.atlas.model.lineage.AtlasLineageInfo.LineageDirection.OUTPUT;
  import static org.apache.atlas.repository.Constants.RELATIONSHIP_GUID_PROPERTY_KEY;
  import static org.apache.atlas.repository.graphdb.AtlasEdgeDirection.IN;
  import static org.apache.atlas.repository.graphdb.AtlasEdgeDirection.OUT;
  import static org.apache.atlas.util.AtlasGremlinQueryProvider.AtlasGremlinQuery.FULL_LINEAGE_DATASET;
  import static org.apache.atlas.util.AtlasGremlinQueryProvider.AtlasGremlinQuery.FULL_LINEAGE_PROCESS;
  import static org.apache.atlas.util.AtlasGremlinQueryProvider.AtlasGremlinQuery.PARTIAL_LINEAGE_DATASET;
  import static org.apache.atlas.util.AtlasGremlinQueryProvider.AtlasGremlinQuery.PARTIAL_LINEAGE_PROCESS;
  
  @Service
  public class EntityLineageService implements AtlasLineageService {
      private static final Logger LOG = LoggerFactory.getLogger(EntityLineageService.class);
  
      private static final String  PROCESS_INPUTS_EDGE                  = "__Process.inputs";
      private static final String  PROCESS_OUTPUTS_EDGE                 = "__Process.outputs";
      private static final String  COLUMNS                              = "columns";
      private static final boolean LINEAGE_USING_GREMLIN                = AtlasConfiguration.LINEAGE_USING_GREMLIN.getBoolean();
      private static final Integer DEFAULT_LINEAGE_MAX_NODE_COUNT       = 9000;
      private static final int     LINEAGE_ON_DEMAND_DEFAULT_DEPTH      = 3;
      private static final int     LINEAGE_ON_DEMAND_DEFAULT_NODE_COUNT = AtlasConfiguration.LINEAGE_ON_DEMAND_DEFAULT_NODE_COUNT.getInt();
      private static final String  SEPARATOR                            = "->";
  
      private final AtlasGraph                graph;
      private final AtlasGremlinQueryProvider gremlinQueryProvider;
      private final EntityGraphRetriever      entityRetriever;
      private final AtlasTypeRegistry         atlasTypeRegistry;
  
      @Inject
      EntityLineageService(AtlasTypeRegistry typeRegistry, AtlasGraph atlasGraph) {
          this.graph                = atlasGraph;
          this.gremlinQueryProvider = AtlasGremlinQueryProvider.getInstance();
          this.entityRetriever      = new EntityGraphRetriever(atlasGraph, typeRegistry);
          this.atlasTypeRegistry    = typeRegistry;
      }
  
      @Override
      @GraphTransaction
      public AtlasLineageInfo getAtlasLineageInfo(String guid, LineageDirection direction, int depth) throws AtlasBaseException {
          AtlasLineageInfo ret;
          boolean          isDataSet = validateEntityTypeAndCheckIfDataSet(guid);
  
          if (LINEAGE_USING_GREMLIN) {
              ret = getLineageInfoV1(guid, direction, depth, isDataSet);
          } else {
              ret = getLineageInfoV2(guid, direction, depth, isDataSet);
          }
  
          return ret;
      }
  
      @Override
      @GraphTransaction
      public AtlasLineageInfo getAtlasLineageInfo(String guid, Map<String, LineageOnDemandConstraints> lineageConstraintsMap) throws AtlasBaseException {
          if (MapUtils.isEmpty(lineageConstraintsMap)) {
              lineageConstraintsMap = new HashMap<>();
  
              lineageConstraintsMap.put(guid, getDefaultLineageConstraints(guid));
          }
  
          boolean          isDataSet = validateEntityTypeAndCheckIfDataSet(guid);
          AtlasLineageInfo ret       = getLineageInfoOnDemand(guid, lineageConstraintsMap, isDataSet);
  
          appendLineageOnDemandPayload(ret, lineageConstraintsMap);
  
          // filtering out on-demand relations which has input & output nodes within the limit
          cleanupRelationsOnDemand(ret);
  
          return ret;
      }
  
      @Override
      @GraphTransaction
      public SchemaDetails getSchemaForHiveTableByName(final String datasetName) throws AtlasBaseException {
          if (StringUtils.isEmpty(datasetName)) {
              // TODO: Complete error handling here
              throw new AtlasBaseException(AtlasErrorCode.BAD_REQUEST);
          }
  
          AtlasEntityType     hiveTable        = atlasTypeRegistry.getEntityTypeByName("hive_table");
          Map<String, Object> lookupAttributes = new HashMap<>();
  
          lookupAttributes.put("qualifiedName", datasetName);
  
          String guid = AtlasGraphUtilsV2.getGuidByUniqueAttributes(hiveTable, lookupAttributes);
  
          return getSchemaForHiveTableByGuid(guid);
      }
  
      @Override
      @GraphTransaction
      public SchemaDetails getSchemaForHiveTableByGuid(final String guid) throws AtlasBaseException {
          if (StringUtils.isEmpty(guid)) {
              throw new AtlasBaseException(AtlasErrorCode.BAD_REQUEST);
          }
  
          SchemaDetails   ret        = new SchemaDetails();
          AtlasEntityType hiveColumn = atlasTypeRegistry.getEntityTypeByName("hive_column");
  
          ret.setDataType(AtlasTypeUtil.toClassTypeDefinition(hiveColumn));
  
          AtlasEntityWithExtInfo entityWithExtInfo = entityRetriever.toAtlasEntityWithExtInfo(guid);
          AtlasEntity            entity            = entityWithExtInfo.getEntity();
  
          AtlasAuthorizationUtils.verifyAccess(new AtlasEntityAccessRequest(atlasTypeRegistry, AtlasPrivilege.ENTITY_READ, new AtlasEntityHeader(entity)), "read entity schema: guid=", guid);
  
          Map<String, AtlasEntity> referredEntities = entityWithExtInfo.getReferredEntities();
          List<String>             columnIds        = getColumnIds(entity);
  
          if (MapUtils.isNotEmpty(referredEntities)) {
              List<Map<String, Object>> rows = referredEntities.entrySet()
                      .stream()
                      .filter(e -> isColumn(columnIds, e))
                      .map(e -> AtlasTypeUtil.toMap(e.getValue()))
                      .collect(Collectors.toList());
              ret.setRows(rows);
          }
  
          return ret;
      }
  
      public boolean isLineageOnDemandEnabled() {
          return AtlasConfiguration.LINEAGE_ON_DEMAND_ENABLED.getBoolean();
      }
  
      private boolean validateEntityTypeAndCheckIfDataSet(String guid) throws AtlasBaseException {
          AtlasEntityHeader entity = entityRetriever.toAtlasEntityHeaderWithClassifications(guid);
  
          AtlasAuthorizationUtils.verifyAccess(new AtlasEntityAccessRequest(atlasTypeRegistry, AtlasPrivilege.ENTITY_READ, entity), "read entity lineage: guid=", guid);
  
          AtlasEntityType entityType = atlasTypeRegistry.getEntityTypeByName(entity.getTypeName());
  
          if (entityType == null) {
              throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_NOT_FOUND, entity.getTypeName());
          }
  
          boolean isDataSet = entityType.getTypeAndAllSuperTypes().contains(DATA_SET_SUPER_TYPE);
  
          if (!isDataSet) {
              boolean isProcess = entityType.getTypeAndAllSuperTypes().contains(PROCESS_SUPER_TYPE);
  
              if (!isProcess) {
                  throw new AtlasBaseException(AtlasErrorCode.INVALID_LINEAGE_ENTITY_TYPE, guid, entity.getTypeName());
              }
          }
  
          return isDataSet;
      }
  
      private void appendLineageOnDemandPayload(AtlasLineageInfo lineageInfo, Map<String, LineageOnDemandConstraints> lineageConstraintsMap) {
          if (lineageInfo == null || MapUtils.isEmpty(lineageConstraintsMap)) {
              return;
          }
  
          lineageInfo.setLineageOnDemandPayload(lineageConstraintsMap);
      }
  
      //Consider only relationsOnDemand which has either more inputs or more outputs than given limit
      private void cleanupRelationsOnDemand(AtlasLineageInfo lineageInfo) {
          if (lineageInfo != null && MapUtils.isNotEmpty(lineageInfo.getRelationsOnDemand())) {
              lineageInfo.getRelationsOnDemand().entrySet().removeIf(x -> !(x.getValue().hasMoreInputs() || x.getValue().hasMoreOutputs()));
          }
      }
  
      private List<String> getColumnIds(AtlasEntity entity) {
          List<String> ret        = new ArrayList<>();
          Object       columnObjs = entity.getAttribute(COLUMNS);
  
          if (columnObjs instanceof List) {
              for (Object pkObj : (List<?>) columnObjs) {
                  if (pkObj instanceof AtlasObjectId) {
                      ret.add(((AtlasObjectId) pkObj).getGuid());
                  }
              }
          }
  
          return ret;
      }
  
      private boolean isColumn(List<String> columnIds, Map.Entry<String, AtlasEntity> e) {
          return columnIds.contains(e.getValue().getGuid());
      }
  
      private AtlasLineageInfo getLineageInfoV1(String guid, LineageDirection direction, int depth, boolean isDataSet) throws AtlasBaseException {
          AtlasLineageInfo ret;
  
          if (direction.equals(INPUT)) {
              ret = getLineageInfo(guid, INPUT, depth, isDataSet);
          } else if (direction.equals(OUTPUT)) {
              ret = getLineageInfo(guid, OUTPUT, depth, isDataSet);
          } else {
              ret = getBothLineageInfoV1(guid, depth, isDataSet);
          }
  
          return ret;
      }
  
      private AtlasLineageInfo getLineageInfo(String guid, LineageDirection direction, int depth, boolean isDataSet) throws AtlasBaseException {
          final Map<String, Object>      bindings     = new HashMap<>();
          String                         lineageQuery = getLineageQuery(guid, direction, depth, isDataSet, bindings);
          List<?>                        results      = executeGremlinScript(bindings, lineageQuery);
          Map<String, AtlasEntityHeader> entities     = new HashMap<>();
          Set<LineageRelation>           relations    = new HashSet<>();
  
          if (CollectionUtils.isNotEmpty(results)) {
              for (Object result : results) {
                  if (result instanceof Map) {
                      for (final Map.Entry<?, ?> entry : ((Map<?, ?>) result).entrySet()) {
                          Object value = entry.getValue();
  
                          if (value instanceof List) {
                              for (Object elem : (List<?>) value) {
                                  if (elem instanceof AtlasEdge) {
                                      processEdge((AtlasEdge) elem, entities, relations);
                                  } else {
                                      LOG.warn("Invalid value of type {} found, ignoring", (elem != null ? elem.getClass().getSimpleName() : "null"));
                                  }
                              }
                          } else if (value instanceof AtlasEdge) {
                              processEdge((AtlasEdge) value, entities, relations);
                          } else {
                              LOG.warn("Invalid value of type {} found, ignoring", (value != null ? value.getClass().getSimpleName() : "null"));
                          }
                      }
                  } else if (result instanceof AtlasEdge) {
                      processEdge((AtlasEdge) result, entities, relations);
                  }
              }
          }
  
          return new AtlasLineageInfo(guid, entities, relations, direction, depth);
      }
  
      private AtlasLineageInfo getLineageInfoV2(String guid, LineageDirection direction, int depth, boolean isDataSet) throws AtlasBaseException {
          AtlasLineageInfo ret = initializeLineageInfo(guid, direction, depth);
  
          if (depth == 0) {
              depth = -1;
          }
  
          if (isDataSet) {
              AtlasVertex datasetVertex = AtlasGraphUtilsV2.findByGuid(this.graph, guid);
  
              if (direction == INPUT || direction == BOTH) {
                  traverseEdges(datasetVertex, true, depth, ret);
              }
  
              if (direction == OUTPUT || direction == BOTH) {
                  traverseEdges(datasetVertex, false, depth, ret);
              }
          } else {
              AtlasVertex processVertex = AtlasGraphUtilsV2.findByGuid(this.graph, guid);
  
              // make one hop to the next dataset vertices from process vertex and traverse with 'depth = depth - 1'
              if (direction == INPUT || direction == BOTH) {
                  Iterable<AtlasEdge> processEdges = processVertex.getEdges(AtlasEdgeDirection.OUT, PROCESS_INPUTS_EDGE);
  
                  for (AtlasEdge processEdge : processEdges) {
                      addEdgeToResult(processEdge, ret);
  
                      AtlasVertex datasetVertex = processEdge.getInVertex();
  
                      traverseEdges(datasetVertex, true, depth - 1, ret);
                  }
              }
  
              if (direction == OUTPUT || direction == BOTH) {
                  Iterable<AtlasEdge> processEdges = processVertex.getEdges(AtlasEdgeDirection.OUT, PROCESS_OUTPUTS_EDGE);
  
                  for (AtlasEdge processEdge : processEdges) {
                      addEdgeToResult(processEdge, ret);
  
                      AtlasVertex datasetVertex = processEdge.getInVertex();
  
                      traverseEdges(datasetVertex, false, depth - 1, ret);
                  }
              }
          }
  
          return ret;
      }
  
      private LineageOnDemandConstraints getDefaultLineageConstraints(String guid) {
          LOG.debug("No lineage on-demand constraints provided for guid: {}, configuring with default values direction: {}, inputRelationsLimit: {}, outputRelationsLimit: {}, depth: {}",
                  guid, BOTH, LINEAGE_ON_DEMAND_DEFAULT_NODE_COUNT, LINEAGE_ON_DEMAND_DEFAULT_NODE_COUNT, LINEAGE_ON_DEMAND_DEFAULT_DEPTH);
  
          return new LineageOnDemandConstraints();
      }
  
      private LineageOnDemandConstraints getAndValidateLineageConstraintsByGuid(String guid, Map<String, LineageOnDemandConstraints> lineageConstraintsMap) {
          if (lineageConstraintsMap == null || !lineageConstraintsMap.containsKey(guid)) {
              return getDefaultLineageConstraints(guid);
          }
  
          LineageOnDemandConstraints lineageConstraintsByGuid = lineageConstraintsMap.get(guid);
  
          if (lineageConstraintsByGuid == null) {
              return getDefaultLineageConstraints(guid);
          }
  
          if (Objects.isNull(lineageConstraintsByGuid.getDirection())) {
              LOG.info("No lineage on-demand direction provided for guid: {}, configuring with default value {}", guid, LineageDirection.BOTH);
  
              lineageConstraintsByGuid.setDirection(BOTH);
          }
  
          if (lineageConstraintsByGuid.getInputRelationsLimit() == 0) {
              LOG.info("No lineage on-demand inputRelationsLimit provided for guid: {}, configuring with default value {}", guid, LINEAGE_ON_DEMAND_DEFAULT_NODE_COUNT);
  
              lineageConstraintsByGuid.setInputRelationsLimit(LINEAGE_ON_DEMAND_DEFAULT_NODE_COUNT);
          }
  
          if (lineageConstraintsByGuid.getOutputRelationsLimit() == 0) {
              LOG.info("No lineage on-demand outputRelationsLimit provided for guid: {}, configuring with default value {}", guid, LINEAGE_ON_DEMAND_DEFAULT_NODE_COUNT);
  
              lineageConstraintsByGuid.setOutputRelationsLimit(LINEAGE_ON_DEMAND_DEFAULT_NODE_COUNT);
          }
  
          if (lineageConstraintsByGuid.getDepth() == 0) {
              LOG.info("No lineage on-demand depth provided for guid: {}, configuring with default value {}", guid, LINEAGE_ON_DEMAND_DEFAULT_DEPTH);
  
              lineageConstraintsByGuid.setDepth(LINEAGE_ON_DEMAND_DEFAULT_DEPTH);
          }
  
          return lineageConstraintsByGuid;
      }
  
      private AtlasLineageInfo getLineageInfoOnDemand(String guid, Map<String, LineageOnDemandConstraints> lineageConstraintsMap, boolean isDataSet) throws AtlasBaseException {
          LineageOnDemandConstraints lineageConstraintsByGuid = getAndValidateLineageConstraintsByGuid(guid, lineageConstraintsMap);
  
          if (lineageConstraintsByGuid == null) {
              throw new AtlasBaseException(AtlasErrorCode.NO_LINEAGE_CONSTRAINTS_FOR_GUID, guid);
          }
  
          LineageDirection direction = lineageConstraintsByGuid.getDirection();
          int              depth     = lineageConstraintsByGuid.getDepth();
  
          AtlasLineageInfo ret = initializeLineageInfo(guid, direction, depth);
  
          if (depth == 0) {
              depth = -1;
          }
  
          if (isDataSet) {
              AtlasVertex datasetVertex = AtlasGraphUtilsV2.findByGuid(this.graph, guid);
  
              if (!ret.getRelationsOnDemand().containsKey(guid)) {
                  ret.getRelationsOnDemand().put(guid, new LineageInfoOnDemand(lineageConstraintsByGuid));
              }
  
              if (direction == INPUT || direction == BOTH) {
                  traverseEdgesOnDemand(datasetVertex, true, depth, new HashSet<>(), lineageConstraintsMap, ret);
              }
  
              if (direction == OUTPUT || direction == BOTH) {
                  traverseEdgesOnDemand(datasetVertex, false, depth, new HashSet<>(), lineageConstraintsMap, ret);
              }
          } else {
              AtlasVertex processVertex = AtlasGraphUtilsV2.findByGuid(this.graph, guid);
  
              // make one hop to the next dataset vertices from process vertex and traverse with 'depth = depth - 1'
              if (direction == INPUT || direction == BOTH) {
                  Iterable<AtlasEdge> processEdges = processVertex.getEdges(AtlasEdgeDirection.OUT, PROCESS_INPUTS_EDGE);
  
                  traverseEdgesOnDemand(processEdges, true, depth, lineageConstraintsMap, ret);
              }
  
              if (direction == OUTPUT || direction == BOTH) {
                  Iterable<AtlasEdge> processEdges = processVertex.getEdges(AtlasEdgeDirection.OUT, PROCESS_OUTPUTS_EDGE);
  
                  traverseEdgesOnDemand(processEdges, false, depth, lineageConstraintsMap, ret);
              }
          }
  
          return ret;
      }
  
      private void traverseEdges(AtlasVertex datasetVertex, boolean isInput, int depth, AtlasLineageInfo ret) throws AtlasBaseException {
          traverseEdges(datasetVertex, isInput, depth, new HashSet<>(), ret);
      }
  
      private void traverseEdges(AtlasVertex datasetVertex, boolean isInput, int depth, Set<String> visitedVertices, AtlasLineageInfo ret) throws AtlasBaseException {
          if (depth != 0) {
              // keep track of visited vertices to avoid circular loop
              visitedVertices.add(getId(datasetVertex));
  
              Iterable<AtlasEdge> incomingEdges = datasetVertex.getEdges(IN, isInput ? PROCESS_OUTPUTS_EDGE : PROCESS_INPUTS_EDGE);
  
              for (AtlasEdge incomingEdge : incomingEdges) {
                  AtlasVertex         processVertex = incomingEdge.getOutVertex();
                  Iterable<AtlasEdge> outgoingEdges = processVertex.getEdges(OUT, isInput ? PROCESS_INPUTS_EDGE : PROCESS_OUTPUTS_EDGE);
  
                  for (AtlasEdge outgoingEdge : outgoingEdges) {
                      AtlasVertex entityVertex = outgoingEdge.getInVertex();
  
                      if (entityVertex != null) {
                          addEdgeToResult(incomingEdge, ret);
                          addEdgeToResult(outgoingEdge, ret);
  
                          if (!visitedVertices.contains(getId(entityVertex))) {
                              traverseEdges(entityVertex, isInput, depth - 1, visitedVertices, ret);
                          }
                      }
                  }
              }
          }
      }
  
      private void traverseEdgesOnDemand(Iterable<AtlasEdge> processEdges, boolean isInput, int depth, Map<String, LineageOnDemandConstraints> lineageConstraintsMap, AtlasLineageInfo ret) throws AtlasBaseException {
          for (AtlasEdge processEdge : processEdges) {
              boolean isInputEdge = processEdge.getLabel().equalsIgnoreCase(PROCESS_INPUTS_EDGE);
  
              if (incrementAndCheckIfRelationsLimitReached(processEdge, isInputEdge, lineageConstraintsMap, ret)) {
                  break;
              } else {
                  addEdgeToResult(processEdge, ret);
              }
  
              AtlasVertex                datasetVertex           = processEdge.getInVertex();
              String                     inGuid                  = AtlasGraphUtilsV2.getIdFromVertex(datasetVertex);
              LineageOnDemandConstraints inGuidLineageConstrains = getAndValidateLineageConstraintsByGuid(inGuid, lineageConstraintsMap);
  
              if (!ret.getRelationsOnDemand().containsKey(inGuid)) {
                  ret.getRelationsOnDemand().put(inGuid, new LineageInfoOnDemand(inGuidLineageConstrains));
              }
  
              traverseEdgesOnDemand(datasetVertex, isInput, depth - 1, new HashSet<>(), lineageConstraintsMap, ret);
          }
      }
  
      private void traverseEdgesOnDemand(AtlasVertex datasetVertex, boolean isInput, int depth, Set<String> visitedVertices, Map<String, LineageOnDemandConstraints> lineageConstraintsMap, AtlasLineageInfo ret) throws AtlasBaseException {
          if (depth != 0) {
              // keep track of visited vertices to avoid circular loop
              visitedVertices.add(getId(datasetVertex));
  
              Iterable<AtlasEdge> incomingEdges = datasetVertex.getEdges(IN, isInput ? PROCESS_OUTPUTS_EDGE : PROCESS_INPUTS_EDGE);
  
              for (AtlasEdge incomingEdge : incomingEdges) {
                  if (incrementAndCheckIfRelationsLimitReached(incomingEdge, !isInput, lineageConstraintsMap, ret)) {
                      break;
                  } else {
                      addEdgeToResult(incomingEdge, ret);
                  }
  
                  AtlasVertex         processVertex = incomingEdge.getOutVertex();
                  Iterable<AtlasEdge> outgoingEdges = processVertex.getEdges(OUT, isInput ? PROCESS_INPUTS_EDGE : PROCESS_OUTPUTS_EDGE);
  
                  for (AtlasEdge outgoingEdge : outgoingEdges) {
                      if (incrementAndCheckIfRelationsLimitReached(outgoingEdge, isInput, lineageConstraintsMap, ret)) {
                          break;
                      } else {
                          addEdgeToResult(outgoingEdge, ret);
                      }
  
                      AtlasVertex entityVertex = outgoingEdge.getInVertex();
  
                      if (entityVertex != null && !visitedVertices.contains(getId(entityVertex))) {
                          traverseEdgesOnDemand(entityVertex, isInput, depth - 1, visitedVertices, lineageConstraintsMap, ret);
                      }
                  }
              }
          }
      }
  
      private boolean incrementAndCheckIfRelationsLimitReached(AtlasEdge atlasEdge, boolean isInput, Map<String, LineageOnDemandConstraints> lineageConstraintsMap, AtlasLineageInfo ret) {
          if (lineageContainsEdge(ret, atlasEdge)) {
              return false;
          }
  
          boolean hasRelationsLimitReached = false;
  
          AtlasVertex                inVertex                 = isInput ? atlasEdge.getOutVertex() : atlasEdge.getInVertex();
          String                     inGuid                   = AtlasGraphUtilsV2.getIdFromVertex(inVertex);
          LineageOnDemandConstraints inGuidLineageConstraints = getAndValidateLineageConstraintsByGuid(inGuid, lineageConstraintsMap);
  
          AtlasVertex                outVertex                 = isInput ? atlasEdge.getInVertex() : atlasEdge.getOutVertex();
          String                     outGuid                   = AtlasGraphUtilsV2.getIdFromVertex(outVertex);
          LineageOnDemandConstraints outGuidLineageConstraints = getAndValidateLineageConstraintsByGuid(outGuid, lineageConstraintsMap);
  
          LineageInfoOnDemand inLineageInfo  = ret.getRelationsOnDemand().containsKey(inGuid) ? ret.getRelationsOnDemand().get(inGuid) : new LineageInfoOnDemand(inGuidLineageConstraints);
          LineageInfoOnDemand outLineageInfo = ret.getRelationsOnDemand().containsKey(outGuid) ? ret.getRelationsOnDemand().get(outGuid) : new LineageInfoOnDemand(outGuidLineageConstraints);
  
          if (inLineageInfo.isInputRelationsReachedLimit()) {
              inLineageInfo.setHasMoreInputs(true);
  
              hasRelationsLimitReached = true;
          } else {
              inLineageInfo.incrementInputRelationsCount();
          }
  
          if (outLineageInfo.isOutputRelationsReachedLimit()) {
              outLineageInfo.setHasMoreOutputs(true);
  
              hasRelationsLimitReached = true;
          } else {
              outLineageInfo.incrementOutputRelationsCount();
          }
  
          if (!hasRelationsLimitReached) {
              ret.getRelationsOnDemand().put(inGuid, inLineageInfo);
              ret.getRelationsOnDemand().put(outGuid, outLineageInfo);
          }
  
          return hasRelationsLimitReached;
      }
  
      private void addEdgeToResult(AtlasEdge edge, AtlasLineageInfo lineageInfo) throws AtlasBaseException {
          if (!lineageContainsEdge(lineageInfo, edge) && !lineageMaxNodeCountReached(lineageInfo.getRelations())) {
              processEdge(edge, lineageInfo);
          }
      }
  
      private int getLineageMaxNodeAllowedCount() {
          return Math.min(DEFAULT_LINEAGE_MAX_NODE_COUNT, AtlasConfiguration.LINEAGE_MAX_NODE_COUNT.getInt());
      }
  
      private boolean lineageMaxNodeCountReached(Set<LineageRelation> relations) {
          return CollectionUtils.isNotEmpty(relations) && relations.size() > getLineageMaxNodeAllowedCount();
      }
  
      private String getVisitedEdgeLabel(String inGuid, String outGuid, String relationGuid) {
          if (isLineageOnDemandEnabled()) {
              return inGuid + SEPARATOR + outGuid;
          }
  
          return relationGuid;
      }
  
      private boolean lineageContainsEdge(AtlasLineageInfo lineageInfo, AtlasEdge edge) {
          boolean ret = false;
  
          if (edge != null && lineageInfo != null && CollectionUtils.isNotEmpty(lineageInfo.getVisitedEdges())) {
              String  inGuid           = AtlasGraphUtilsV2.getIdFromVertex(edge.getInVertex());
              String  outGuid          = AtlasGraphUtilsV2.getIdFromVertex(edge.getOutVertex());
              String  relationGuid     = AtlasGraphUtilsV2.getEncodedProperty(edge, RELATIONSHIP_GUID_PROPERTY_KEY, String.class);
              boolean isInputEdge      = edge.getLabel().equalsIgnoreCase(PROCESS_INPUTS_EDGE);
              String  visitedEdgeLabel = isInputEdge ? getVisitedEdgeLabel(inGuid, outGuid, relationGuid) : getVisitedEdgeLabel(outGuid, inGuid, relationGuid);
  
              if (lineageInfo.getVisitedEdges().contains(visitedEdgeLabel)) {
                  ret = true;
              }
          }
  
          return ret;
      }
  
      private void processEdge(final AtlasEdge edge, final AtlasLineageInfo lineageInfo) throws AtlasBaseException {
          processEdge(edge, lineageInfo.getGuidEntityMap(), lineageInfo.getRelations(), lineageInfo.getVisitedEdges());
      }
  
      private AtlasLineageInfo initializeLineageInfo(String guid, LineageDirection direction, int depth) {
          return new AtlasLineageInfo(guid, new HashMap<>(), new HashSet<>(), new HashSet<>(), new HashMap<>(), direction, depth);
      }
  
      private static String getId(AtlasVertex vertex) {
          return vertex.getIdForDisplay();
      }
  
      private List<?> executeGremlinScript(Map<String, Object> bindings, String lineageQuery) throws AtlasBaseException {
          List<?>      ret;
          ScriptEngine engine = graph.getGremlinScriptEngine();
  
          try {
              ret = (List<?>) graph.executeGremlinScript(engine, bindings, lineageQuery, false);
          } catch (ScriptException e) {
              throw new AtlasBaseException(INSTANCE_LINEAGE_QUERY_FAILED, lineageQuery);
          } finally {
              graph.releaseGremlinScriptEngine(engine);
          }
  
          return ret;
      }
  
      private void processEdge(final AtlasEdge edge, final Map<String, AtlasEntityHeader> entities, final Set<LineageRelation> relations) throws AtlasBaseException {
          processEdge(edge, entities, relations, null);
      }
  
      private void processEdge(final AtlasEdge edge, final Map<String, AtlasEntityHeader> entities, final Set<LineageRelation> relations, final Set<String> visitedEdges) throws AtlasBaseException {
          AtlasVertex inVertex     = edge.getInVertex();
          AtlasVertex outVertex    = edge.getOutVertex();
          String      inGuid       = AtlasGraphUtilsV2.getIdFromVertex(inVertex);
          String      outGuid      = AtlasGraphUtilsV2.getIdFromVertex(outVertex);
          String      relationGuid = AtlasGraphUtilsV2.getEncodedProperty(edge, RELATIONSHIP_GUID_PROPERTY_KEY, String.class);
          boolean     isInputEdge  = edge.getLabel().equalsIgnoreCase(PROCESS_INPUTS_EDGE);
  
          if (!entities.containsKey(inGuid)) {
              AtlasEntityHeader entityHeader = entityRetriever.toAtlasEntityHeader(inVertex);
  
              entities.put(inGuid, entityHeader);
          }
  
          if (!entities.containsKey(outGuid)) {
              AtlasEntityHeader entityHeader = entityRetriever.toAtlasEntityHeader(outVertex);
  
              entities.put(outGuid, entityHeader);
          }
  
          if (isInputEdge) {
              relations.add(new LineageRelation(inGuid, outGuid, relationGuid));
          } else {
              relations.add(new LineageRelation(outGuid, inGuid, relationGuid));
          }
  
          if (visitedEdges != null) {
              String visitedEdgeLabel = isInputEdge ? getVisitedEdgeLabel(inGuid, outGuid, relationGuid) : getVisitedEdgeLabel(outGuid, inGuid, relationGuid);
  
              visitedEdges.add(visitedEdgeLabel);
          }
      }
  
      private AtlasLineageInfo getBothLineageInfoV1(String guid, int depth, boolean isDataSet) throws AtlasBaseException {
          AtlasLineageInfo inputLineage  = getLineageInfo(guid, INPUT, depth, isDataSet);
          AtlasLineageInfo outputLineage = getLineageInfo(guid, OUTPUT, depth, isDataSet);
          AtlasLineageInfo ret           = inputLineage;
  
          ret.getRelations().addAll(outputLineage.getRelations());
          ret.getGuidEntityMap().putAll(outputLineage.getGuidEntityMap());
          ret.setLineageDirection(BOTH);
  
          return ret;
      }
  
      private String getLineageQuery(String entityGuid, LineageDirection direction, int depth, boolean isDataSet, Map<String, Object> bindings) {
          String incomingFrom = null;
          String outgoingTo   = null;
          String ret;
  
          if (direction.equals(INPUT)) {
              incomingFrom = PROCESS_OUTPUTS_EDGE;
              outgoingTo   = PROCESS_INPUTS_EDGE;
          } else if (direction.equals(OUTPUT)) {
              incomingFrom = PROCESS_INPUTS_EDGE;
              outgoingTo   = PROCESS_OUTPUTS_EDGE;
          }
  
          bindings.put("guid", entityGuid);
          bindings.put("incomingEdgeLabel", incomingFrom);
          bindings.put("outgoingEdgeLabel", outgoingTo);
          bindings.put("dataSetDepth", depth);
          bindings.put("processDepth", depth - 1);
  
          if (depth < 1) {
              ret = isDataSet ? gremlinQueryProvider.getQuery(FULL_LINEAGE_DATASET) : gremlinQueryProvider.getQuery(FULL_LINEAGE_PROCESS);
          } else {
              ret = isDataSet ? gremlinQueryProvider.getQuery(PARTIAL_LINEAGE_DATASET) : gremlinQueryProvider.getQuery(PARTIAL_LINEAGE_PROCESS);
          }
  
          return ret;
      }
  }
  