package org.embulk.output.s3_parquet.aws import com.amazonaws.client.builder.AwsClientBuilder import com.amazonaws.services.glue.{AWSGlue, AWSGlueClientBuilder} import com.amazonaws.services.s3.{AmazonS3, AmazonS3ClientBuilder} import com.amazonaws.services.s3.transfer.{ TransferManager, TransferManagerBuilder } object Aws { trait Task extends AwsCredentials.Task with AwsEndpointConfiguration.Task with AwsClientConfiguration.Task with AwsS3Configuration.Task def apply(task: Task): Aws = { new Aws(task) } } class Aws(task: Aws.Task) { def withS3[A](f: AmazonS3 => A): A = { val builder: AmazonS3ClientBuilder = AmazonS3ClientBuilder.standard() AwsS3Configuration(task).configureAmazonS3ClientBuilder(builder) val svc = createService(builder) try f(svc) finally svc.shutdown() } def withTransferManager[A](f: TransferManager => A): A = { withS3 { s3 => val svc = TransferManagerBuilder.standard().withS3Client(s3).build() try f(svc) finally svc.shutdownNow(false) } } def withGlue[A](f: AWSGlue => A): A = { val builder: AWSGlueClientBuilder = AWSGlueClientBuilder.standard() val svc = createService(builder) try f(svc) finally svc.shutdown() } def createService[S <: AwsClientBuilder[S, T], T]( builder: AwsClientBuilder[S, T] ): T = { AwsEndpointConfiguration(task).configureAwsClientBuilder(builder) AwsClientConfiguration(task).configureAwsClientBuilder(builder) builder.setCredentials(AwsCredentials(task).createAwsCredentialsProvider) builder.build() } }