embulk-output-s3_parquet 0.0.2 → 0.0.3

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

@@ -11,55 +11,63 @@ 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 {
+    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) {
+                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()
+    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)
     }
-    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

@@ -2,44 +2,62 @@ 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
+object Aws
+{
 
-  def apply(task: Task): Aws = new Aws(task)
+    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)
+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)
+    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()
-  }
+        builder.build()
+    }
 }

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

@@ -8,27 +8,35 @@ import com.amazonaws.client.builder.AwsClientBuilder
 import org.embulk.config.{Config, ConfigDefault}
 import org.embulk.output.s3_parquet.aws.AwsClientConfiguration.Task
 
-object AwsClientConfiguration {
 
-  trait Task {
+object AwsClientConfiguration
+{
 
-    @Config("http_proxy")
-    @ConfigDefault("null")
-    def getHttpProxy: Optional[HttpProxy.Task]
+    trait 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

@@ -9,120 +9,129 @@ import org.embulk.config.{Config, ConfigDefault, ConfigException}
 import org.embulk.output.s3_parquet.aws.AwsCredentials.Task
 import org.embulk.spi.unit.LocalFile
 
-object AwsCredentials {
 
-  trait Task {
+object AwsCredentials
+{
 
-    @Config("auth_method")
-    @ConfigDefault("\"default\"")
-    def getAuthMethod: String
+    trait Task
+    {
 
-    @Config("access_key_id")
-    @ConfigDefault("null")
-    def getAccessKeyId: Optional[String]
+        @Config("auth_method")
+        @ConfigDefault("\"default\"")
+        def getAuthMethod: String
 
-    @Config("secret_access_key")
-    @ConfigDefault("null")
-    def getSecretAccessKey: Optional[String]
+        @Config("access_key_id")
+        @ConfigDefault("null")
+        def getAccessKeyId: Optional[String]
 
-    @Config("session_token")
-    @ConfigDefault("null")
-    def getSessionToken: Optional[String]
+        @Config("secret_access_key")
+        @ConfigDefault("null")
+        def getSecretAccessKey: Optional[String]
 
-    @Config("profile_file")
-    @ConfigDefault("null")
-    def getProfileFile: Optional[LocalFile]
+        @Config("session_token")
+        @ConfigDefault("null")
+        def getSessionToken: Optional[String]
 
-    @Config("profile_name")
-    @ConfigDefault("\"default\"")
-    def getProfileName: String
+        @Config("profile_file")
+        @ConfigDefault("null")
+        def getProfileFile: Optional[LocalFile]
 
-    @Config("role_arn")
-    @ConfigDefault("null")
-    def getRoleArn: Optional[String]
+        @Config("profile_name")
+        @ConfigDefault("\"default\"")
+        def getProfileName: String
 
-    @Config("role_session_name")
-    @ConfigDefault("null")
-    def getRoleSessionName: Optional[String]
+        @Config("role_arn")
+        @ConfigDefault("null")
+        def getRoleArn: Optional[String]
 
-    @Config("role_external_id")
-    @ConfigDefault("null")
-    def getRoleExternalId: Optional[String]
+        @Config("role_session_name")
+        @ConfigDefault("null")
+        def getRoleSessionName: Optional[String]
 
-    @Config("role_session_duration_seconds")
-    @ConfigDefault("null")
-    def getRoleSessionDurationSeconds: Optional[Int]
+        @Config("role_external_id")
+        @ConfigDefault("null")
+        def getRoleExternalId: Optional[String]
 
-    @Config("scope_down_policy")
-    @ConfigDefault("null")
-    def getScopeDownPolicy: Optional[String]
+        @Config("role_session_duration_seconds")
+        @ConfigDefault("null")
+        def getRoleSessionDurationSeconds: Optional[Int]
 
-  }
+        @Config("scope_down_policy")
+        @ConfigDefault("null")
+        def getScopeDownPolicy: Optional[String]
 
-  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
+    def apply(task: Task): AwsCredentials =
+    {
+        new AwsCredentials(task)
+    }
+}
 
-      case "profile" =>
-        if (task.getProfileFile.isPresent) {
-          val pf: ProfilesConfigFile = new ProfilesConfigFile(task.getProfileFile.get().getFile)
-          new ProfileCredentialsProvider(pf, task.getProfileName)
+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 "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'].")
         }
-        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 "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}."))
+    }
 
 
 }