Merge pull request #7 from guardian/get-all

Support getAll for getting batches of items
scanamo · Mar 24, 2016 · 7c61382 · 7c61382
2 parents 2da534e + b49a600
commit 7c61382
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 2 deletions.
diff --git a/src/main/scala/com/gu/scanamo/Scanamo.scala b/src/main/scala/com/gu/scanamo/Scanamo.scala
@@ -62,9 +62,37 @@ object Scanamo {
     * Some(Valid(Farmer(Maggot,75,Farm(List(dog)))))
     * }}}
     */
-  def get[K, T](client: AmazonDynamoDB)(tableName: String)(key: (Symbol, K)*)
+  def get[K, T](client: AmazonDynamoDB)(tableName: String)(key: (Symbol, K))
     (implicit fk: DynamoFormat[K], ft: DynamoFormat[T]): Option[ValidatedNel[DynamoReadError, T]] =
-    Option(client.getItem(getRequest(tableName)(key: _*)).getItem).map(read[T])
+    Option(client.getItem(getRequest(tableName)(key)).getItem).map(read[T])
+
+  /**
+    * Returns all the items in the table with matching keys
+    *
+    * Results are returned in the same order as the keys are provided
+    *
+    * {{{
+    * >>> val client = LocalDynamoDB.client()
+    *
+    * >>> case class Farm(animals: List[String])
+    * >>> case class Farmer(name: String, age: Long, farm: Farm)
+    *
+    * >>> val putResult = Scanamo.putAll(client)("farmers")(List(
+    * ...   Farmer("Boggis", 43L, Farm(List("chicken"))), Farmer("Bunce", 52L, Farm(List("goose"))), Farmer("Bean", 55L, Farm(List("turkey")))
+    * ... ))
+    * >>> Scanamo.getAll[String, Farmer](client)("farmers")('name -> List("Boggis", "Bean"))
+    * List(Valid(Farmer(Boggis,43,Farm(List(chicken)))), Valid(Farmer(Bean,55,Farm(List(turkey)))))
+    * }}}
+    */
+  def getAll[K, T](client: AmazonDynamoDB)(tableName: String)(keys: (Symbol, List[K]))
+    (implicit fk: DynamoFormat[K], ft: DynamoFormat[T]): List[ValidatedNel[DynamoReadError, T]] = {
+    import collection.convert.decorateAsScala._
+    val keyValueOptions = keys._2.map(Option(_))
+    def keyValueOption(avMap: java.util.Map[String, AttributeValue]) = fk.read(avMap.get(keys._1.name)).toOption
+
+    client.batchGetItem(batchGetRequest(tableName)(keys)).getResponses.get(tableName).asScala
+      .sortBy(i => keyValueOptions.indexOf(keyValueOption(i))).map(read[T]).toList
+  }
 
   /**
     * {{{

diff --git a/src/main/scala/com/gu/scanamo/ScanamoRequest.scala b/src/main/scala/com/gu/scanamo/ScanamoRequest.scala
@@ -39,6 +39,11 @@ object ScanamoRequest {
   def getRequest[K](tableName: String)(key: (Symbol, K)*)(implicit fk: DynamoFormat[K]): GetItemRequest =
     new GetItemRequest().withTableName(tableName).withKey(asAVMap(key: _*))
 
+  def batchGetRequest[K](tableName: String)(keys: (Symbol, List[K]))(implicit fk: DynamoFormat[K]): BatchGetItemRequest =
+    new BatchGetItemRequest().withRequestItems(Map(tableName ->
+      new KeysAndAttributes().withKeys((keys._2.map(k => Map(keys._1.name -> fk.write(k)).asJava).asJava))
+    ).asJava)
+
   /**
     * {{{
     * prop> import collection.convert.decorateAsJava._