code: /*
   * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.integrations.destination.s3.parquet;
  
  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.avro.AvroRecordFactory;
  import io.airbyte.integrations.destination.s3.credential.S3AccessKeyCredentialConfig;
  import io.airbyte.integrations.destination.s3.template.S3FilenameTemplateParameterObject;
  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.net.URI;
  import java.net.URISyntaxException;
  import java.sql.Timestamp;
  import java.util.UUID;
  import org.apache.avro.Schema;
  import org.apache.avro.generic.GenericData;
  import org.apache.avro.generic.GenericData.Record;
  import org.apache.hadoop.conf.Configuration;
  import org.apache.hadoop.fs.Path;
  import org.apache.hadoop.fs.s3a.Constants;
  import org.apache.parquet.avro.AvroParquetWriter;
  import org.apache.parquet.hadoop.ParquetWriter;
  import org.apache.parquet.hadoop.util.HadoopOutputFile;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import tech.allegro.schema.json2avro.converter.JsonAvroConverter;
  
  public class S3ParquetWriter extends BaseS3Writer implements DestinationFileWriter {
  
    private static final Logger LOGGER = LoggerFactory.getLogger(S3ParquetWriter.class);
  
    private final ParquetWriter<Record> parquetWriter;
    private final AvroRecordFactory avroRecordFactory;
    private final Schema schema;
    private final String outputFilename;
    private final String objectKey;
    private final String gcsFileLocation;
  
    public S3ParquetWriter(final S3DestinationConfig config,
                           final AmazonS3 s3Client,
                           final ConfiguredAirbyteStream configuredStream,
                           final Timestamp uploadTimestamp,
                           final Schema schema,
                           final JsonAvroConverter converter)
        throws URISyntaxException, IOException {
      super(config, s3Client, configuredStream);
  
      outputFilename = determineOutputFilename(S3FilenameTemplateParameterObject
          .builder()
          .s3Format(S3Format.PARQUET)
          .timestamp(uploadTimestamp)
          .fileExtension(S3Format.PARQUET.getFileExtension())
          .fileNamePattern(config.getFileNamePattern())
          .build());
  
      objectKey = String.join("/", outputPrefix, outputFilename);
  
      LOGGER.info("Full S3 path for stream '{}': s3://{}/{}", stream.getName(), config.getBucketName(), objectKey);
      gcsFileLocation = String.format("s3a://%s/%s/%s", config.getBucketName(), outputPrefix, outputFilename);
  
      final URI uri = new URI(
          String.format("s3a://%s/%s/%s", config.getBucketName(), outputPrefix, outputFilename));
      final Path path = new Path(uri);
  
      final S3ParquetFormatConfig formatConfig = (S3ParquetFormatConfig) config.getFormatConfig();
      final Configuration hadoopConfig = getHadoopConfig(config);
      this.parquetWriter = AvroParquetWriter.<GenericData.Record>builder(HadoopOutputFile.fromPath(path, hadoopConfig))
          .withSchema(schema)
          .withCompressionCodec(formatConfig.getCompressionCodec())
          .withRowGroupSize(formatConfig.getBlockSize())
          .withMaxPaddingSize(formatConfig.getMaxPaddingSize())
          .withPageSize(formatConfig.getPageSize())
          .withDictionaryPageSize(formatConfig.getDictionaryPageSize())
          .withDictionaryEncoding(formatConfig.isDictionaryEncoding())
          .build();
      this.avroRecordFactory = new AvroRecordFactory(schema, converter);
      this.schema = schema;
    }
  
    public static Configuration getHadoopConfig(final S3DestinationConfig config) {
      final Configuration hadoopConfig = new Configuration();
      final S3AccessKeyCredentialConfig credentialConfig = (S3AccessKeyCredentialConfig) config.getS3CredentialConfig();
      hadoopConfig.set(Constants.ACCESS_KEY, credentialConfig.getAccessKeyId());
      hadoopConfig.set(Constants.SECRET_KEY, credentialConfig.getSecretAccessKey());
      if (config.getEndpoint().isEmpty()) {
        hadoopConfig.set(Constants.ENDPOINT, String.format("s3.%s.amazonaws.com", config.getBucketRegion()));
      } else {
        hadoopConfig.set(Constants.ENDPOINT, config.getEndpoint());
        hadoopConfig.set(Constants.PATH_STYLE_ACCESS, "true");
      }
      hadoopConfig.set(Constants.AWS_CREDENTIALS_PROVIDER,
          "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider");
      return hadoopConfig;
    }
  
    public Schema getSchema() {
      return schema;
    }
  
    /**
     * The file path includes prefix and filename, but does not include the bucket name.
     */
    public String getOutputFilePath() {
      return outputPrefix + "/" + outputFilename;
    }
  
    public String getOutputFilename() {
      return outputFilename;
    }
  
    @Override
    public void write(final UUID id, final AirbyteRecordMessage recordMessage) throws IOException {
      parquetWriter.write(avroRecordFactory.getAvroRecord(id, recordMessage));
    }
  
    @Override
    protected void closeWhenSucceed() throws IOException {
      parquetWriter.close();
    }
  
    @Override
    protected void closeWhenFail() throws IOException {
      parquetWriter.close();
    }
  
    @Override
    public String getOutputPath() {
      return objectKey;
    }
  
    @Override
    public String getFileLocation() {
      return gcsFileLocation;
    }
  
    @Override
    public S3Format getFileFormat() {
      return S3Format.PARQUET;
    }
  
    @Override
    public void write(final JsonNode formattedData) throws IOException {
      parquetWriter.write(avroRecordFactory.getAvroRecord(formattedData));
    }
  
  }
  