embulk-output-s3_parquet 0.1.0 → 0.2.0

data/src/main/scala/org/embulk/output/s3_parquet/S3ParquetPageOutput.scala

@@ -1,6 +1,5 @@
 package org.embulk.output.s3_parquet
 
-
 import java.io.File
 import java.nio.file.{Files, Paths}
 
@@ -11,63 +10,61 @@ import org.embulk.config.TaskReport
 import org.embulk.output.s3_parquet.aws.Aws
 import org.embulk.spi.{Exec, Page, PageReader, TransactionalPageOutput}
 
+case class S3ParquetPageOutput(
+    outputLocalFile: String,
+    reader: PageReader,
+    writer: ParquetWriter[PageReader],
+    aws: Aws,
+    destBucket: String,
+    destKey: String
+) extends TransactionalPageOutput {
 
-case class S3ParquetPageOutput(outputLocalFile: String,
-                               reader: PageReader,
-                               writer: ParquetWriter[PageReader],
-                               aws: Aws,
-                               destBucket: String,
-                               destKey: String)
-    extends TransactionalPageOutput
-{
-
-    private var isClosed: Boolean = false
+  private var isClosed: Boolean = false
 
-    override def add(page: Page): Unit =
-    {
-        reader.setPage(page)
-        while (reader.nextRecord()) {
-            writer.write(reader)
-        }
+  override def add(page: Page): Unit = {
+    reader.setPage(page)
+    while (reader.nextRecord()) {
+      writer.write(reader)
     }
+  }
 
-    override def finish(): Unit =
-    {
-    }
+  override def finish(): Unit = {}
 
-    override def close(): Unit =
-    {
-        synchronized {
-            if (!isClosed) {
-                writer.close()
-                isClosed = true
-            }
+  override def close(): Unit = {
+    synchronized {
+      if (!isClosed) {
+        ContextClassLoaderSwapper.usingPluginClass {
+          writer.close()
         }
+        isClosed = true
+      }
     }
+  }
 
-    override def abort(): Unit =
-    {
-        close()
-        cleanup()
-    }
+  override def abort(): Unit = {
+    close()
+    cleanup()
+  }
 
-    override def commit(): TaskReport =
-    {
-        close()
-        val result: UploadResult = aws.withTransferManager { xfer: TransferManager =>
-            val upload: Upload = xfer.upload(destBucket, destKey, new File(outputLocalFile))
-            upload.waitForUploadResult()
-        }
-        cleanup()
-        Exec.newTaskReport()
-            .set("bucket", result.getBucketName)
-            .set("key", result.getKey)
-            .set("etag", result.getETag)
-            .set("version_id", result.getVersionId)
+  override def commit(): TaskReport = {
+    close()
+    val result: UploadResult = ContextClassLoaderSwapper.usingPluginClass {
+      aws.withTransferManager { xfer: TransferManager =>
+        val upload: Upload =
+          xfer.upload(destBucket, destKey, new File(outputLocalFile))
+        upload.waitForUploadResult()
+      }
     }
+    cleanup()
+    Exec
+      .newTaskReport()
+      .set("bucket", result.getBucketName)
+      .set("key", result.getKey)
+      .set("etag", result.getETag)
+      .set("version_id", result.getVersionId)
+  }
 
-    private def cleanup(): Unit =
-    {
-        Files.delete(Paths.get(outputLocalFile))
-    }
+  private def cleanup(): Unit = {
+    Files.delete(Paths.get(outputLocalFile))
+  }
 }

data/src/main/scala/org/embulk/output/s3_parquet/aws/Aws.scala

@@ -1,63 +1,59 @@
 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)
-    }
-
+import com.amazonaws.services.s3.transfer.{
+  TransferManager,
+  TransferManagerBuilder
 }
 
-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()
-    }
+object Aws {
 
-    def withTransferManager[A](f: TransferManager => A): A =
-    {
-        withS3 { s3 =>
-            val svc = TransferManagerBuilder.standard().withS3Client(s3).build()
-            try f(svc)
-            finally svc.shutdownNow(false)
-        }
-    }
+  trait Task
+      extends AwsCredentials.Task
+      with AwsEndpointConfiguration.Task
+      with AwsClientConfiguration.Task
+      with AwsS3Configuration.Task
 
-    def withGlue[A](f: AWSGlue => A): A =
-    {
-        val builder: AWSGlueClientBuilder = AWSGlueClientBuilder.standard()
-        val svc = createService(builder)
-        try f(svc)
-        finally svc.shutdown()
-    }
+  def apply(task: Task): Aws = {
+    new Aws(task)
+  }
 
-    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()
+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()
+  }
 }

data/src/main/scala/org/embulk/output/s3_parquet/aws/AwsClientConfiguration.scala

@@ -1,6 +1,5 @@
 package org.embulk.output.s3_parquet.aws
 
-
 import java.util.Optional
 
 import com.amazonaws.ClientConfiguration
@@ -8,35 +7,31 @@ import com.amazonaws.client.builder.AwsClientBuilder
 import org.embulk.config.{Config, ConfigDefault}
 import org.embulk.output.s3_parquet.aws.AwsClientConfiguration.Task
 
+object AwsClientConfiguration {
 
-object AwsClientConfiguration
-{
-
-    trait Task
-    {
+  trait Task {
 
-        @Config("http_proxy")
-        @ConfigDefault("null")
-        def getHttpProxy: Optional[HttpProxy.Task]
+    @Config("http_proxy")
+    @ConfigDefault("null")
+    def getHttpProxy: Optional[HttpProxy.Task]
 
-    }
+  }
 
-    def apply(task: Task): AwsClientConfiguration =
-    {
-        new AwsClientConfiguration(task)
-    }
+  def apply(task: Task): AwsClientConfiguration = {
+    new AwsClientConfiguration(task)
+  }
 }
 
-class AwsClientConfiguration(task: Task)
-{
+class AwsClientConfiguration(task: Task) {
 
-    def configureAwsClientBuilder[S <: AwsClientBuilder[S, T], T](builder: AwsClientBuilder[S, T]): Unit =
-    {
-        task.getHttpProxy.ifPresent { v =>
-            val cc = new ClientConfiguration
-            HttpProxy(v).configureClientConfiguration(cc)
-            builder.setClientConfiguration(cc)
-        }
+  def configureAwsClientBuilder[S <: AwsClientBuilder[S, T], T](
+      builder: AwsClientBuilder[S, T]
+  ): Unit = {
+    task.getHttpProxy.ifPresent { v =>
+      val cc = new ClientConfiguration
+      HttpProxy(v).configureClientConfiguration(cc)
+      builder.setClientConfiguration(cc)
     }
+  }
 
 }

data/src/main/scala/org/embulk/output/s3_parquet/aws/AwsCredentials.scala

@@ -1,147 +1,174 @@
 package org.embulk.output.s3_parquet.aws
 
-
 import java.util.Optional
 
-import com.amazonaws.auth.{AnonymousAWSCredentials, AWSCredentialsProvider, AWSStaticCredentialsProvider, BasicAWSCredentials, BasicSessionCredentials, DefaultAWSCredentialsProviderChain, EC2ContainerCredentialsProviderWrapper, EnvironmentVariableCredentialsProvider, STSAssumeRoleSessionCredentialsProvider, SystemPropertiesCredentialsProvider, WebIdentityTokenCredentialsProvider}
-import com.amazonaws.auth.profile.{ProfileCredentialsProvider, ProfilesConfigFile}
+import com.amazonaws.auth.{
+  AnonymousAWSCredentials,
+  AWSCredentialsProvider,
+  AWSStaticCredentialsProvider,
+  BasicAWSCredentials,
+  BasicSessionCredentials,
+  DefaultAWSCredentialsProviderChain,
+  EC2ContainerCredentialsProviderWrapper,
+  EnvironmentVariableCredentialsProvider,
+  STSAssumeRoleSessionCredentialsProvider,
+  SystemPropertiesCredentialsProvider,
+  WebIdentityTokenCredentialsProvider
+}
+import com.amazonaws.auth.profile.{
+  ProfileCredentialsProvider,
+  ProfilesConfigFile
+}
 import org.embulk.config.{Config, ConfigDefault, ConfigException}
 import org.embulk.output.s3_parquet.aws.AwsCredentials.Task
 import org.embulk.spi.unit.LocalFile
 
+object AwsCredentials {
 
-object AwsCredentials
-{
+  trait Task {
 
-    trait Task
-    {
+    @Config("auth_method")
+    @ConfigDefault("\"default\"")
+    def getAuthMethod: String
 
-        @Config("auth_method")
-        @ConfigDefault("\"default\"")
-        def getAuthMethod: String
+    @Config("access_key_id")
+    @ConfigDefault("null")
+    def getAccessKeyId: Optional[String]
 
-        @Config("access_key_id")
-        @ConfigDefault("null")
-        def getAccessKeyId: Optional[String]
+    @Config("secret_access_key")
+    @ConfigDefault("null")
+    def getSecretAccessKey: Optional[String]
 
-        @Config("secret_access_key")
-        @ConfigDefault("null")
-        def getSecretAccessKey: Optional[String]
+    @Config("session_token")
+    @ConfigDefault("null")
+    def getSessionToken: Optional[String]
 
-        @Config("session_token")
-        @ConfigDefault("null")
-        def getSessionToken: Optional[String]
+    @Config("profile_file")
+    @ConfigDefault("null")
+    def getProfileFile: Optional[LocalFile]
 
-        @Config("profile_file")
-        @ConfigDefault("null")
-        def getProfileFile: Optional[LocalFile]
+    @Config("profile_name")
+    @ConfigDefault("\"default\"")
+    def getProfileName: String
 
-        @Config("profile_name")
-        @ConfigDefault("\"default\"")
-        def getProfileName: String
+    @Config("role_arn")
+    @ConfigDefault("null")
+    def getRoleArn: Optional[String]
 
-        @Config("role_arn")
-        @ConfigDefault("null")
-        def getRoleArn: Optional[String]
+    @Config("role_session_name")
+    @ConfigDefault("null")
+    def getRoleSessionName: Optional[String]
 
-        @Config("role_session_name")
-        @ConfigDefault("null")
-        def getRoleSessionName: Optional[String]
+    @Config("role_external_id")
+    @ConfigDefault("null")
+    def getRoleExternalId: Optional[String]
 
-        @Config("role_external_id")
-        @ConfigDefault("null")
-        def getRoleExternalId: Optional[String]
+    @Config("role_session_duration_seconds")
+    @ConfigDefault("null")
+    def getRoleSessionDurationSeconds: Optional[Int]
 
-        @Config("role_session_duration_seconds")
-        @ConfigDefault("null")
-        def getRoleSessionDurationSeconds: Optional[Int]
+    @Config("scope_down_policy")
+    @ConfigDefault("null")
+    def getScopeDownPolicy: Optional[String]
 
-        @Config("scope_down_policy")
-        @ConfigDefault("null")
-        def getScopeDownPolicy: Optional[String]
-
-        @Config("web_identity_token_file")
-        @ConfigDefault("null")
-        def getWebIdentityTokenFile: Optional[String]
-    }
+    @Config("web_identity_token_file")
+    @ConfigDefault("null")
+    def getWebIdentityTokenFile: Optional[String]
+  }
 
-    def apply(task: Task): AwsCredentials =
-    {
-        new AwsCredentials(task)
-    }
+  def apply(task: Task): AwsCredentials = {
+    new AwsCredentials(task)
+  }
 }
 
-class AwsCredentials(task: Task)
-{
-
-    def createAwsCredentialsProvider: AWSCredentialsProvider =
-    {
-        task.getAuthMethod match {
-            case "basic" =>
-                new AWSStaticCredentialsProvider(new BasicAWSCredentials(
-                    getRequiredOption(task.getAccessKeyId, "access_key_id"),
-                    getRequiredOption(task.getAccessKeyId, "secret_access_key")
-                    ))
-
-            case "env" =>
-                new EnvironmentVariableCredentialsProvider
-
-            case "instance" =>
-                // NOTE: combination of InstanceProfileCredentialsProvider and ContainerCredentialsProvider
-                new EC2ContainerCredentialsProviderWrapper
-
-            case "profile" =>
-                if (task.getProfileFile.isPresent) {
-                    val pf: ProfilesConfigFile = new ProfilesConfigFile(task.getProfileFile.get().getFile)
-                    new ProfileCredentialsProvider(pf, task.getProfileName)
-                }
-                else new ProfileCredentialsProvider(task.getProfileName)
-
-            case "properties" =>
-                new SystemPropertiesCredentialsProvider
-
-            case "anonymous" =>
-                new AWSStaticCredentialsProvider(new AnonymousAWSCredentials)
-
-            case "session" =>
-                new AWSStaticCredentialsProvider(new BasicSessionCredentials(
-                    getRequiredOption(task.getAccessKeyId, "access_key_id"),
-                    getRequiredOption(task.getSecretAccessKey, "secret_access_key"),
-                    getRequiredOption(task.getSessionToken, "session_token")
-                    ))
-
-            case "assume_role" =>
-                // NOTE: Are http_proxy, endpoint, region required when assuming role?
-                val builder = new STSAssumeRoleSessionCredentialsProvider.Builder(
-                    getRequiredOption(task.getRoleArn, "role_arn"),
-                    getRequiredOption(task.getRoleSessionName, "role_session_name")
-                    )
-                task.getRoleExternalId.ifPresent(v => builder.withExternalId(v))
-                task.getRoleSessionDurationSeconds.ifPresent(v => builder.withRoleSessionDurationSeconds(v))
-                task.getScopeDownPolicy.ifPresent(v => builder.withScopeDownPolicy(v))
-
-                builder.build()
-
-            case "web_identity_token" =>
-                WebIdentityTokenCredentialsProvider.builder()
-                    .roleArn(getRequiredOption(task.getRoleArn, "role_arn"))
-                    .roleSessionName(getRequiredOption(task.getRoleSessionName, "role_session_name"))
-                    .webIdentityTokenFile(getRequiredOption(task.getWebIdentityTokenFile, "web_identity_token_file"))
-                    .build()
-
-            case "default" =>
-                new DefaultAWSCredentialsProviderChain
-
-            case am =>
-                throw new ConfigException(s"'$am' is unsupported: `auth_method` must be one of ['basic', 'env', 'instance', 'profile', 'properties', 'anonymous', 'session', 'assume_role', 'default'].")
+class AwsCredentials(task: Task) {
+
+  def createAwsCredentialsProvider: AWSCredentialsProvider = {
+    task.getAuthMethod match {
+      case "basic" =>
+        new AWSStaticCredentialsProvider(
+          new BasicAWSCredentials(
+            getRequiredOption(task.getAccessKeyId, "access_key_id"),
+            getRequiredOption(task.getSecretAccessKey, "secret_access_key")
+          )
+        )
+
+      case "env" =>
+        new EnvironmentVariableCredentialsProvider
+
+      case "instance" =>
+        // NOTE: combination of InstanceProfileCredentialsProvider and ContainerCredentialsProvider
+        new EC2ContainerCredentialsProviderWrapper
+
+      case "profile" =>
+        if (task.getProfileFile.isPresent) {
+          val pf: ProfilesConfigFile = new ProfilesConfigFile(
+            task.getProfileFile.get().getFile
+          )
+          new ProfileCredentialsProvider(pf, task.getProfileName)
         }
+        else new ProfileCredentialsProvider(task.getProfileName)
+
+      case "properties" =>
+        new SystemPropertiesCredentialsProvider
+
+      case "anonymous" =>
+        new AWSStaticCredentialsProvider(new AnonymousAWSCredentials)
+
+      case "session" =>
+        new AWSStaticCredentialsProvider(
+          new BasicSessionCredentials(
+            getRequiredOption(task.getAccessKeyId, "access_key_id"),
+            getRequiredOption(task.getSecretAccessKey, "secret_access_key"),
+            getRequiredOption(task.getSessionToken, "session_token")
+          )
+        )
+
+      case "assume_role" =>
+        // NOTE: Are http_proxy, endpoint, region required when assuming role?
+        val builder = new STSAssumeRoleSessionCredentialsProvider.Builder(
+          getRequiredOption(task.getRoleArn, "role_arn"),
+          getRequiredOption(task.getRoleSessionName, "role_session_name")
+        )
+        task.getRoleExternalId.ifPresent(v => builder.withExternalId(v))
+        task.getRoleSessionDurationSeconds.ifPresent(v =>
+          builder.withRoleSessionDurationSeconds(v)
+        )
+        task.getScopeDownPolicy.ifPresent(v => builder.withScopeDownPolicy(v))
+
+        builder.build()
+
+      case "web_identity_token" =>
+        WebIdentityTokenCredentialsProvider
+          .builder()
+          .roleArn(getRequiredOption(task.getRoleArn, "role_arn"))
+          .roleSessionName(
+            getRequiredOption(task.getRoleSessionName, "role_session_name")
+          )
+          .webIdentityTokenFile(
+            getRequiredOption(
+              task.getWebIdentityTokenFile,
+              "web_identity_token_file"
+            )
+          )
+          .build()
+
+      case "default" =>
+        new DefaultAWSCredentialsProviderChain
+
+      case am =>
+        throw new ConfigException(
+          s"'$am' is unsupported: `auth_method` must be one of ['basic', 'env', 'instance', 'profile', 'properties', 'anonymous', 'session', 'assume_role', 'default']."
+        )
     }
-
-    private def getRequiredOption[A](o: Optional[A],
-                                     name: String): A =
-    {
-        o.orElseThrow(() => new ConfigException(s"`$name` must be set when `auth_method` is ${task.getAuthMethod}."))
-    }
-
+  }
+
+  private def getRequiredOption[A](o: Optional[A], name: String): A = {
+    o.orElseThrow(() =>
+      new ConfigException(
+        s"`$name` must be set when `auth_method` is ${task.getAuthMethod}."
+      )
+    )
+  }
 
 }