code: /*
   * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.integrations.destination.s3.csv;
  
  import alex.mojaki.s3upload.MultiPartOutputStream;
  import alex.mojaki.s3upload.StreamTransferManager;
  import com.amazonaws.services.s3.AmazonS3;
  import com.fasterxml.jackson.databind.JsonNode;
  import io.airbyte.integrations.destination.s3.S3DestinationConfig;
  import io.airbyte.integrations.destination.s3.S3Format;
  import io.airbyte.integrations.destination.s3.template.S3FilenameTemplateParameterObject;
  import io.airbyte.integrations.destination.s3.util.StreamTransferManagerFactory;
  import io.airbyte.integrations.destination.s3.writer.BaseS3Writer;
  import io.airbyte.integrations.destination.s3.writer.DestinationFileWriter;
  import io.airbyte.protocol.models.AirbyteRecordMessage;
  import io.airbyte.protocol.models.ConfiguredAirbyteStream;
  import java.io.IOException;
  import java.io.PrintWriter;
  import java.nio.charset.StandardCharsets;
  import java.sql.Timestamp;
  import java.util.UUID;
  import org.apache.commons.csv.CSVFormat;
  import org.apache.commons.csv.CSVPrinter;
  import org.apache.commons.csv.QuoteMode;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  public class S3CsvWriter extends BaseS3Writer implements DestinationFileWriter {
  
    private static final Logger LOGGER = LoggerFactory.getLogger(S3CsvWriter.class);
  
    private final CsvSheetGenerator csvSheetGenerator;
    private final StreamTransferManager uploadManager;
    private final MultiPartOutputStream outputStream;
    private final CSVPrinter csvPrinter;
    private final String objectKey;
    private final String gcsFileLocation;
  
    private S3CsvWriter(final S3DestinationConfig config,
                        final AmazonS3 s3Client,
                        final ConfiguredAirbyteStream configuredStream,
                        final Timestamp uploadTimestamp,
                        final int uploadThreads,
                        final int queueCapacity,
                        final boolean writeHeader,
                        CSVFormat csvSettings,
                        final CsvSheetGenerator csvSheetGenerator)
        throws IOException {
      super(config, s3Client, configuredStream);
  
      this.csvSheetGenerator = csvSheetGenerator;
  
      final String fileSuffix = "_" + UUID.randomUUID();
      final String outputFilename = determineOutputFilename(S3FilenameTemplateParameterObject
          .builder()
          .customSuffix(fileSuffix)
          .s3Format(S3Format.CSV)
          .fileExtension(S3Format.CSV.getFileExtension())
          .fileNamePattern(config.getFileNamePattern())
          .timestamp(uploadTimestamp)
          .build());
      this.objectKey = String.join("/", outputPrefix, outputFilename);
  
      LOGGER.info("Full S3 path for stream '{}': s3://{}/{}", stream.getName(), config.getBucketName(),
          objectKey);
      gcsFileLocation = String.format("gs://%s/%s", config.getBucketName(), objectKey);
  
      this.uploadManager = StreamTransferManagerFactory
          .create(config.getBucketName(), objectKey, s3Client)
          .get()
          .numUploadThreads(uploadThreads)
          .queueCapacity(queueCapacity);
      // We only need one output stream as we only have one input stream. This is reasonably performant.
      this.outputStream = uploadManager.getMultiPartOutputStreams().get(0);
      if (writeHeader) {
        csvSettings = csvSettings.withHeader(csvSheetGenerator.getHeaderRow().toArray(new String[0]));
      }
      this.csvPrinter = new CSVPrinter(new PrintWriter(outputStream, true, StandardCharsets.UTF_8), csvSettings);
    }
  
    public static class Builder {
  
      private final S3DestinationConfig config;
      private final AmazonS3 s3Client;
      private final ConfiguredAirbyteStream configuredStream;
      private final Timestamp uploadTimestamp;
      private int uploadThreads = StreamTransferManagerFactory.DEFAULT_UPLOAD_THREADS;
      private int queueCapacity = StreamTransferManagerFactory.DEFAULT_QUEUE_CAPACITY;
      private boolean withHeader = true;
      private CSVFormat csvSettings = CSVFormat.DEFAULT.withQuoteMode(QuoteMode.ALL);
      private CsvSheetGenerator csvSheetGenerator;
  
      public Builder(final S3DestinationConfig config,
                     final AmazonS3 s3Client,
                     final ConfiguredAirbyteStream configuredStream,
                     final Timestamp uploadTimestamp) {
        this.config = config;
        this.s3Client = s3Client;
        this.configuredStream = configuredStream;
        this.uploadTimestamp = uploadTimestamp;
      }
  
      public Builder uploadThreads(final int uploadThreads) {
        this.uploadThreads = uploadThreads;
        return this;
      }
  
      public Builder queueCapacity(final int queueCapacity) {
        this.queueCapacity = queueCapacity;
        return this;
      }
  
      public Builder withHeader(final boolean withHeader) {
        this.withHeader = withHeader;
        return this;
      }
  
      public Builder csvSettings(final CSVFormat csvSettings) {
        this.csvSettings = csvSettings;
        return this;
      }
  
      public Builder csvSheetGenerator(final CsvSheetGenerator csvSheetGenerator) {
        this.csvSheetGenerator = csvSheetGenerator;
        return this;
      }
  
      public S3CsvWriter build() throws IOException {
        if (csvSheetGenerator == null) {
          final S3CsvFormatConfig formatConfig = (S3CsvFormatConfig) config.getFormatConfig();
          csvSheetGenerator = CsvSheetGenerator.Factory.create(configuredStream.getStream().getJsonSchema(), formatConfig);
        }
        return new S3CsvWriter(config,
            s3Client,
            configuredStream,
            uploadTimestamp,
            uploadThreads,
            queueCapacity,
            withHeader,
            csvSettings,
            csvSheetGenerator);
      }
  
    }
  
    @Override
    public void write(final UUID id, final AirbyteRecordMessage recordMessage) throws IOException {
      csvPrinter.printRecord(csvSheetGenerator.getDataRow(id, recordMessage));
    }
  
    @Override
    protected void closeWhenSucceed() throws IOException {
      csvPrinter.close();
      outputStream.close();
      uploadManager.complete();
    }
  
    @Override
    protected void closeWhenFail() throws IOException {
      csvPrinter.close();
      outputStream.close();
      uploadManager.abort();
    }
  
    @Override
    public String getOutputPath() {
      return objectKey;
    }
  
    @Override
    public String getFileLocation() {
      return gcsFileLocation;
    }
  
    @Override
    public S3Format getFileFormat() {
      return S3Format.CSV;
    }
  
    @Override
    public void write(final JsonNode formattedData) throws IOException {
      csvPrinter.printRecord(csvSheetGenerator.getDataRow(formattedData));
    }
  
  }
  