Ld/aws dynamo upgrade 3

common-lib/src/main/scala/com/gu/mediaservice/lib/aws/DynamoDB.scala

@@ -15,8 +15,9 @@ import play.api.libs.json._
 import software.amazon.awssdk.enhanced.dynamodb._
 import software.amazon.awssdk.enhanced.dynamodb.document.EnhancedDocument
 import software.amazon.awssdk.enhanced.dynamodb.model
+import software.amazon.awssdk.enhanced.dynamodb.model.{BatchGetItemEnhancedRequest, ReadBatch}
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient
-import software.amazon.awssdk.services.dynamodb.model.{UpdateItemRequest, AttributeValue => AttributeValueV2, ReturnValue => ReturnValueV2}
+import software.amazon.awssdk.services.dynamodb.model.{UpdateItemRequest, AttributeValue => AttributeValueV2, ReturnValue => ReturnValueV2, QueryRequest => QueryRequestV2}
 
 import scala.annotation.tailrec
 import scala.concurrent.{ExecutionContext, Future}
@@ -131,6 +132,60 @@ class DynamoDB[T](config: CommonConfig, tableName: String, lastModifiedKey: Opti
   def setAddV2(id: String, key: String, value: List[String])(implicit ex: ExecutionContext): Future[JsObject] = Future {
     updateV2(id, DynamoDB.addExpr(key, lastModifiedKey), AttributeValueV2.fromSs(value.asJava))
   }
+
+  def batchGetV2(ids: List[String], attributeKey: String)(implicit ex: ExecutionContext, rjs: Reads[T]): Future[Map[String, T]] = {
+    val chunks =
+      ids.grouped(100).toList.zipWithIndex
+
+    Future
+      .traverse(chunks) { case (chunk, idx) =>
+        logger.info(s"Fetching records for chunk $idx of ${chunks.size}")
+        Future {
+
+          val readBatchBuilder =
+            ReadBatch.builder(classOf[EnhancedDocument])
+              .mappedTableResource(table2)
+
+          chunk.foreach { id =>
+            readBatchBuilder.addGetItem(
+              Key.builder()
+                .partitionValue(id)
+                .build()
+            )
+          }
+
+          val results =
+            dynamo2.batchGetItem(
+              BatchGetItemEnhancedRequest.builder()
+                .readBatches(readBatchBuilder.build())
+                .build()
+            )
+
+          results
+            .resultsForTable(table2)
+            .asScala
+            .toList
+            .flatMap { doc =>
+
+              logger.info(s"Obtained document $doc")
+
+              val json = asJsObject(doc)
+
+              val maybeT =
+                (json \ attributeKey).asOpt[T]
+
+              logger.info(s"Obtained a T of $maybeT from json $json")
+
+              maybeT.map(
+                doc.getString(IdKey) -> _
+              )
+            }
+            .toMap
+        }
+      }
+      .map(_.foldLeft(Map.empty[String, T])(_ ++ _))
+  }
+
   def batchGet(ids: List[String], attributeKey: String)
               (implicit ex: ExecutionContext, rjs: Reads[T]): Future[Map[String, T]] = {
     val keyChunkList = ids
@@ -200,6 +255,36 @@ class DynamoDB[T](config: CommonConfig, tableName: String, lastModifiedKey: Opti
     items map (a => a.getString("id"))
   }
 
+  def scanForIdV2(
+                 indexName: String,
+                 keyName: String,
+                 key: String
+               )(implicit ex: ExecutionContext): Future[List[String]] =
+    Future {
+
+      val response =
+        client2.query(
+          QueryRequestV2.builder()
+            .tableName(tableName)
+            .indexName(indexName)
+            .keyConditionExpression(s"$keyName = :key")
+            .expressionAttributeValues(
+              Map(
+                ":key" ->
+                  AttributeValueV2.builder()
+                    .s(key)
+                    .build()
+              ).asJava
+            )
+            .projectionExpression("id")
+            .build()
+        )
+
+      response.items().asScala.toList.map { item =>
+        item.get("id").s()
+      }
+    }
+
   private def updateRequestBuilder(id: String, expression: String) = {
     UpdateItemRequest.builder()
       .key(Map(IdKey -> AttributeValueV2.fromS(id)).asJava)

metadata-editor/app/lib/MetadataSqsMessageConsumer.scala

@@ -15,6 +15,6 @@ class MetadataSqsMessageConsumer(config: EditsConfig, metadataEditorMetrics: Met
     }
 
   def processDeletedImage(message: JsValue) = Future {
-      withImageId(message)(id => store.deleteItem(id))
+      withImageId(message)(id => store.deleteItemV2(id))
   }
 }

metadata-editor/app/lib/Syndication.scala

@@ -48,8 +48,8 @@ trait Syndication extends Edit with MessageSubjects with GridLogging {
                                 (implicit ec: ExecutionContext): Future[Unit] =
     publishChangedSyndicationRightsForPhotoshoot[Unit](id, unchangedPhotoshoot = false) { () =>
       for {
-        edits <- editsStore.removeKey(id, Edits.Photoshoot)
-        _ <- editsStore.removeKey(id, Edits.PhotoshootTitle)
+        edits <- editsStore.removeKeyV2(id, Edits.Photoshoot)
+        _ <- editsStore.removeKeyV2(id, Edits.PhotoshootTitle)
         _ = publish(id, UpdateImagePhotoshootMetadata)(edits)
       } yield ()
     }
@@ -135,7 +135,7 @@ trait Syndication extends Edit with MessageSubjects with GridLogging {
   def getAllImageRightsInPhotoshoot(photoshoot: Photoshoot)
                                    (implicit ec: ExecutionContext): Future[Map[String, SyndicationRights]] = for {
     imageIds <- getImagesInPhotoshoot(photoshoot)
-    allNonInferredRights <- syndicationStore.batchGet(imageIds, syndicationRightsFieldName)
+    allNonInferredRights <- syndicationStore.batchGetV2(imageIds, syndicationRightsFieldName)
   } yield {
     logger.info(s"Found non-inferred rights for ${allNonInferredRights.size} of ${imageIds.size} images in photoshoot ${photoshoot.title}")
     val mostRecentInferrableRightsMaybe = getMostRecentInferrableSyndicationRights(allNonInferredRights.values.toList)
@@ -144,7 +144,7 @@ trait Syndication extends Edit with MessageSubjects with GridLogging {
 
   def getImagesInPhotoshoot(photoshoot: Photoshoot)
                            (implicit ec: ExecutionContext): Future[List[String]] =
-      editsStore.scanForId(config.editsTablePhotoshootIndex, Edits.PhotoshootTitle, photoshoot.title)
+      editsStore.scanForIdV2(config.editsTablePhotoshootIndex, Edits.PhotoshootTitle, photoshoot.title)
         .recover { case NoItemFound => Nil }
 
   def getChangedRights(before: Map[String, SyndicationRights], after: Map[String, SyndicationRights]): Map[String, Option[SyndicationRights]] = {