embulk-output-s3_parquet 0.0.2

data/gradlew.bat

@@ -0,0 +1,84 @@
+@if "%DEBUG%" == "" @echo off
+@rem ##########################################################################
+@rem
+@rem  Gradle startup script for Windows
+@rem
+@rem ##########################################################################
+
+@rem Set local scope for the variables with windows NT shell
+if "%OS%"=="Windows_NT" setlocal
+
+set DIRNAME=%~dp0
+if "%DIRNAME%" == "" set DIRNAME=.
+set APP_BASE_NAME=%~n0
+set APP_HOME=%DIRNAME%
+
+@rem Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.
+set DEFAULT_JVM_OPTS=
+
+@rem Find java.exe
+if defined JAVA_HOME goto findJavaFromJavaHome
+
+set JAVA_EXE=java.exe
+%JAVA_EXE% -version >NUL 2>&1
+if "%ERRORLEVEL%" == "0" goto init
+
+echo.
+echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH.
+echo.
+echo Please set the JAVA_HOME variable in your environment to match the
+echo location of your Java installation.
+
+goto fail
+
+:findJavaFromJavaHome
+set JAVA_HOME=%JAVA_HOME:"=%
+set JAVA_EXE=%JAVA_HOME%/bin/java.exe
+
+if exist "%JAVA_EXE%" goto init
+
+echo.
+echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME%
+echo.
+echo Please set the JAVA_HOME variable in your environment to match the
+echo location of your Java installation.
+
+goto fail
+
+:init
+@rem Get command-line arguments, handling Windows variants
+
+if not "%OS%" == "Windows_NT" goto win9xME_args
+
+:win9xME_args
+@rem Slurp the command line arguments.
+set CMD_LINE_ARGS=
+set _SKIP=2
+
+:win9xME_args_slurp
+if "x%~1" == "x" goto execute
+
+set CMD_LINE_ARGS=%*
+
+:execute
+@rem Setup the command line
+
+set CLASSPATH=%APP_HOME%\gradle\wrapper\gradle-wrapper.jar
+
+@rem Execute Gradle
+"%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -classpath "%CLASSPATH%" org.gradle.wrapper.GradleWrapperMain %CMD_LINE_ARGS%
+
+:end
+@rem End local scope for the variables with windows NT shell
+if "%ERRORLEVEL%"=="0" goto mainEnd
+
+:fail
+rem Set variable GRADLE_EXIT_CONSOLE if you need the _script_ return code instead of
+rem the _cmd.exe /c_ return code!
+if  not "" == "%GRADLE_EXIT_CONSOLE%" exit 1
+exit /b 1
+
+:mainEnd
+if "%OS%"=="Windows_NT" endlocal
+
+:omega

data/lib/embulk/output/s3_parquet.rb

@@ -0,0 +1,3 @@
+Embulk::JavaPlugin.register_output(
+  "s3_parquet", "org.embulk.output.s3_parquet.S3ParquetOutputPlugin",
+  File.expand_path('../../../../classpath', __FILE__))

data/settings.gradle

@@ -0,0 +1 @@
+rootProject.name = 'embulk-output-s3_parquet'

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

@@ -0,0 +1,199 @@
+package org.embulk.output.s3_parquet
+
+
+import java.nio.file.{Files, Paths}
+import java.util.{IllegalFormatException, Locale, Optional, List => JList, Map => JMap}
+
+import com.amazonaws.services.s3.model.CannedAccessControlList
+import org.apache.parquet.column.ParquetProperties
+import org.apache.parquet.hadoop.ParquetWriter
+import org.apache.parquet.hadoop.metadata.CompressionCodecName
+import org.embulk.config.{Config, ConfigDefault, ConfigDiff, ConfigException, ConfigSource, Task, TaskReport, TaskSource}
+import org.embulk.output.s3_parquet.S3ParquetOutputPlugin.PluginTask
+import org.embulk.output.s3_parquet.aws.Aws
+import org.embulk.output.s3_parquet.parquet.ParquetFileWriter
+import org.embulk.spi.{Exec, OutputPlugin, PageReader, Schema, TransactionalPageOutput}
+import org.embulk.spi.time.TimestampFormatter
+import org.embulk.spi.time.TimestampFormatter.TimestampColumnOption
+import org.embulk.spi.util.Timestamps
+import org.slf4j.Logger
+
+object S3ParquetOutputPlugin {
+
+  trait PluginTask
+    extends Task
+    with TimestampFormatter.Task
+    with Aws.Task {
+
+    @Config("bucket")
+    def getBucket: String
+
+    @Config("path_prefix")
+    @ConfigDefault("\"\"")
+    def getPathPrefix: String
+
+    @Config("sequence_format")
+    @ConfigDefault("\"%03d.%02d.\"")
+    def getSequenceFormat: String
+
+    @Config("file_ext")
+    @ConfigDefault("\"parquet\"")
+    def getFileExt: String
+
+    @Config("compression_codec")
+    @ConfigDefault("\"uncompressed\"")
+    def getCompressionCodecString: String
+
+    def setCompressionCodec(v: CompressionCodecName): Unit
+
+    def getCompressionCodec: CompressionCodecName
+
+    @Config("column_options")
+    @ConfigDefault("{}")
+    def getColumnOptions: JMap[String, TimestampColumnOption]
+
+    @Config("canned_acl")
+    @ConfigDefault("\"private\"")
+    def getCannedAclString: String
+
+    def setCannedAcl(v: CannedAccessControlList): Unit
+
+    def getCannedAcl: CannedAccessControlList
+
+    @Config("block_size")
+    @ConfigDefault("null")
+    def getBlockSize: Optional[Int]
+
+    @Config("page_size")
+    @ConfigDefault("null")
+    def getPageSize: Optional[Int]
+
+    @Config("max_padding_size")
+    @ConfigDefault("null")
+    def getMaxPaddingSize: Optional[Int]
+
+    @Config("enable_dictionary_encoding")
+    @ConfigDefault("null")
+    def getEnableDictionaryEncoding: Optional[Boolean]
+
+    @Config("buffer_dir")
+    @ConfigDefault("null")
+    def getBufferDir: Optional[String]
+
+  }
+
+}
+
+class S3ParquetOutputPlugin
+  extends OutputPlugin {
+
+  val logger: Logger = Exec.getLogger(classOf[S3ParquetOutputPlugin])
+
+  private def withPluginContextClassLoader[A](f: => A): A = {
+    val original: ClassLoader = Thread.currentThread.getContextClassLoader
+    Thread.currentThread.setContextClassLoader(classOf[S3ParquetOutputPlugin].getClassLoader)
+    try f
+    finally Thread.currentThread.setContextClassLoader(original)
+  }
+
+  override def transaction(config: ConfigSource,
+                           schema: Schema,
+                           taskCount: Int,
+                           control: OutputPlugin.Control): ConfigDiff = {
+    val task: PluginTask = config.loadConfig(classOf[PluginTask])
+
+    withPluginContextClassLoader {
+      configure(task, schema)
+      control.run(task.dump)
+    }
+
+    Exec.newConfigDiff
+  }
+
+  private def configure(task: PluginTask,
+                        schema: Schema): Unit = {
+    // sequence_format
+    try String.format(task.getSequenceFormat, 0: Integer, 0: Integer)
+    catch {
+      case e: IllegalFormatException => throw new ConfigException(s"Invalid sequence_format: ${task.getSequenceFormat}", e)
+    }
+
+    // compression_codec
+    CompressionCodecName.values().find(v => v.name().toLowerCase(Locale.ENGLISH).equals(task.getCompressionCodecString)) match {
+      case Some(v) => task.setCompressionCodec(v)
+      case None    =>
+        val unsupported: String = task.getCompressionCodecString
+        val supported: String = CompressionCodecName.values().map(v => s"'${v.name().toLowerCase}'").mkString(", ")
+        throw new ConfigException(s"'$unsupported' is unsupported: `compression_codec` must be one of [$supported].")
+    }
+
+    // column_options
+    task.getColumnOptions.forEach { (k: String,
+                                     _) =>
+      val c = schema.lookupColumn(k)
+      if (!c.getType.getName.equals("timestamp")) throw new ConfigException(s"column:$k is not 'timestamp' type.")
+    }
+
+    // canned_acl
+    CannedAccessControlList.values().find(v => v.toString.equals(task.getCannedAclString)) match {
+      case Some(v) => task.setCannedAcl(v)
+      case None    =>
+        val unsupported: String = task.getCannedAclString
+        val supported: String = CannedAccessControlList.values().map(v => s"'${v.toString}'").mkString(", ")
+        throw new ConfigException(s"'$unsupported' is unsupported: `canned_acl` must be one of [$supported].")
+    }
+  }
+
+  override def resume(taskSource: TaskSource,
+                      schema: Schema,
+                      taskCount: Int,
+                      control: OutputPlugin.Control): ConfigDiff = {
+    throw new UnsupportedOperationException("s3_parquet output plugin does not support resuming")
+  }
+
+  override def cleanup(taskSource: TaskSource,
+                       schema: Schema,
+                       taskCount: Int,
+                       successTaskReports: JList[TaskReport]): Unit = {
+    successTaskReports.forEach { tr =>
+      logger.info(
+        s"Created: s3://${tr.get(classOf[String], "bucket")}/${tr.get(classOf[String], "key")}, "
+          + s"version_id: ${tr.get(classOf[String], "version_id", null)}, "
+          + s"etag: ${tr.get(classOf[String], "etag", null)}")
+    }
+  }
+
+  override def open(taskSource: TaskSource,
+                    schema: Schema,
+                    taskIndex: Int): TransactionalPageOutput = {
+    val task = taskSource.loadTask(classOf[PluginTask])
+    val bufferDir: String = task.getBufferDir.orElse(Files.createTempDirectory("embulk-output-s3_parquet-").toString)
+    val bufferFile: String = Paths.get(bufferDir, s"embulk-output-s3_parquet-task-$taskIndex-0.parquet").toString
+    val destS3bucket: String = task.getBucket
+    val destS3Key: String = task.getPathPrefix + String.format(task.getSequenceFormat, taskIndex: Integer, 0: Integer) + task.getFileExt
+
+
+    val pageReader: PageReader = new PageReader(schema)
+    val aws: Aws = Aws(task)
+    val timestampFormatters: Seq[TimestampFormatter] = Timestamps.newTimestampColumnFormatters(task, schema, task.getColumnOptions)
+    val parquetWriter: ParquetWriter[PageReader] = ParquetFileWriter.builder()
+      .withPath(bufferFile)
+      .withSchema(schema)
+      .withTimestampFormatters(timestampFormatters)
+      .withCompressionCodec(task.getCompressionCodec)
+      .withDictionaryEncoding(task.getEnableDictionaryEncoding.orElse(ParquetProperties.DEFAULT_IS_DICTIONARY_ENABLED))
+      .withDictionaryPageSize(task.getPageSize.orElse(ParquetProperties.DEFAULT_DICTIONARY_PAGE_SIZE))
+      .withMaxPaddingSize(task.getMaxPaddingSize.orElse(ParquetWriter.MAX_PADDING_SIZE_DEFAULT))
+      .withPageSize(task.getPageSize.orElse(ParquetProperties.DEFAULT_PAGE_SIZE))
+      .withRowGroupSize(task.getBlockSize.orElse(ParquetWriter.DEFAULT_BLOCK_SIZE))
+      .withValidation(ParquetWriter.DEFAULT_IS_VALIDATING_ENABLED)
+      .withWriteMode(org.apache.parquet.hadoop.ParquetFileWriter.Mode.CREATE)
+      .withWriterVersion(ParquetProperties.DEFAULT_WRITER_VERSION)
+      .build()
+
+    logger.info(s"Local Buffer File: $bufferFile, Destination: s3://$destS3bucket/$destS3Key")
+
+    S3ParquetPageOutput(bufferFile, pageReader, parquetWriter, aws, destS3bucket, destS3Key)
+  }
+
+}

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

@@ -0,0 +1,65 @@
+package org.embulk.output.s3_parquet
+
+
+import java.io.File
+import java.nio.file.{Files, Paths}
+
+import com.amazonaws.services.s3.transfer.{TransferManager, Upload}
+import com.amazonaws.services.s3.transfer.model.UploadResult
+import org.apache.parquet.hadoop.ParquetWriter
+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 {
+
+  private var isClosed: Boolean = false
+
+  override def add(page: Page): Unit = {
+    reader.setPage(page)
+    while (reader.nextRecord()) {
+      writer.write(reader)
+    }
+  }
+
+  override def finish(): Unit = {
+  }
+
+  override def close(): Unit = {
+    synchronized {
+      if (!isClosed) {
+        writer.close()
+        isClosed = true
+      }
+    }
+  }
+
+  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)
+  }
+
+  private def cleanup(): Unit = {
+    Files.delete(Paths.get(outputLocalFile))
+  }
+}

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

@@ -0,0 +1,45 @@
+package org.embulk.output.s3_parquet.aws
+
+
+import com.amazonaws.client.builder.AwsClientBuilder
+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 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

@@ -0,0 +1,34 @@
+package org.embulk.output.s3_parquet.aws
+
+
+import java.util.Optional
+
+import com.amazonaws.ClientConfiguration
+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 {
+
+    @Config("http_proxy")
+    @ConfigDefault("null")
+    def getHttpProxy: Optional[HttpProxy.Task]
+
+  }
+
+  def apply(task: Task): AwsClientConfiguration = new AwsClientConfiguration(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)
+    }
+  }
+
+}

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

@@ -0,0 +1,128 @@
+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}
+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 {
+
+  trait Task {
+
+    @Config("auth_method")
+    @ConfigDefault("\"default\"")
+    def getAuthMethod: String
+
+    @Config("access_key_id")
+    @ConfigDefault("null")
+    def getAccessKeyId: Optional[String]
+
+    @Config("secret_access_key")
+    @ConfigDefault("null")
+    def getSecretAccessKey: Optional[String]
+
+    @Config("session_token")
+    @ConfigDefault("null")
+    def getSessionToken: Optional[String]
+
+    @Config("profile_file")
+    @ConfigDefault("null")
+    def getProfileFile: Optional[LocalFile]
+
+    @Config("profile_name")
+    @ConfigDefault("\"default\"")
+    def getProfileName: String
+
+    @Config("role_arn")
+    @ConfigDefault("null")
+    def getRoleArn: Optional[String]
+
+    @Config("role_session_name")
+    @ConfigDefault("null")
+    def getRoleSessionName: Optional[String]
+
+    @Config("role_external_id")
+    @ConfigDefault("null")
+    def getRoleExternalId: 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
+
+      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'].")
+    }
+  }
+
+  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}."))
+  }
+
+
+}