embulk-input-dynamodb 0.2.0 → 0.3.0

data/src/main/scala/org/embulk/input/dynamodb/{ope → deprecated/ope}/ScanOperation.scala

@@ -1,28 +1,35 @@
-package org.embulk.input.dynamodb.ope
+package org.embulk.input.dynamodb.deprecated.ope
 
 import java.util.{List => JList, Map => JMap}
 
-import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient
-import com.amazonaws.services.dynamodbv2.model.{AttributeValue, Condition, ScanRequest, ScanResult}
+import com.amazonaws.services.dynamodbv2.{AmazonDynamoDB, AmazonDynamoDBClient}
+import com.amazonaws.services.dynamodbv2.model.{
+  AttributeValue,
+  Condition,
+  ScanRequest,
+  ScanResult
+}
 import org.embulk.input.dynamodb.PluginTask
 import org.embulk.spi.{BufferAllocator, PageBuilder, PageOutput, Schema}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
+
+class ScanOperation(client: AmazonDynamoDB) extends AbstractOperation {
 
-class ScanOperation(client: AmazonDynamoDBClient) extends AbstractOperation {
   override def execute(
-            task: PluginTask,
-            schema: Schema,
-            output: PageOutput): Unit =
-  {
+      task: PluginTask,
+      schema: Schema,
+      output: PageOutput
+  ): Unit = {
     val allocator: BufferAllocator = task.getBufferAllocator
     val pageBuilder: PageBuilder = new PageBuilder(allocator, schema, output)
 
-    val attributes: JList[String] = schema.getColumns.asScala.map(_.getName).asJava
+    val attributes: JList[String] =
+      schema.getColumns.asScala.map(_.getName).asJava
     val scanFilter: JMap[String, Condition] = createFilters(task).asJava
     var evaluateKey: JMap[String, AttributeValue] = null
 
-    val scanLimit: Long   = task.getScanLimit
+    val scanLimit: Long = task.getScanLimit
     val recordLimit: Long = task.getRecordLimit
     var recordCount: Long = 0
 
@@ -42,9 +49,9 @@ class ScanOperation(client: AmazonDynamoDBClient) extends AbstractOperation {
       val result: ScanResult = client.scan(request)
       evaluateKey = result.getLastEvaluatedKey
 
-      val items = result.getItems.asScala.map(_.asScala.toMap)
+      val items = result.getItems.asScala.map(_.asScala.toMap).toSeq
       recordCount += write(pageBuilder, schema, items)
-    } while(evaluateKey != null && (recordLimit == 0 || recordLimit > recordCount))
+    } while (evaluateKey != null && (recordLimit == 0 || recordLimit > recordCount))
 
     pageBuilder.finish()
   }

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

@@ -0,0 +1,154 @@
+package org.embulk.input.dynamodb.item
+
+import java.nio.ByteBuffer
+import java.nio.charset.StandardCharsets
+import java.util.{Optional, List => JList, Map => JMap}
+
+import com.amazonaws.services.dynamodbv2.model.AttributeValue
+import org.embulk.config.{Config, ConfigDefault, Task => EmbulkTask}
+
+import scala.jdk.CollectionConverters._
+import scala.util.chaining._
+
+/**
+  * TODO: I want to bind directly `org.embulk.config.Config`` to `com.amazonaws.services.dynamodbv2.model.AttributeValue`.
+  * Should I implement `com.amazonaws.transform.JsonUnmarshallerContext`?
+ **/
+object DynamodbAttributeValue {
+
+  trait Task extends EmbulkTask {
+
+    @Config("S")
+    @ConfigDefault("null")
+    def getS: Optional[String]
+
+    @Config("N")
+    @ConfigDefault("null")
+    def getN: Optional[String]
+
+    @Config("B")
+    @ConfigDefault("null")
+    def getB: Optional[String]
+
+    @Config("SS")
+    @ConfigDefault("null")
+    def getSS: Optional[JList[String]]
+
+    @Config("NS")
+    @ConfigDefault("null")
+    def getNS: Optional[JList[String]]
+
+    @Config("BS")
+    @ConfigDefault("null")
+    def getBS: Optional[JList[String]]
+
+    @Config("M")
+    @ConfigDefault("null")
+    def getM: Optional[JMap[String, DynamodbAttributeValue.Task]]
+
+    @Config("L")
+    @ConfigDefault("null")
+    def getL: Optional[JList[DynamodbAttributeValue.Task]]
+
+    @Config("NULL")
+    @ConfigDefault("null")
+    def getNULL: Optional[Boolean]
+
+    @Config("BOOL")
+    @ConfigDefault("null")
+    def getBOOL: Optional[Boolean]
+  }
+
+  def apply(task: Task): DynamodbAttributeValue = {
+    val original = new AttributeValue()
+      .tap(a => task.getS.ifPresent(v => a.setS(v)))
+      .tap(a => task.getN.ifPresent(v => a.setN(v)))
+      .tap { a =>
+        task.getB.ifPresent { v =>
+          a.setB(ByteBuffer.wrap(v.getBytes(StandardCharsets.UTF_8)))
+        }
+      }
+      .tap(a => task.getSS.ifPresent(v => a.setSS(v)))
+      .tap(a => task.getNS.ifPresent(v => a.setNS(v)))
+      .tap { a =>
+        task.getBS.ifPresent { v =>
+          a.setBS(
+            v.asScala
+              .map(e => ByteBuffer.wrap(e.getBytes(StandardCharsets.UTF_8)))
+              .asJava
+          )
+        }
+      }
+      .tap { a =>
+        task.getM.ifPresent { v =>
+          a.setM(v.asScala.map(x => (x._1, apply(x._2).getOriginal)).asJava)
+        }
+      }
+      .tap(a =>
+        task.getL.ifPresent(v =>
+          a.setL(v.asScala.map(apply).map(_.getOriginal).asJava)
+        )
+      )
+      .tap(a => task.getNULL.ifPresent(v => a.setNULL(v)))
+      .tap(a => task.getBOOL.ifPresent(v => a.setBOOL(v)))
+    new DynamodbAttributeValue(original)
+  }
+
+  def apply(original: AttributeValue): DynamodbAttributeValue = {
+    new DynamodbAttributeValue(original)
+  }
+
+  def apply(item: Map[String, AttributeValue]): DynamodbAttributeValue = {
+    val original = new AttributeValue().withM(item.asJava)
+    new DynamodbAttributeValue(original)
+  }
+}
+
+class DynamodbAttributeValue(original: AttributeValue) {
+
+  require(
+    message =
+      s"Invalid AttributeValue: ${original} which must have 1 attribute value.",
+    requirement = {
+      Seq(hasS, hasN, hasB, hasSS, hasNS, hasBS, hasM, hasL, hasNULL, hasBOOL)
+        .count(has => has) == 1
+    }
+  )
+
+  def getOriginal: AttributeValue = original
+  def isNull: Boolean = Option[Boolean](getOriginal.getNULL).getOrElse(false)
+  def hasS: Boolean = Option(getOriginal.getS).isDefined
+  def hasN: Boolean = Option(getOriginal.getN).isDefined
+  def hasB: Boolean = Option(getOriginal.getB).isDefined
+  def hasSS: Boolean = Option(getOriginal.getSS).isDefined
+  def hasNS: Boolean = Option(getOriginal.getNS).isDefined
+  def hasBS: Boolean = Option(getOriginal.getBS).isDefined
+  def hasM: Boolean = Option(getOriginal.getM).isDefined
+  def hasL: Boolean = Option(getOriginal.getL).isDefined
+  def hasNULL: Boolean = Option(getOriginal.getNULL).isDefined
+  def hasBOOL: Boolean = Option(getOriginal.getBOOL).isDefined
+  def getS: String = getOriginal.getS
+  def getN: String = getOriginal.getN
+  def getB: ByteBuffer = getOriginal.getB
+  def getSS: JList[String] = getOriginal.getSS
+  def getNS: JList[String] = getOriginal.getNS
+  def getBS: JList[ByteBuffer] = getOriginal.getBS
+  def getM: JMap[String, AttributeValue] = getOriginal.getM
+  def getL: JList[AttributeValue] = getOriginal.getL
+  def getNULL: Boolean = getOriginal.getNULL
+  def getBOOL: Boolean = getOriginal.getBOOL
+
+  def getType: DynamodbAttributeValueType = {
+    if (hasS) return DynamodbAttributeValueType.S
+    if (hasN) return DynamodbAttributeValueType.N
+    if (hasB) return DynamodbAttributeValueType.B
+    if (hasSS) return DynamodbAttributeValueType.SS
+    if (hasNS) return DynamodbAttributeValueType.NS
+    if (hasBS) return DynamodbAttributeValueType.BS
+    if (hasM) return DynamodbAttributeValueType.M
+    if (hasL) return DynamodbAttributeValueType.L
+    if (hasNULL) return DynamodbAttributeValueType.NULL
+    if (hasBOOL) return DynamodbAttributeValueType.BOOL
+    DynamodbAttributeValueType.UNKNOWN
+  }
+}

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

@@ -0,0 +1,245 @@
+package org.embulk.input.dynamodb.item
+
+import java.nio.ByteBuffer
+import java.nio.charset.StandardCharsets
+
+import org.embulk.input.dynamodb.logger
+import org.embulk.spi.time.{Timestamp, TimestampParser}
+import org.msgpack.value.{Value, ValueFactory}
+
+import scala.jdk.CollectionConverters._
+import scala.util.chaining._
+
+object DynamodbAttributeValueEmbulkTypeTransformable {
+
+  val TRUTHY_STRINGS: Set[String] = Set(
+    "true",
+    "True",
+    "TRUE",
+    "yes",
+    "Yes",
+    "YES",
+    "t",
+    "T",
+    "y",
+    "Y",
+    "on",
+    "On",
+    "ON",
+    "1"
+  )
+
+  val FALSY_STRINGS: Set[String] = Set(
+    "false",
+    "False",
+    "FALSE",
+    "no",
+    "No",
+    "NO",
+    "f",
+    "F",
+    "n",
+    "N",
+    "off",
+    "Off",
+    "OFF",
+    "0"
+  )
+}
+
+case class DynamodbAttributeValueEmbulkTypeTransformable(
+    attributeValue: DynamodbAttributeValue,
+    typeEnforcer: Option[DynamodbAttributeValueType] = None,
+    timestampParser: Option[TimestampParser] = None
+) {
+
+  private def fromAttributeValueType: DynamodbAttributeValueType =
+    typeEnforcer.getOrElse(attributeValue.getType)
+
+  private def convertNAsLongOrDouble(n: String): Either[Long, Double] = {
+    n.toLongOption match {
+      case Some(l) => Left(l)
+      case None    => Right(n.toDouble)
+    }
+  }
+
+  private def convertBAsString(b: ByteBuffer): String = {
+    new String(b.array(), StandardCharsets.UTF_8)
+  }
+
+  private def hasAttributeValueType: Boolean = {
+    fromAttributeValueType.equals(attributeValue.getType)
+  }
+
+  def asMessagePack: Option[Value] = {
+    if (!hasAttributeValueType) return None
+    if (attributeValue.isNull) return None
+
+    Option(fromAttributeValueType match {
+      case DynamodbAttributeValueType.S =>
+        ValueFactory.newString(attributeValue.getS)
+      case DynamodbAttributeValueType.N =>
+        convertNAsLongOrDouble(attributeValue.getN) match {
+          case Left(v)  => ValueFactory.newInteger(v)
+          case Right(v) => ValueFactory.newFloat(v)
+        }
+      case DynamodbAttributeValueType.B =>
+        ValueFactory.newBinary(attributeValue.getB.array())
+      case DynamodbAttributeValueType.SS =>
+        ValueFactory.newArray(
+          attributeValue.getSS.asScala.map(ValueFactory.newString).asJava
+        )
+      case DynamodbAttributeValueType.NS =>
+        ValueFactory.newArray(
+          attributeValue.getNS.asScala
+            .map(convertNAsLongOrDouble(_) match {
+              case Left(v)  => ValueFactory.newInteger(v)
+              case Right(v) => ValueFactory.newFloat(v)
+            })
+            .asJava
+        )
+      case DynamodbAttributeValueType.BS =>
+        ValueFactory.newArray(
+          attributeValue.getBS.asScala
+            .map(b => ValueFactory.newBinary(b.array()))
+            .asJava
+        )
+      case DynamodbAttributeValueType.M =>
+        ValueFactory
+          .newMapBuilder()
+          .tap { builder =>
+            attributeValue.getM.asScala.foreach { x =>
+              builder.put(
+                ValueFactory.newString(x._1),
+                DynamodbAttributeValueEmbulkTypeTransformable(
+                  DynamodbAttributeValue(x._2)
+                ).asMessagePack.getOrElse(ValueFactory.newNil())
+              )
+            }
+          }
+          .build()
+      case DynamodbAttributeValueType.L =>
+        ValueFactory.newArray(
+          attributeValue.getL.asScala.map { av =>
+            DynamodbAttributeValueEmbulkTypeTransformable(
+              DynamodbAttributeValue(av)
+            ).asMessagePack.getOrElse(ValueFactory.newNil())
+          }.asJava
+        )
+      case DynamodbAttributeValueType.BOOL =>
+        ValueFactory.newBoolean(attributeValue.getBOOL)
+      case _ =>
+        logger.warn(
+          s"Unsupported AttributeValue: ${attributeValue.getOriginal.toString}"
+        )
+        return None
+    })
+  }
+
+  def asBoolean: Option[Boolean] = {
+    if (!hasAttributeValueType) return None
+    if (attributeValue.isNull) return None
+
+    Option(fromAttributeValueType match {
+      case DynamodbAttributeValueType.S =>
+        if (DynamodbAttributeValueEmbulkTypeTransformable.TRUTHY_STRINGS
+              .contains(attributeValue.getS)) true
+        else if (DynamodbAttributeValueEmbulkTypeTransformable.FALSY_STRINGS
+                   .contains(attributeValue.getS)) false
+        else return None
+      case DynamodbAttributeValueType.N =>
+        convertNAsLongOrDouble(attributeValue.getN) match {
+          case Left(v)  => v > 0
+          case Right(v) => v > 0.0
+        }
+      case DynamodbAttributeValueType.B =>
+        val s = convertBAsString(attributeValue.getB)
+        if (DynamodbAttributeValueEmbulkTypeTransformable.TRUTHY_STRINGS
+              .contains(s))
+          true
+        else if (DynamodbAttributeValueEmbulkTypeTransformable.FALSY_STRINGS
+                   .contains(s)) false
+        else return None
+      case DynamodbAttributeValueType.BOOL => attributeValue.getBOOL
+      case unsupported =>
+        logger.debug(s"cannot convert ${unsupported.toString} as boolean.")
+        return None
+    })
+  }
+
+  def asLong: Option[Long] = {
+    if (!hasAttributeValueType) return None
+    if (attributeValue.isNull) return None
+
+    Option(fromAttributeValueType match {
+      case DynamodbAttributeValueType.S =>
+        convertNAsLongOrDouble(attributeValue.getS) match {
+          case Left(v)  => v
+          case Right(v) => v.toLong
+        }
+      case DynamodbAttributeValueType.N =>
+        convertNAsLongOrDouble(attributeValue.getN) match {
+          case Left(v)  => v
+          case Right(v) => v.toLong
+        }
+      case DynamodbAttributeValueType.B =>
+        convertNAsLongOrDouble(convertBAsString(attributeValue.getB)) match {
+          case Left(v)  => v
+          case Right(v) => v.toLong
+        }
+      case DynamodbAttributeValueType.BOOL =>
+        if (attributeValue.getBOOL) 1L
+        else 0L
+      case unsupported =>
+        logger.debug(s"cannot convert ${unsupported.toString} as long.")
+        return None
+    })
+  }
+
+  def asDouble: Option[Double] = {
+    if (!hasAttributeValueType) return None
+    if (attributeValue.isNull) return None
+
+    fromAttributeValueType match {
+      case DynamodbAttributeValueType.S => attributeValue.getS.toDoubleOption
+      case DynamodbAttributeValueType.N => attributeValue.getN.toDoubleOption
+      case DynamodbAttributeValueType.B =>
+        convertBAsString(attributeValue.getB).toDoubleOption
+      case DynamodbAttributeValueType.BOOL =>
+        Option(
+          if (attributeValue.getBOOL) 1.0d
+          else 0.0d
+        )
+      case unsupported =>
+        logger.debug(s"cannot convert ${unsupported.toString} as double.")
+        None
+    }
+  }
+
+  def asString: Option[String] = {
+    if (!hasAttributeValueType) return None
+    if (attributeValue.isNull) return None
+
+    Option(fromAttributeValueType match {
+      case DynamodbAttributeValueType.S    => attributeValue.getS
+      case DynamodbAttributeValueType.N    => attributeValue.getN
+      case DynamodbAttributeValueType.B    => convertBAsString(attributeValue.getB)
+      case DynamodbAttributeValueType.SS   => asMessagePack.map(_.toJson).get
+      case DynamodbAttributeValueType.NS   => asMessagePack.map(_.toJson).get
+      case DynamodbAttributeValueType.BS   => asMessagePack.map(_.toJson).get
+      case DynamodbAttributeValueType.M    => asMessagePack.map(_.toJson).get
+      case DynamodbAttributeValueType.L    => asMessagePack.map(_.toJson).get
+      case DynamodbAttributeValueType.BOOL => attributeValue.getBOOL.toString
+      case _ =>
+        logger.warn(
+          s"Unsupported AttributeValue: ${attributeValue.getOriginal.toString}"
+        )
+        return None
+    })
+  }
+
+  def asTimestamp: Option[Timestamp] = {
+    timestampParser.flatMap(p => asString.map(p.parse))
+  }
+
+}