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.impexp;
  
  import org.apache.atlas.exception.AtlasBaseException;
  import org.apache.atlas.model.instance.AtlasEntity;
  import org.apache.atlas.repository.graphdb.AtlasGraph;
  import org.apache.atlas.repository.util.UniqueList;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import javax.inject.Inject;
  import javax.script.ScriptEngine;
  import javax.script.ScriptException;
  
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  
  public class IncrementalExportEntityProvider implements ExtractStrategy {
      private static final Logger LOG = LoggerFactory.getLogger(IncrementalExportEntityProvider.class);
  
      private static final String QUERY_PARAMETER_START_GUID             = "startGuid";
      private static final String QUERY_PARAMETER_MODIFICATION_TIMESTAMP = "modificationTimestamp";
      private static final String QUERY_DB                               = "g.V().has('__guid', startGuid)";
      private static final String QUERY_TABLE                            = QUERY_DB + ".in('__hive_table.db')";
      private static final String QUERY_SD                               = QUERY_TABLE + ".out('__hive_table.sd')";
      private static final String QUERY_COLUMN                           = QUERY_TABLE + ".out('__hive_table.columns')";
      private static final String TRANSFORM_CLAUSE                       = ".project('__guid').by('__guid').dedup().toList()";
      private static final String TIMESTAMP_CLAUSE                       = ".has('__modificationTimestamp', gt(modificationTimestamp))";
      private static final String QUERY_TABLE_DB                         = QUERY_DB + ".out('__hive_table.db')";
      private static final String QUERY_TABLE_SD                         = QUERY_DB + ".out('__hive_table.sd')";
      private static final String QUERY_TABLE_COLUMNS                    = QUERY_DB + ".out('__hive_table.columns')";
  
      private final AtlasGraph   atlasGraph;
      private final ScriptEngine scriptEngine;
  
      @Inject
      public IncrementalExportEntityProvider(AtlasGraph atlasGraph) {
          this.atlasGraph = atlasGraph;
  
          ScriptEngine scriptEngine = null;
  
          try {
              scriptEngine = atlasGraph.getGremlinScriptEngine();
          } catch (AtlasBaseException e) {
              LOG.error("Error instantiating script engine.", e);
          }
  
          this.scriptEngine = scriptEngine;
      }
  
      @Override
      public void connectedFetch(AtlasEntity entity, ExportService.ExportContext context) {
          //starting entity is hive_table
          context.guidsToProcess.addAll(fetchGuids(entity.getGuid(), QUERY_TABLE_DB, context.changeMarker));
          context.guidsToProcess.addAll(fetchGuids(entity.getGuid(), QUERY_TABLE_SD, context.changeMarker));
          context.guidsToProcess.addAll(fetchGuids(entity.getGuid(), QUERY_TABLE_COLUMNS, context.changeMarker));
      }
  
      @Override
      public void fullFetch(AtlasEntity entity, ExportService.ExportContext context) {
          populate(entity.getGuid(), context.changeMarker, context.guidsToProcess);
      }
  
      @Override
      public void close() {
          if (scriptEngine != null) {
              atlasGraph.releaseGremlinScriptEngine(scriptEngine);
          }
      }
  
      public void populate(String dbEntityGuid, long timeStamp, UniqueList<String> guidsToProcess) {
          if (timeStamp == 0L) {
              full(dbEntityGuid, guidsToProcess);
          } else {
              partial(dbEntityGuid, timeStamp, guidsToProcess);
          }
      }
  
      private void partial(String dbEntityGuid, long timeStamp, UniqueList<String> guidsToProcess) {
          guidsToProcess.addAll(fetchGuids(dbEntityGuid, QUERY_TABLE, timeStamp));
          guidsToProcess.addAll(fetchGuids(dbEntityGuid, QUERY_SD, timeStamp));
          guidsToProcess.addAll(fetchGuids(dbEntityGuid, QUERY_COLUMN, timeStamp));
      }
  
      private void full(String dbEntityGuid, UniqueList<String> guidsToProcess) {
          guidsToProcess.addAll(fetchGuids(dbEntityGuid, QUERY_TABLE, 0L));
      }
  
      private List<String> fetchGuids(final String dbEntityGuid, String query, long timeStamp) {
          Map<String, Object> bindings        = new HashMap<>(Collections.singletonMap(QUERY_PARAMETER_START_GUID, dbEntityGuid));
          String              queryWithClause = query;
  
          if (timeStamp > 0L) {
              bindings.put(QUERY_PARAMETER_MODIFICATION_TIMESTAMP, timeStamp);
  
              queryWithClause = queryWithClause.concat(TIMESTAMP_CLAUSE);
          }
  
          return executeGremlinQuery(queryWithClause, bindings);
      }
  
      private List<String> executeGremlinQuery(String query, Map<String, Object> bindings) {
          try {
              List<String> guids              = new ArrayList<>();
              String       queryWithTransform = query + TRANSFORM_CLAUSE;
  
              List<Map<String, Object>> result = (List<Map<String, Object>>) atlasGraph.executeGremlinScript(scriptEngine, bindings, queryWithTransform, false);
  
              if (result == null) {
                  return guids;
              }
  
              for (Map<String, Object> item : result) {
                  guids.add((String) item.get(EntitiesExtractor.PROPERTY_GUID));
              }
  
              return guids;
          } catch (ScriptException e) {
              LOG.error("error executing query: {}: bindings: {}", query, bindings, e);
  
              return null;
          }
      }
  }
  