code: /*
   * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.integrations.destination.redshift.copiers;
  
  import com.amazonaws.services.s3.AmazonS3;
  import com.fasterxml.jackson.databind.ObjectMapper;
  import com.google.common.annotations.VisibleForTesting;
  import io.airbyte.commons.lang.Exceptions;
  import io.airbyte.db.jdbc.JdbcDatabase;
  import io.airbyte.integrations.destination.ExtendedNameTransformer;
  import io.airbyte.integrations.destination.jdbc.SqlOperations;
  import io.airbyte.integrations.destination.jdbc.copy.s3.S3CopyConfig;
  import io.airbyte.integrations.destination.jdbc.copy.s3.S3StreamCopier;
  import io.airbyte.integrations.destination.redshift.manifest.Entry;
  import io.airbyte.integrations.destination.redshift.manifest.Manifest;
  import io.airbyte.integrations.destination.s3.S3DestinationConfig;
  import io.airbyte.integrations.destination.s3.credential.S3AccessKeyCredentialConfig;
  import io.airbyte.protocol.models.ConfiguredAirbyteStream;
  import java.sql.Timestamp;
  import java.time.Instant;
  import java.util.Optional;
  import java.util.UUID;
  import java.util.stream.Collectors;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  public class RedshiftStreamCopier extends S3StreamCopier {
  
    private static final Logger LOGGER = LoggerFactory.getLogger(RedshiftStreamCopier.class);
    // From https://docs.aws.amazon.com/redshift/latest/dg/t_loading-tables-from-s3.html
    // "Split your load data files so that the files are about equal size, between 1 MB and 1 GB after
    // compression"
    public static final int MAX_PARTS_PER_FILE = 4;
  
    private final ObjectMapper objectMapper;
    private String manifestFilePath = null;
  
    public RedshiftStreamCopier(final String stagingFolder,
                                final String schema,
                                final AmazonS3 client,
                                final JdbcDatabase db,
                                final S3CopyConfig config,
                                final ExtendedNameTransformer nameTransformer,
                                final SqlOperations sqlOperations,
                                final ConfiguredAirbyteStream configuredAirbyteStream) {
      this(
          stagingFolder,
          schema,
          client,
          db,
          config,
          nameTransformer,
          sqlOperations,
          Timestamp.from(Instant.now()),
          configuredAirbyteStream);
    }
  
    @VisibleForTesting
    RedshiftStreamCopier(final String stagingFolder,
                         final String schema,
                         final AmazonS3 client,
                         final JdbcDatabase db,
                         final S3CopyConfig config,
                         final ExtendedNameTransformer nameTransformer,
                         final SqlOperations sqlOperations,
                         final Timestamp uploadTime,
                         final ConfiguredAirbyteStream configuredAirbyteStream) {
      super(stagingFolder,
          schema,
          client,
          db,
          config,
          nameTransformer,
          sqlOperations,
          configuredAirbyteStream,
          uploadTime,
          MAX_PARTS_PER_FILE);
      objectMapper = new ObjectMapper();
    }
  
    @Override
    public void copyStagingFileToTemporaryTable() {
      final var possibleManifest = Optional.ofNullable(createManifest());
      LOGGER.info("Starting copy to tmp table: {} in destination for stream: {}, schema: {}, .", tmpTableName, streamName, schemaName);
      possibleManifest.stream()
          .map(this::putManifest)
          .forEach(this::executeCopy);
      LOGGER.info("Copy to tmp table {} in destination for stream {} complete.", tmpTableName, streamName);
    }
  
    @Override
    public void copyS3CsvFileIntoTable(final JdbcDatabase database,
                                       final String s3FileLocation,
                                       final String schema,
                                       final String tableName,
                                       final S3DestinationConfig s3Config) {
      throw new RuntimeException("Redshift Stream Copier should not copy individual files without use of a manifest");
    }
  
    @Override
    public void removeFileAndDropTmpTable() throws Exception {
      super.removeFileAndDropTmpTable();
      if (manifestFilePath != null) {
        LOGGER.info("Begin cleaning s3 manifest file {}.", manifestFilePath);
        if (s3Client.doesObjectExist(s3Config.getBucketName(), manifestFilePath)) {
          s3Client.deleteObject(s3Config.getBucketName(), manifestFilePath);
        }
        LOGGER.info("S3 manifest file {} cleaned.", manifestFilePath);
      }
    }
  
    /**
     * Creates the contents of a manifest file given the `s3StagingFiles`. There must be at least one
     * entry in a manifest file otherwise it is not considered valid for the COPY command.
     *
     * @return null if no stagingFiles exist otherwise the manifest body String
     */
    private String createManifest() {
      if (getStagingFiles().isEmpty()) {
        return null;
      }
  
      final var s3FileEntries = getStagingFiles().stream()
          .map(filePath -> new Entry(getFullS3Path(s3Config.getBucketName(), filePath)))
          .collect(Collectors.toList());
      final var manifest = new Manifest(s3FileEntries);
  
      return Exceptions.toRuntime(() -> objectMapper.writeValueAsString(manifest));
    }
  
    /**
     * Upload the supplied manifest file to S3
     *
     * @param manifestContents the manifest contents, never null
     * @return the path where the manifest file was placed in S3
     */
    private String putManifest(final String manifestContents) {
      manifestFilePath =
          String.join("/", s3Config.getBucketPath(), stagingFolder, schemaName, String.format("%s.manifest", UUID.randomUUID()));
  
      s3Client.putObject(s3Config.getBucketName(), manifestFilePath, manifestContents);
  
      return manifestFilePath;
    }
  
    /**
     * Run Redshift COPY command with the given manifest file
     *
     * @param manifestPath the path in S3 to the manifest file
     */
    private void executeCopy(final String manifestPath) {
      final S3AccessKeyCredentialConfig credentialConfig = (S3AccessKeyCredentialConfig) s3Config.getS3CredentialConfig();
      final var copyQuery = String.format(
          "COPY %s.%s FROM '%s'\n"
              + "CREDENTIALS 'aws_access_key_id=%s;aws_secret_access_key=%s'\n"
              + "CSV REGION '%s' TIMEFORMAT 'auto'\n"
              + "STATUPDATE OFF\n"
              + "MANIFEST;",
          schemaName,
          tmpTableName,
          getFullS3Path(s3Config.getBucketName(), manifestPath),
          credentialConfig.getAccessKeyId(),
          credentialConfig.getSecretAccessKey(),
          s3Config.getBucketRegion());
  
      Exceptions.toRuntime(() -> db.execute(copyQuery));
    }
  
  }
  