embulk-input-dynamodb 0.2.0 → 0.3.0

data/src/main/scala/org/embulk/input/dynamodb/item/DynamodbAttributeValueType.scala

@@ -0,0 +1,33 @@
+package org.embulk.input.dynamodb.item
+
+sealed abstract class DynamodbAttributeValueType
+
+object DynamodbAttributeValueType {
+  final case object S extends DynamodbAttributeValueType
+  final case object N extends DynamodbAttributeValueType
+  final case object B extends DynamodbAttributeValueType
+  final case object SS extends DynamodbAttributeValueType
+  final case object NS extends DynamodbAttributeValueType
+  final case object BS extends DynamodbAttributeValueType
+  final case object M extends DynamodbAttributeValueType
+  final case object L extends DynamodbAttributeValueType
+  final case object NULL extends DynamodbAttributeValueType
+  final case object BOOL extends DynamodbAttributeValueType
+  final case object UNKNOWN extends DynamodbAttributeValueType
+
+  def apply(typeName: String): DynamodbAttributeValueType = {
+    typeName match {
+      case "S"    => S
+      case "N"    => N
+      case "B"    => B
+      case "SS"   => SS
+      case "NS"   => NS
+      case "BS"   => BS
+      case "M"    => M
+      case "L"    => L
+      case "NULL" => NULL
+      case "BOOL" => BOOL
+      case _      => UNKNOWN
+    }
+  }
+}

data/src/main/scala/org/embulk/input/dynamodb/item/DynamodbItemColumnVisitor.scala

@@ -0,0 +1,50 @@
+package org.embulk.input.dynamodb.item
+
+import org.embulk.spi.{Column, ColumnVisitor, PageBuilder}
+
+case class DynamodbItemColumnVisitor(
+    itemReader: DynamodbItemReader,
+    pageBuilder: PageBuilder
+) extends ColumnVisitor {
+
+  override def booleanColumn(column: Column): Unit = {
+    itemReader.getBoolean(column) match {
+      case Some(v) => pageBuilder.setBoolean(column, v)
+      case None    => pageBuilder.setNull(column)
+    }
+  }
+
+  override def longColumn(column: Column): Unit = {
+    itemReader.getLong(column) match {
+      case Some(v) => pageBuilder.setLong(column, v)
+      case None    => pageBuilder.setNull(column)
+    }
+  }
+
+  override def doubleColumn(column: Column): Unit =
+    itemReader.getDouble(column) match {
+      case Some(v) => pageBuilder.setDouble(column, v)
+      case None    => pageBuilder.setNull(column)
+    }
+
+  override def stringColumn(column: Column): Unit = {
+    itemReader.getString(column) match {
+      case Some(v) => pageBuilder.setString(column, v)
+      case None    => pageBuilder.setNull(column)
+    }
+  }
+
+  override def timestampColumn(column: Column): Unit = {
+    itemReader.getTimestamp(column) match {
+      case Some(v) => pageBuilder.setTimestamp(column, v)
+      case None    => pageBuilder.setNull(column)
+    }
+  }
+
+  override def jsonColumn(column: Column): Unit =
+    itemReader.getJson(column) match {
+      case Some(v) => pageBuilder.setJson(column, v)
+      case None    => pageBuilder.setNull(column)
+    }
+
+}

data/src/main/scala/org/embulk/input/dynamodb/item/DynamodbItemConsumer.scala

@@ -0,0 +1,40 @@
+package org.embulk.input.dynamodb.item
+
+import com.amazonaws.services.dynamodbv2.model.AttributeValue
+import org.embulk.spi.PageBuilder
+
+object DynamodbItemConsumer {
+
+  def consumeItemsAsJson(
+      schema: DynamodbItemSchema,
+      pageBuilder: PageBuilder
+  ): Seq[Map[String, AttributeValue]] => Unit = {
+    val column = schema.getEmbulkSchema.getColumn(0)
+    items: Seq[Map[String, AttributeValue]] =>
+      items.foreach { item =>
+        val transformable = DynamodbAttributeValueEmbulkTypeTransformable(
+          DynamodbAttributeValue(item)
+        )
+        transformable.asMessagePack match {
+          case Some(v) => pageBuilder.setJson(column, v)
+          case None    => pageBuilder.setNull(column)
+        }
+        pageBuilder.addRecord()
+      }
+  }
+
+  def consumeItemsByEmbulkSchema(
+      schema: DynamodbItemSchema,
+      pageBuilder: PageBuilder
+  ): Seq[Map[String, AttributeValue]] => Unit = {
+    items: Seq[Map[String, AttributeValue]] =>
+      val itemReader = DynamodbItemReader(schema, DynamodbItemIterator(items))
+      val visitor = DynamodbItemColumnVisitor(itemReader, pageBuilder)
+
+      while (itemReader.nextItem) {
+        schema.visitColumns(visitor)
+        pageBuilder.addRecord()
+      }
+  }
+
+}

data/src/main/scala/org/embulk/input/dynamodb/item/DynamodbItemIterator.scala

@@ -0,0 +1,19 @@
+package org.embulk.input.dynamodb.item
+
+import com.amazonaws.services.dynamodbv2.model.AttributeValue
+
+object DynamodbItemIterator {
+
+  def apply(data: Seq[Map[String, AttributeValue]]): DynamodbItemIterator =
+    new DynamodbItemIterator {
+
+      private val delegated: Iterator[Map[String, AttributeValue]] =
+        data.iterator
+      override def hasNext: Boolean = delegated.hasNext
+
+      override def next(): Map[String, DynamodbAttributeValue] =
+        delegated.next().map(x => (x._1, DynamodbAttributeValue(x._2)))
+    }
+}
+
+trait DynamodbItemIterator extends Iterator[Map[String, DynamodbAttributeValue]]

data/src/main/scala/org/embulk/input/dynamodb/item/DynamodbItemReader.scala

@@ -0,0 +1,64 @@
+package org.embulk.input.dynamodb.item
+
+import org.embulk.spi.Column
+import org.embulk.spi.time.Timestamp
+import org.msgpack.value.Value
+
+case class DynamodbItemReader(
+    private val schema: DynamodbItemSchema,
+    private val ite: DynamodbItemIterator
+) {
+  private var currentItem: Map[String, DynamodbAttributeValue] = _
+
+  def nextItem: Boolean = {
+    ite.nextOption() match {
+      case Some(v) =>
+        currentItem = v
+        true
+      case None => false
+    }
+  }
+
+  def getSchema: DynamodbItemSchema = schema
+
+  def getTransformable(
+      name: String,
+      value: DynamodbAttributeValue
+  ): DynamodbAttributeValueEmbulkTypeTransformable = {
+    DynamodbAttributeValueEmbulkTypeTransformable(
+      value,
+      typeEnforcer = schema.getAttributeType(name),
+      timestampParser = schema.getTimestampParser(name)
+    )
+  }
+
+  def getBoolean(column: Column): Option[Boolean] =
+    currentItem
+      .get(column.getName)
+      .flatMap(v => getTransformable(column.getName, v).asBoolean)
+
+  def getString(column: Column): Option[String] =
+    currentItem
+      .get(column.getName)
+      .flatMap(v => getTransformable(column.getName, v).asString)
+
+  def getLong(column: Column): Option[Long] =
+    currentItem
+      .get(column.getName)
+      .flatMap(v => getTransformable(column.getName, v).asLong)
+
+  def getDouble(column: Column): Option[Double] =
+    currentItem
+      .get(column.getName)
+      .flatMap(v => getTransformable(column.getName, v).asDouble)
+
+  def getTimestamp(column: Column): Option[Timestamp] =
+    currentItem
+      .get(column.getName)
+      .flatMap(v => getTransformable(column.getName, v).asTimestamp)
+
+  def getJson(column: Column): Option[Value] =
+    currentItem
+      .get(column.getName)
+      .flatMap(v => getTransformable(column.getName, v).asMessagePack)
+}

data/src/main/scala/org/embulk/input/dynamodb/item/DynamodbItemSchema.scala

@@ -0,0 +1,135 @@
+package org.embulk.input.dynamodb.item
+
+import java.util.{Optional, List => JList}
+
+import com.amazonaws.services.dynamodbv2.model.AttributeValue
+import com.fasterxml.jackson.annotation.{JsonCreator, JsonValue}
+import org.embulk.config.{Config, ConfigDefault, Task => EmbulkTask}
+import org.embulk.spi.{Column, PageBuilder, Schema}
+import org.embulk.spi.`type`.{Type, Types}
+import org.embulk.spi.time.TimestampParser
+
+import scala.jdk.CollectionConverters._
+import scala.util.chaining._
+import scala.util.Try
+
+object DynamodbItemSchema {
+
+  trait ColumnTask
+      extends EmbulkTask
+      with TimestampParser.TimestampColumnOption {
+
+    @Config("name")
+    def getName: String
+
+    @Config("type")
+    def getType: Type
+
+    @Config("attribute_type")
+    @ConfigDefault("null")
+    def getAttributeType: Optional[String]
+  }
+
+  @deprecated(
+    message = "for DeprecatedDynamodbInputPlugin",
+    since = "0.3.0"
+  )
+  case class SchemaConfigCompat(columnTasks: Seq[ColumnTask]) {
+    @JsonCreator
+    def this(columnTasks: JList[ColumnTask]) =
+      this(columnTasks.asScala.toSeq)
+
+    @JsonValue
+    def getColumnTasks: JList[ColumnTask] = columnTasks.asJava
+
+    def toSchema: Schema =
+      Schema
+        .builder()
+        .tap { b =>
+          columnTasks.foreach { t =>
+            b.add(t.getName, t.getType)
+          }
+        }
+        .build()
+
+    def isEmpty: Boolean = columnTasks.isEmpty
+  }
+
+  trait Task extends EmbulkTask with TimestampParser.Task {
+
+    @Config("json_column_name")
+    @ConfigDefault("\"record\"")
+    def getJsonColumnName: String
+
+    @Config("columns")
+    @ConfigDefault("[]")
+    def getColumns: SchemaConfigCompat
+  }
+
+}
+
+case class DynamodbItemSchema(task: DynamodbItemSchema.Task) {
+
+  // TODO: build in this class after removing SchemaConfigCompat.
+  private lazy val embulkSchema: Schema =
+    if (!isItemAsJson) task.getColumns.toSchema
+    else
+      Schema
+        .builder()
+        .add(task.getJsonColumnName, Types.JSON)
+        .build()
+
+  private lazy val timestampParsers: Map[String, TimestampParser] =
+    task.getColumns.columnTasks.map { columnTask =>
+      columnTask.getName -> TimestampParser.of(task, columnTask)
+    }.toMap
+
+  private lazy val attributeTypes: Map[String, DynamodbAttributeValueType] =
+    task.getColumns.columnTasks
+      .filter(_.getAttributeType.isPresent)
+      .map { columnTask =>
+        columnTask.getName -> DynamodbAttributeValueType(
+          columnTask.getAttributeType.get()
+        )
+      }
+      .toMap
+
+  private lazy val embulkColumns: Map[String, Column] =
+    getEmbulkSchema.getColumns.asScala
+      .map(column => column.getName -> column)
+      .toMap
+
+  def getEmbulkSchema: Schema = embulkSchema
+
+  def getTimestampParser(column: Column): Option[TimestampParser] =
+    timestampParsers.get(column.getName)
+
+  def getTimestampParser(columnName: String): Option[TimestampParser] =
+    getEmbulkColumn(columnName).flatMap(getTimestampParser)
+
+  def getAttributeType(column: Column): Option[DynamodbAttributeValueType] =
+    attributeTypes.get(column.getName)
+
+  def getAttributeType(
+      columnName: String
+  ): Option[DynamodbAttributeValueType] =
+    getEmbulkColumn(columnName).flatMap(getAttributeType)
+
+  def getEmbulkColumn(columnName: String): Option[Column] =
+    embulkColumns.get(columnName)
+
+  def getEmbulkColumn(columnIndex: Int): Option[Column] =
+    Try(getEmbulkSchema.getColumn(columnIndex)).toOption
+
+  def isItemAsJson: Boolean = task.getColumns.isEmpty
+
+  def visitColumns(visitor: DynamodbItemColumnVisitor): Unit =
+    getEmbulkSchema.visitColumns(visitor)
+
+  def getItemsConsumer(
+      pageBuilder: PageBuilder
+  ): Seq[Map[String, AttributeValue]] => Unit = {
+    if (isItemAsJson) DynamodbItemConsumer.consumeItemsAsJson(this, pageBuilder)
+    else DynamodbItemConsumer.consumeItemsByEmbulkSchema(this, pageBuilder)
+  }
+}

data/src/main/scala/org/embulk/input/dynamodb/operation/AbstractDynamodbOperation.scala

@@ -0,0 +1,169 @@
+package org.embulk.input.dynamodb.operation
+
+import java.lang.{
+  Boolean => JBoolean,
+  Integer => JInteger,
+  Long => JLong,
+  String => JString
+}
+import java.util.{Optional, Map => JMap}
+
+import com.amazonaws.services.dynamodbv2.model.{
+  AttributeValue,
+  ReturnConsumedCapacity,
+  Select
+}
+import org.embulk.config.{
+  Config,
+  ConfigDefault,
+  ConfigException,
+  Task => EmbulkTask
+}
+import org.embulk.input.dynamodb.item.DynamodbAttributeValue
+
+import scala.jdk.CollectionConverters._
+import scala.language.reflectiveCalls
+
+object AbstractDynamodbOperation {
+
+  trait Task extends EmbulkTask {
+
+    // ref. https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Query.html#Query.ReadConsistency
+    // ref. https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Scan.html#Scan.ReadConsistency
+    @Config("consistent_read")
+    @ConfigDefault("false")
+    def getConsistentRead: Boolean
+
+    // ref. https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Scan.html#Scan.Pagination
+    // ref. https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Query.Pagination.html
+    @Config("exclusive_start_key")
+    @ConfigDefault("{}")
+    def getExclusiveStartKey: JMap[String, DynamodbAttributeValue.Task]
+
+    // ref. https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Expressions.ExpressionAttributeNames.html
+    @Config("expression_attribute_names")
+    @ConfigDefault("{}")
+    def getExpressionAttributeNames: JMap[String, String]
+
+    // ref. https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Expressions.ExpressionAttributeValues.html
+    @Config("expression_attribute_values")
+    @ConfigDefault("{}")
+    def getExpressionAttributeValues: JMap[String, DynamodbAttributeValue.Task]
+
+    // ref. https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Query.html#Query.FilterExpression
+    // ref. https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Scan.html#Scan.FilterExpression
+    @Config("filter_expression")
+    @ConfigDefault("null")
+    def getFilterExpression: Optional[String]
+
+    // ref. https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/SecondaryIndexes.html
+    @Config("index_name")
+    @ConfigDefault("null")
+    def getIndexName: Optional[String]
+
+    // NOTE: Use batch_size for Query/Scan limit per 1 Query/Scan.
+    // ref. https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Query.html#Query.Limit
+    // ref. https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Scan.html#Scan.Limit
+    @Config("batch_size")
+    @ConfigDefault("null")
+    def getBatchSize: Optional[Int]
+
+    // NOTE: This limit is total records limit, not the limit of Query/Scan request.
+    @Config("limit")
+    @ConfigDefault("null")
+    def getLimit: Optional[JLong]
+
+    // ref. https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Expressions.ProjectionExpressions.html
+    @Config("projection_expression")
+    @ConfigDefault("null")
+    def getProjectionExpression: Optional[String]
+
+    // TODO: just reporting ?
+    @Config("return_consumed_capacity")
+    @ConfigDefault("null")
+    def getReturnConsumedCapacity: Optional[ReturnConsumedCapacity]
+
+    // ref. https://docs.aws.amazon.com/amazondynamodb/latest/APIReference/API_Query.html#DDB-Query-request-Select
+    // ref. https://docs.aws.amazon.com/amazondynamodb/latest/APIReference/API_Scan.html#DDB-Scan-request-Select
+    @Config("select")
+    @ConfigDefault("null")
+    def getSelect: Optional[Select]
+
+    def getTableName: String
+    def setTableName(tableName: String): Unit
+  }
+
+  type RequestBuilderMethods = {
+    def setConsistentRead(v: JBoolean): Unit
+    def setExclusiveStartKey(v: JMap[JString, AttributeValue]): Unit
+    def setExpressionAttributeNames(v: JMap[JString, JString]): Unit
+    def setExpressionAttributeValues(v: JMap[JString, AttributeValue]): Unit
+    def setFilterExpression(v: JString): Unit
+    def setIndexName(v: JString): Unit
+    def setLimit(v: JInteger): Unit
+    def setProjectionExpression(v: JString): Unit
+    def setReturnConsumedCapacity(v: ReturnConsumedCapacity): Unit
+    def setSelect(v: Select): Unit
+    def setTableName(v: JString): Unit
+  }
+
+}
+
+abstract class AbstractDynamodbOperation(
+    task: AbstractDynamodbOperation.Task
+) extends EmbulkDynamodbOperation {
+
+  protected def calculateLoadableRecords(loadedRecords: Long): Option[Long] = {
+    if (!task.getLimit.isPresent) return None
+    val loadableRecords = task.getLimit.get() - loadedRecords
+    if (loadableRecords <= 0) Option(0L)
+    else Option(loadableRecords)
+  }
+
+  protected def configureRequest[
+      A <: AbstractDynamodbOperation.RequestBuilderMethods
+  ](
+      req: A,
+      lastEvaluatedKey: Option[Map[String, AttributeValue]]
+  ): Unit = {
+    def attributeValueTaskToAttributeValue(
+        x: (String, DynamodbAttributeValue.Task)
+    ): (String, AttributeValue) = {
+      (x._1, DynamodbAttributeValue(x._2).getOriginal)
+    }
+
+    req.setConsistentRead(task.getConsistentRead)
+    lastEvaluatedKey match {
+      case Some(v) => req.setExclusiveStartKey(v.asJava)
+      case None =>
+        if (!task.getExclusiveStartKey.isEmpty)
+          req.setExclusiveStartKey(
+            task.getExclusiveStartKey.asScala
+              .map(attributeValueTaskToAttributeValue)
+              .asJava
+          )
+    }
+
+    if (!task.getExpressionAttributeNames.isEmpty)
+      req.setExpressionAttributeNames(task.getExpressionAttributeNames)
+    if (!task.getExpressionAttributeValues.isEmpty)
+      req.setExpressionAttributeValues(
+        task.getExpressionAttributeValues.asScala
+          .map(attributeValueTaskToAttributeValue)
+          .asJava
+      )
+    task.getFilterExpression.ifPresent(req.setFilterExpression)
+    task.getIndexName.ifPresent(req.setIndexName)
+    task.getBatchSize.ifPresent { v =>
+      if (v <= 0)
+        throw new ConfigException(
+          "\"batch_size\" must be greater than or equal to 1."
+        )
+      req.setLimit(JInteger.valueOf(v)) // Note: Use BatchSize for the limit per a request.
+    }
+    task.getProjectionExpression.ifPresent(req.setProjectionExpression)
+    task.getReturnConsumedCapacity.ifPresent(req.setReturnConsumedCapacity)
+    task.getSelect.ifPresent(req.setSelect)
+    req.setTableName(task.getTableName)
+  }
+}