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.migration;
  
  import org.apache.atlas.ApplicationProperties;
  import org.apache.atlas.AtlasException;
  import org.apache.atlas.RequestContext;
  import org.apache.atlas.exception.AtlasBaseException;
  import org.apache.atlas.model.impexp.AtlasImportRequest;
  import org.apache.atlas.model.migration.MigrationImportStatus;
  import org.apache.atlas.repository.graph.AtlasGraphProvider;
  import org.apache.atlas.repository.impexp.ImportService;
  import org.apache.atlas.type.AtlasType;
  import org.apache.commons.codec.digest.DigestUtils;
  import org.apache.commons.collections.CollectionUtils;
  import org.apache.commons.io.filefilter.WildcardFileFilter;
  import org.apache.commons.lang.ArrayUtils;
  import org.apache.commons.lang.StringUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.File;
  import java.io.FileFilter;
  import java.io.FileInputStream;
  import java.io.IOException;
  import java.io.InputStream;
  import java.net.InetAddress;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.List;
  import java.util.Map;
  import java.util.zip.ZipFile;
  
  public class ZipFileMigrationImporter implements Runnable {
      private static final Logger LOG = LoggerFactory.getLogger(ZipFileMigrationImporter.class);
  
      private static final String APPLICATION_PROPERTY_MIGRATION_NUMER_OF_WORKERS = "atlas.migration.mode.workers";
      private static final String APPLICATION_PROPERTY_MIGRATION_BATCH_SIZE       = "atlas.migration.mode.batch.size";
      private static final String DEFAULT_NUMBER_OF_WORKERS                       = "4";
      private static final String DEFAULT_BATCH_SIZE                              = "100";
      private static final String ZIP_FILE_COMMENT_ENTITIES_COUNT                 = "entitiesCount";
      private static final String ZIP_FILE_COMMENT_TOTAL_COUNT                    = "total";
      private static final String FILE_EXTENSION_ZIP                              = ".zip";
      private static final String ENV_USER_NAME                                   = "user.name";
      private static final String ARCHIVE_DIR                                     = "archive";
  
      private final ImportService              importService;
      private       List<String>               filesToImport;
      private final DataMigrationStatusService dataMigrationStatusService;
      private       File                       archiveDir;
  
      /**
       * Input:
       * fileName : can support wildcards. If it contains wildcards then all matching files will be imported
       */
      public ZipFileMigrationImporter(ImportService importService, String fileName) {
          this.importService              = importService;
          this.dataMigrationStatusService = new DataMigrationStatusService(AtlasGraphProvider.getGraphInstance());
  
          initialize(fileName);
      }
  
      @Override
      public void run() {
          for (String fileToImport : filesToImport) {
              try {
                  detectFileToImport(fileToImport);
  
                  int streamSize = getStreamSizeFromComment(fileToImport);
  
                  MigrationImportStatus migrationImportStatus = getCreateMigrationStatus(fileToImport, streamSize);
  
                  performImport(fileToImport, streamSize, Long.toString(migrationImportStatus.getCurrentIndex()));
  
                  dataMigrationStatusService.setStatus("DONE");
  
                  moveZipFileToArchiveDir(fileToImport);
              } catch (IOException e) {
                  LOG.error("Migration Import: IO Error!", e);
  
                  dataMigrationStatusService.setStatus("FAIL");
              } catch (AtlasBaseException e) {
                  LOG.error("Migration Import: Error!", e);
  
                  dataMigrationStatusService.setStatus("FAIL");
              }
          }
      }
  
      private void initialize(String fileName) {
          this.filesToImport = getAllFilesToImport(fileName);
  
          if (CollectionUtils.isNotEmpty(this.filesToImport)) {
              createArchiveDirectory(fileName);
          }
      }
  
      /**
       * Input:
       * fileName : If it contains wildcards then all matching files will be discovered
       */
      private List<String> getAllFilesToImport(String fileName) {
          List<String> ret          = new ArrayList<>();
          File         fileToImport = new File(fileName);
  
          if (fileToImport.exists() && fileToImport.isFile()) {
              //Input file present so no need to expand
              LOG.info("Migration Import: zip file for import: {}", fileToImport);
  
              ret.add(fileToImport.getAbsolutePath());
          } else {
              //The fileName might have wildcard
              String dirPath       = new File(fileToImport.getParent()).getAbsolutePath();
              File   importDataDir = new File(dirPath);
  
              if (importDataDir.exists() && importDataDir.isDirectory()) {
                  String     fileNameWithWildcard = fileToImport.getName();
                  FileFilter fileFilter           = new WildcardFileFilter(fileNameWithWildcard);
                  File[]     importFiles          = importDataDir.listFiles(fileFilter);
  
                  if (ArrayUtils.isNotEmpty(importFiles)) {
                      Arrays.sort(importFiles);
  
                      LOG.info("Migration Import: zip files for import: ");
  
                      for (File importFile : importFiles) {
                          if (isValidImportFile(importFile)) {
                              LOG.info("{} with absolute path - {}", importFile.getName(), importFile.getAbsolutePath());
  
                              ret.add(importFile.getAbsolutePath());
                          } else {
                              LOG.warn("Ignoring {} as it is not a file or does not end with extension {}", FILE_EXTENSION_ZIP, importFile.getAbsolutePath());
                          }
                      }
                  } else {
                      LOG.warn("Migration Import: No files to import");
                  }
              }
          }
  
          return ret;
      }
  
      private boolean isValidImportFile(File importFile) {
          return importFile.isFile() && StringUtils.endsWithIgnoreCase(importFile.getName(), FILE_EXTENSION_ZIP);
      }
  
      private void createArchiveDirectory(String fileName) {
          File   fileToImport = new File(fileName);
          String parentPath   = new File(fileToImport.getParent()).getAbsolutePath();
  
          this.archiveDir = new File(parentPath + File.separator + ARCHIVE_DIR);
  
          if (this.archiveDir.exists() && !this.archiveDir.canWrite()) {
              LOG.warn("Migration Import: No write permission to archive directory {}", this.archiveDir.getAbsolutePath());
  
              this.archiveDir = null;
          } else if (!this.archiveDir.exists() && !this.archiveDir.getParentFile().canWrite()) {
              LOG.warn("Migration Import: No permission to create archive directory {}", this.archiveDir.getAbsolutePath());
  
              this.archiveDir = null;
          } else {
              this.archiveDir.mkdirs();
  
              LOG.info("Migration Import: archive directory for zip files: {}", this.archiveDir.getAbsolutePath());
          }
      }
  
      private void moveZipFileToArchiveDir(String srcFilePath) {
          if (this.archiveDir == null) {
              return;
          }
  
          File   sourceFile = new File(srcFilePath);
          String newFile    = archiveDir.getAbsolutePath() + File.separator + sourceFile.getName();
  
          if (!sourceFile.canWrite()) {
              LOG.warn("Migration Import: No permission to archive the zip file {}", sourceFile.getAbsolutePath());
  
              this.archiveDir = null;
          } else {
              if (sourceFile.renameTo(new File(newFile))) {
                  sourceFile.delete();
  
                  LOG.info("Migration Import: Successfully archived the zip file: {} to {}", srcFilePath, this.archiveDir.getAbsolutePath());
              } else {
                  LOG.warn("Migration Import: Failed to archive the zip file: {}", srcFilePath);
              }
          }
      }
  
      private MigrationImportStatus getCreateMigrationStatus(String fileName, int streamSize) {
          MigrationImportStatus status = null;
  
          try {
              status = new MigrationImportStatus(fileName, DigestUtils.md5Hex(new FileInputStream(fileName)));
  
              status.setTotalCount(streamSize);
          } catch (IOException e) {
              LOG.error("Exception occurred while creating migration import", e);
          }
  
          MigrationImportStatus statusRetrieved = dataMigrationStatusService.getCreate(status);
  
          LOG.info("DataMigrationStatusService: Position: {}", statusRetrieved.getCurrentIndex());
  
          dataMigrationStatusService.setStatus("STARTED");
  
          return statusRetrieved;
      }
  
      private void detectFileToImport(String fileToImport) throws IOException {
          FileWatcher fileWatcher = new FileWatcher(fileToImport);
  
          fileWatcher.start();
      }
  
      private int getStreamSizeFromComment(String fileToImport) {
          int ret = 1;
  
          try {
              ZipFile zipFile = new ZipFile(fileToImport);
              String  comment = zipFile.getComment();
  
              ret = processZipFileStreamSizeComment(comment);
  
              zipFile.close();
          } catch (IOException e) {
              LOG.error("Error opening ZIP file: {}", fileToImport, e);
          }
  
          return ret;
      }
  
      private int processZipFileStreamSizeComment(String comment) {
          if (StringUtils.isEmpty(comment)) {
              return 1;
          }
  
          Map map           = AtlasType.fromJson(comment, Map.class);
          int entitiesCount = (int) map.get(ZIP_FILE_COMMENT_ENTITIES_COUNT);
          int totalCount    = (int) map.get(ZIP_FILE_COMMENT_TOTAL_COUNT);
  
          LOG.info("ZipFileMigrationImporter: Zip file: Comment: streamSize: {}: total: {}", entitiesCount, totalCount);
  
          return entitiesCount;
      }
  
      private void performImport(String fileToImport, int streamSize, String startPosition) throws AtlasBaseException {
          try {
              LOG.info("Migration Import: {}: Starting at: {}...", fileToImport, startPosition);
  
              InputStream fs = new FileInputStream(fileToImport);
  
              RequestContext.get().setUser(getUserNameFromEnvironment(), null);
  
              importService.run(fs, getImportRequest(fileToImport, streamSize, startPosition),
                      getUserNameFromEnvironment(), InetAddress.getLocalHost().getHostName(), InetAddress.getLocalHost().getHostAddress());
          } catch (Exception ex) {
              LOG.error("Migration Import: Error loading zip for migration!", ex);
  
              throw new AtlasBaseException(ex);
          } finally {
              LOG.info("Migration Import: {}: Done!", fileToImport);
          }
      }
  
      private String getUserNameFromEnvironment() {
          return System.getProperty(ENV_USER_NAME);
      }
  
      private AtlasImportRequest getImportRequest(String fileToImport, int streamSize, String position) throws AtlasException {
          AtlasImportRequest request = new AtlasImportRequest();
  
          request.setOption(AtlasImportRequest.OPTION_KEY_MIGRATION_FILE_NAME, fileToImport);
          request.setSizeOption(streamSize);
          request.setOption(AtlasImportRequest.OPTION_KEY_MIGRATION, "true");
          request.setOption(AtlasImportRequest.OPTION_KEY_NUM_WORKERS, getPropertyValue(APPLICATION_PROPERTY_MIGRATION_NUMER_OF_WORKERS, DEFAULT_NUMBER_OF_WORKERS));
          request.setOption(AtlasImportRequest.OPTION_KEY_BATCH_SIZE, getPropertyValue(APPLICATION_PROPERTY_MIGRATION_BATCH_SIZE, DEFAULT_BATCH_SIZE));
          request.setOption(AtlasImportRequest.START_POSITION_KEY, (StringUtils.isEmpty(position) ? "0" : position));
  
          return request;
      }
  
      private String getPropertyValue(String property, String defaultValue) throws AtlasException {
          return ApplicationProperties.get().getString(property, defaultValue);
      }
  }
  