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 com.google.common.annotations.VisibleForTesting;
  import org.apache.atlas.AtlasConfiguration;
  import org.apache.atlas.AtlasErrorCode;
  import org.apache.atlas.RequestContext;
  import org.apache.atlas.entitytransform.BaseEntityHandler;
  import org.apache.atlas.entitytransform.TransformerContext;
  import org.apache.atlas.exception.AtlasBaseException;
  import org.apache.atlas.model.PList;
  import org.apache.atlas.model.audit.AtlasAuditEntry;
  import org.apache.atlas.model.impexp.AsyncImportStatus;
  import org.apache.atlas.model.impexp.AtlasAsyncImportRequest;
  import org.apache.atlas.model.impexp.AtlasExportRequest;
  import org.apache.atlas.model.impexp.AtlasImportRequest;
  import org.apache.atlas.model.impexp.AtlasImportResult;
  import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
  import org.apache.atlas.model.instance.AtlasObjectId;
  import org.apache.atlas.model.instance.EntityMutationResponse;
  import org.apache.atlas.model.typedef.AtlasTypesDef;
  import org.apache.atlas.repository.audit.AtlasAuditService;
  import org.apache.atlas.repository.store.graph.BulkImporter;
  import org.apache.atlas.repository.store.graph.v2.AsyncImportTaskExecutor;
  import org.apache.atlas.repository.store.graph.v2.EntityImportStream;
  import org.apache.atlas.store.AtlasTypeDefStore;
  import org.apache.atlas.type.AtlasType;
  import org.apache.atlas.type.AtlasTypeRegistry;
  import org.apache.atlas.utils.AtlasAsyncImportTestUtil;
  import org.apache.atlas.utils.AtlasJson;
  import org.apache.atlas.utils.AtlasStringUtil;
  import org.apache.atlas.v1.typesystem.types.utils.TypesUtil;
  import org.apache.commons.collections.CollectionUtils;
  import org.apache.commons.collections.MapUtils;
  import org.apache.commons.lang3.StringUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import org.springframework.stereotype.Component;
  
  import javax.inject.Inject;
  
  import java.io.File;
  import java.io.FileInputStream;
  import java.io.FileNotFoundException;
  import java.io.IOException;
  import java.io.InputStream;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.stream.Collectors;
  
  import static org.apache.atlas.model.impexp.AtlasAsyncImportRequest.ImportStatus.FAILED;
  import static org.apache.atlas.model.impexp.AtlasImportRequest.TRANSFORMERS_KEY;
  import static org.apache.atlas.model.impexp.AtlasImportRequest.TRANSFORMS_KEY;
  import static org.apache.atlas.model.impexp.AtlasImportResult.OperationStatus.SUCCESS;
  
  @Component
  public class ImportService implements AsyncImporter {
      private static final Logger LOG = LoggerFactory.getLogger(ImportService.class);
  
      private static final String ATLAS_TYPE_HIVE_TABLE = "hive_table";
      private static final String OPERATION_STATUS      = "operationStatus";
  
      private final AtlasTypeDefStore       typeDefStore;
      private final AtlasTypeRegistry       typeRegistry;
      private final BulkImporter            bulkImporter;
      private final AuditsWriter            auditsWriter;
      private final ImportTransformsShaper  importTransformsShaper;
      private final AsyncImportTaskExecutor asyncImportTaskExecutor;
      private final AsyncImportService      asyncImportService;
      private final AtlasAuditService       auditService;
  
      private final TableReplicationRequestProcessor tableReplicationRequestProcessor;
  
      @Inject
      public ImportService(AtlasTypeDefStore typeDefStore, AtlasTypeRegistry typeRegistry, BulkImporter bulkImporter,
              AuditsWriter auditsWriter, ImportTransformsShaper importTransformsShaper,
              TableReplicationRequestProcessor tableReplicationRequestProcessor, AsyncImportTaskExecutor asyncImportTaskExecutor,
              AsyncImportService asyncImportService, AtlasAuditService auditService) {
          this.typeDefStore                     = typeDefStore;
          this.typeRegistry                     = typeRegistry;
          this.bulkImporter                     = bulkImporter;
          this.auditsWriter                     = auditsWriter;
          this.importTransformsShaper           = importTransformsShaper;
          this.tableReplicationRequestProcessor = tableReplicationRequestProcessor;
          this.asyncImportTaskExecutor = asyncImportTaskExecutor;
          this.asyncImportService = asyncImportService;
          this.auditService  = auditService;
      }
  
      public AtlasImportResult run(InputStream inputStream, String userName, String hostName, String requestingIP) throws AtlasBaseException {
          return run(inputStream, null, userName, hostName, requestingIP);
      }
  
      public AtlasImportResult run(InputStream inputStream, AtlasImportRequest request, String userName, String hostName, String requestingIP) throws AtlasBaseException {
          if (request == null) {
              request = new AtlasImportRequest();
          }
  
          EntityImportStream source = createZipSource(request, inputStream, AtlasConfiguration.IMPORT_TEMP_DIRECTORY.getString());
  
          return run(source, request, userName, hostName, requestingIP);
      }
  
      public AtlasImportResult run(AtlasImportRequest request, String userName, String hostName, String requestingIP) throws AtlasBaseException {
          String fileName = request.getFileName();
  
          if (StringUtils.isBlank(fileName)) {
              throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "FILENAME parameter not found");
          }
  
          AtlasImportResult result = null;
  
          try {
              LOG.info("==> import(user={}, from={}, fileName={})", userName, requestingIP, fileName);
  
              File file = new File(fileName);
  
              result = run(new FileInputStream(file), request, userName, hostName, requestingIP);
          } catch (AtlasBaseException excp) {
              LOG.error("import(user={}, from={}, fileName={}): failed", userName, requestingIP, excp);
  
              throw excp;
          } catch (FileNotFoundException excp) {
              LOG.error("import(user={}, from={}, fileName={}): file not found", userName, requestingIP, excp);
  
              throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, fileName + ": file not found");
          } catch (Exception excp) {
              LOG.error("import(user={}, from={}, fileName={}): failed", userName, requestingIP, excp);
  
              throw new AtlasBaseException(excp);
          } finally {
              LOG.info("<== import(user={}, from={}, fileName={}): status={}", userName, requestingIP, fileName,
                      (result == null ? AtlasImportResult.OperationStatus.FAIL : result.getOperationStatus()));
          }
  
          return result;
      }
  
      public AtlasAsyncImportRequest run(AtlasImportRequest request, InputStream inputStream, String userName, String hostName, String requestingIP) throws AtlasBaseException {
          if (request == null) {
              request = new AtlasImportRequest();
          }
  
          try {
              LOG.info("==> asyncImport(user={}, from={}, request={})", userName, requestingIP, request);
  
              EntityImportStream source     = createZipSource(inputStream, AtlasConfiguration.IMPORT_TEMP_DIRECTORY.getString());
              String             transforms = MapUtils.isNotEmpty(request.getOptions()) ? request.getOptions().get(TRANSFORMS_KEY) : null;
  
              setImportTransform(source, transforms);
  
              String transformers = MapUtils.isNotEmpty(request.getOptions()) ? request.getOptions().get(TRANSFORMERS_KEY) : null;
  
              setEntityTransformerHandlers(source, transformers);
  
              AtlasImportResult result = new AtlasImportResult(request, userName, requestingIP, hostName, System.currentTimeMillis());
  
              result.setExportResult(source.getExportResult());
  
              return asyncImportTaskExecutor.run(result, source);
          } finally {
              LOG.info("<== asyncImport(user={}, from={}, request={})", userName, requestingIP, request);
          }
      }
  
      public PList<AsyncImportStatus> getAsyncImportsStatus(int offset, int limit) throws AtlasBaseException {
          return asyncImportService.getAsyncImportsStatus(offset, limit);
      }
  
      public AtlasAsyncImportRequest getAsyncImportRequest(String importId) throws AtlasBaseException {
          return asyncImportService.getAsyncImportRequest(importId);
      }
  
      public void abortAsyncImport(String importId) throws AtlasBaseException {
          asyncImportTaskExecutor.abortAsyncImportRequest(importId);
      }
  
      @VisibleForTesting
      AtlasImportResult run(EntityImportStream source, AtlasImportRequest request, String userName, String hostName, String requestingIP) throws AtlasBaseException {
          AtlasImportResult result = new AtlasImportResult(request, userName, requestingIP, hostName, System.currentTimeMillis());
  
          long startTimestamp = System.currentTimeMillis();
  
          try {
              LOG.info("==> import(user={}, from={}, request={})", userName, requestingIP, request);
  
              RequestContext.get().setImportInProgress(true);
  
              String transforms = AtlasStringUtil.getOption(request.getOptions(), TRANSFORMS_KEY);
  
              setImportTransform(source, transforms);
  
              String transformers = AtlasStringUtil.getOption(request.getOptions(), TRANSFORMERS_KEY);
  
              setEntityTransformerHandlers(source, transformers);
  
              processTypes(source.getTypesDef(), result);
              setStartPosition(request, source);
  
              processEntities(userName, source, result);
  
              processReplicationDeletion(source.getExportResult().getRequest(), request);
          } catch (AtlasBaseException excp) {
              LOG.error("import(user={}, from={}): failed", userName, requestingIP, excp);
  
              throw excp;
          } catch (Exception excp) {
              LOG.error("import(user={}, from={}): failed", userName, requestingIP, excp);
  
              throw new AtlasBaseException(excp);
          } finally {
              RequestContext.get().setImportInProgress(false);
  
              long endTimestamp = System.currentTimeMillis();
  
              result.incrementMeticsCounter("duration", getDuration(endTimestamp, startTimestamp));
  
              if (source != null) {
                  source.close();
              }
  
              LOG.info("<== import(user={}, from={}): status={}", userName, requestingIP, result.getOperationStatus());
          }
  
          return result;
      }
  
      @Override
      public void onImportTypeDef(AtlasTypesDef typesDef, String importId) throws AtlasBaseException {
          LOG.info("==> onImportTypeDef(importId={})", importId);
  
          AtlasAsyncImportRequest importRequest = asyncImportService.fetchImportRequestByImportId(importId);
  
          if (importRequest == null) {
              throw new AtlasBaseException(AtlasErrorCode.IMPORT_NOT_FOUND, importId);
          }
  
          AtlasImportResult result = importRequest.getImportResult();
  
          try {
              RequestContext.get().setImportInProgress(true);
  
              processTypes(typesDef, result);
          } catch (AtlasBaseException abe) {
              importRequest.setStatus(FAILED);
  
              throw abe;
          } finally {
              RequestContext.get().setImportInProgress(false);
  
              importRequest.setImportResult(result);
  
              asyncImportService.updateImportRequest(importRequest);
  
              LOG.info("<== onImportTypeDef()");
          }
      }
  
      @Override
      public Boolean onImportEntity(AtlasEntityWithExtInfo entityWithExtInfo, String importId, int position) throws AtlasBaseException {
          LOG.info("==> onImportEntity(importId={}, position={})", importId, position);
  
          AtlasAsyncImportRequest importRequest = asyncImportService.fetchImportRequestByImportId(importId);
  
          if (importRequest == null) {
              throw new AtlasBaseException(AtlasErrorCode.IMPORT_NOT_FOUND, importId);
          }
  
          AtlasImportResult      result                  = importRequest.getImportResult();
          float                  importProgress          = importRequest.getImportDetails().getImportProgress();
          int                    importedEntitiesCounter = importRequest.getImportDetails().getImportedEntitiesCount();
          int                    failedEntitiesCounter   = importRequest.getImportDetails().getFailedEntitiesCount();
          Set<String>            processedEntities       = new HashSet<>(result.getProcessedEntities());
          List<String>           failedEntities          = importRequest.getImportDetails().getFailedEntities();
          EntityMutationResponse entityMutationResponse  = null;
          long                   startTimestamp          = System.currentTimeMillis();
  
          try {
              RequestContext.get().setImportInProgress(true);
  
              TypesUtil.Pair<EntityMutationResponse, Float> resp = this.bulkImporter.asyncImport(entityWithExtInfo, entityMutationResponse,
                      result, processedEntities, failedEntities, position, importRequest.getImportDetails().getTotalEntitiesCount(), importProgress);
  
              importedEntitiesCounter += 1;
  
              importRequest.getImportDetails().setImportedEntitiesCount(importedEntitiesCounter);
  
              result.setProcessedEntities(new ArrayList<>(processedEntities));
  
              importRequest.getImportDetails().setImportProgress(resp.right);
          } catch (AtlasBaseException abe) {
              LOG.warn("Failed to import entity: {} at position: {} for import: {}", entityWithExtInfo.getEntity().getGuid(), position, importId, abe);
              failedEntitiesCounter += 1;
  
              importRequest.getImportDetails().setFailedEntitiesCount(failedEntitiesCounter);
              failedEntities.add(entityWithExtInfo.getEntity().getGuid());
              importRequest.getImportDetails().setFailedEntities(failedEntities);
              importRequest.getImportDetails().addFailure(entityWithExtInfo.getEntity().getGuid(), abe.getMessage());
          } finally {
              RequestContext.get().setImportInProgress(false);
  
              result.incrementMeticsCounter("duration", getDuration(System.currentTimeMillis(), startTimestamp));
              importRequest.setImportResult(result);
              importRequest.setCompletedTime(System.currentTimeMillis());
  
              asyncImportService.updateImportRequest(importRequest);
  
              LOG.info("<== onImportEntity(importId={}, position={})", importId, position);
          }
  
          if (importRequest.getImportDetails().getPublishedEntityCount() <=
                  importRequest.getImportDetails().getImportedEntitiesCount() + importRequest.getImportDetails().getFailedEntitiesCount()) {
              onImportComplete(importId);
              return true;
          }
          return false;
      }
  
      @Override
      public void onImportComplete(String importId) throws AtlasBaseException {
          LOG.info("==> onImportComplete(importId={})", importId);
  
          try {
              AtlasAsyncImportRequest importRequest = asyncImportService.fetchImportRequestByImportId(importId);
  
              if (importRequest == null) {
                  throw new AtlasBaseException(AtlasErrorCode.IMPORT_NOT_FOUND, importId);
              }
  
              AtlasAsyncImportTestUtil.intercept(importRequest);
  
              if (importRequest.getImportDetails().getTotalEntitiesCount() == importRequest.getImportDetails().getImportedEntitiesCount()) {
                  importRequest.setStatus(AtlasAsyncImportRequest.ImportStatus.SUCCESSFUL);
                  importRequest.getImportResult().setOperationStatus(SUCCESS);
              } else if (importRequest.getImportDetails().getImportedEntitiesCount() > 0) {
                  importRequest.setStatus(AtlasAsyncImportRequest.ImportStatus.PARTIAL_SUCCESS);
                  importRequest.getImportResult().setOperationStatus(AtlasImportResult.OperationStatus.PARTIAL_SUCCESS);
              } else {
                  importRequest.setStatus(FAILED);
              }
  
              asyncImportService.updateImportRequest(importRequest);
  
              AtlasImportResult result = importRequest.getImportResult();
  
              processReplicationDeletion(result.getExportResult().getRequest(), result.getRequest());
  
              auditsWriter.write(result.getUserName(), result, result.getTimeStamp(), System.currentTimeMillis(), importRequest.getImportDetails().getCreationOrder());
  
              addToImportOperationAudits(result);
          }
          finally {
              LOG.info("<== onImportComplete(importId={})", importId);
          }
      }
  
      @Override
      public void onCompleteImportRequest(String importId) {
          LOG.info("==> onCompleteImportRequest(importId={})", importId);
  
          asyncImportTaskExecutor.onCompleteImportRequest(importId);
  
          LOG.info("<== onCompleteImportRequest(importId={})", importId);
      }
  
      @VisibleForTesting
      void setImportTransform(EntityImportStream source, String transforms) throws AtlasBaseException {
          ImportTransforms importTransform = ImportTransforms.fromJson(transforms);
  
          if (importTransform == null) {
              return;
          }
  
          importTransformsShaper.shape(importTransform, source.getExportResult().getRequest());
  
          source.setImportTransform(importTransform);
  
          if (LOG.isDebugEnabled()) {
              debugLog("   => transforms: {}", AtlasType.toJson(importTransform));
          }
      }
  
      @VisibleForTesting
      void setEntityTransformerHandlers(EntityImportStream source, String transformersJson) throws AtlasBaseException {
          if (StringUtils.isEmpty(transformersJson)) {
              return;
          }
  
          TransformerContext      context        = new TransformerContext(typeRegistry, typeDefStore, source.getExportResult().getRequest());
          List<BaseEntityHandler> entityHandlers = BaseEntityHandler.fromJson(transformersJson, context);
  
          if (CollectionUtils.isEmpty(entityHandlers)) {
              return;
          }
  
          source.setEntityHandlers(entityHandlers);
      }
  
      @VisibleForTesting
      boolean checkHiveTableIncrementalSkipLineage(AtlasImportRequest importRequest, AtlasExportRequest exportRequest) {
          if (exportRequest == null || CollectionUtils.isEmpty(exportRequest.getItemsToExport())) {
              return false;
          }
  
          for (AtlasObjectId itemToExport : exportRequest.getItemsToExport()) {
              if (!itemToExport.getTypeName().equalsIgnoreCase(ATLAS_TYPE_HIVE_TABLE)) {
                  return false;
              }
          }
  
          return importRequest.isReplicationOptionSet() && exportRequest.isReplicationOptionSet() &&
                  exportRequest.getFetchTypeOptionValue().equalsIgnoreCase(AtlasExportRequest.FETCH_TYPE_INCREMENTAL) &&
                  exportRequest.getSkipLineageOptionValue();
      }
  
      private void debugLog(String s, Object... params) {
          if (LOG.isDebugEnabled()) {
              LOG.debug(s, params);
          }
      }
  
      private void setStartPosition(AtlasImportRequest request, EntityImportStream source) throws AtlasBaseException {
          if (request.getStartGuid() != null) {
              source.setPositionUsingEntityGuid(request.getStartGuid());
          } else if (request.getStartPosition() != null) {
              source.setPosition(Integer.parseInt(request.getStartPosition()));
          }
      }
  
      @VisibleForTesting
      void processTypes(AtlasTypesDef typeDefinitionMap, AtlasImportResult result) throws AtlasBaseException {
          if (result.getRequest().getUpdateTypeDefs() != null && !result.getRequest().getUpdateTypeDefs().equals("true")) {
              return;
          }
  
          ImportTypeDefProcessor importTypeDefProcessor = new ImportTypeDefProcessor(this.typeDefStore, this.typeRegistry);
  
          importTypeDefProcessor.processTypes(typeDefinitionMap, result);
      }
  
      private void processEntities(String userName, EntityImportStream importSource, AtlasImportResult result) throws AtlasBaseException {
          result.setExportResult(importSource.getExportResult());
  
          long startTimestamp = System.currentTimeMillis();
  
          this.bulkImporter.bulkImport(importSource, result);
  
          long endTimestamp = System.currentTimeMillis();
  
          result.incrementMeticsCounter("duration", getDuration(endTimestamp, startTimestamp));
          result.setOperationStatus(AtlasImportResult.OperationStatus.SUCCESS);
  
          if (isMigrationMode(result.getRequest())) {
              return;
          }
  
          auditsWriter.write(userName, result, startTimestamp, endTimestamp, importSource.getCreationOrder());
      }
  
      private EntityImportStream createZipSource(InputStream inputStream, String configuredTemporaryDirectory) throws AtlasBaseException {
          try {
              if (StringUtils.isEmpty(configuredTemporaryDirectory)) {
                  return new ZipSource(inputStream);
              } else {
                  return new ZipSourceWithBackingDirectory(inputStream, configuredTemporaryDirectory);
              }
          } catch (IOException ex) {
              throw new AtlasBaseException(ex);
          }
      }
  
      @VisibleForTesting
      void processReplicationDeletion(AtlasExportRequest exportRequest, AtlasImportRequest importRequest) throws AtlasBaseException {
          if (checkHiveTableIncrementalSkipLineage(importRequest, exportRequest)) {
              tableReplicationRequestProcessor.process(exportRequest, importRequest);
          }
      }
  
      private int getDuration(long endTime, long startTime) {
          return (int) (endTime - startTime);
      }
  
      private EntityImportStream createZipSource(AtlasImportRequest request, InputStream inputStream, String configuredTemporaryDirectory) throws AtlasBaseException {
          try {
              if (isMigrationMode(request) || AtlasStringUtil.optionEquals(request.getOptions(), AtlasImportRequest.OPTION_KEY_FORMAT, AtlasImportRequest.OPTION_KEY_FORMAT_ZIP_DIRECT)) {
                  LOG.info("ZipSource Format: ZipDirect: Size: {}", AtlasStringUtil.getOption(request.getOptions(), "size"));
  
                  return getZipDirectEntityImportStream(request, inputStream);
              } else if (StringUtils.isEmpty(configuredTemporaryDirectory)) {
                  return new ZipSource(inputStream);
              } else {
                  return new ZipSourceWithBackingDirectory(inputStream, configuredTemporaryDirectory);
              }
          } catch (IOException ex) {
              throw new AtlasBaseException(ex);
          }
      }
  
      private EntityImportStream getZipDirectEntityImportStream(AtlasImportRequest request, InputStream inputStream) throws IOException, AtlasBaseException {
          ZipSourceDirect zipSourceDirect = new ZipSourceDirect(inputStream, request.getSizeOption());
  
          LOG.info("Using ZipSourceDirect: Size: {} entities", zipSourceDirect.size());
  
          return zipSourceDirect;
      }
  
      private boolean isMigrationMode(AtlasImportRequest request) {
          return AtlasStringUtil.hasOption(request.getOptions(), AtlasImportRequest.OPTION_KEY_MIGRATION);
      }
  
      @VisibleForTesting
      void addToImportOperationAudits(AtlasImportResult result) throws AtlasBaseException {
          String params = AtlasJson.toJson(Collections.singletonMap(OPERATION_STATUS, result.getOperationStatus().name()));
  
          if (result.getExportResult().getRequest() == null) {
              int resultCount = result.getProcessedEntities().size();
  
              auditService.add(AtlasAuditEntry.AuditOperation.IMPORT, params, AtlasJson.toJson(result.getMetrics()), resultCount);
          } else {
              List<AtlasObjectId> objectIds         = result.getExportResult().getRequest().getItemsToExport();
              Map<String, Long>   entityCountByType = objectIds.stream().collect(Collectors.groupingBy(AtlasObjectId::getTypeName, Collectors.counting()));
              int                 resultCount       = objectIds.size();
  
              auditService.add(AtlasAuditEntry.AuditOperation.IMPORT, params, AtlasJson.toJson(entityCountByType), resultCount);
          }
      }
  }
  