marker

PaginatedAWSRequest.scala

package com.gu.emr.util

import java.util.concurrent.{Future => JavaFuture}

import com.amazonaws.services.elasticmapreduce.model._
import com.amazonaws.{AmazonWebServiceRequest, AmazonWebServiceResult}

import scala.concurrent.{ExecutionContext, Future}

trait Marking[Request, Result] {

  def getMarker(result: Result): String

  def withMarker(request: Request, marker: String): Request
}

case class PaginatedAWSRequest[Request <: AmazonWebServiceRequest, Result <: AmazonWebServiceResult[_], Item](
  asyncAwsCall: Request => JavaFuture[Result])(getItems: Result => java.util.List[Item])(
  implicit marking: Marking[Request, Result]) {

  import com.gu.emr.syntax.javaFuture._

  import scala.collection.JavaConverters._

  private def recurse(request: Request, results: List[Item])(implicit ec: ExecutionContext): Future[List[Item]] =
    asyncAwsCall(request).asScala.flatMap { result =>
      val newResults = results ++ getItems(result).asScala.toList

      Option(marking.getMarker(result)) match {
        case None | Some("") =>
          Future.successful(newResults)

        case Some(m) =>
          recurse(marking.withMarker(request, m), newResults)
      }
    }

  def run(request: Request)(implicit ec: ExecutionContext): Future[List[Item]] =
    recurse(request, List.empty)
}

object PaginatedAWSRequest {

  implicit def listClustersMarker: Marking[ListClustersRequest, ListClustersResult] =
    new Marking[ListClustersRequest, ListClustersResult] {
      override def getMarker(result: ListClustersResult): String = result.getMarker
      override def withMarker(request: ListClustersRequest, marker: String): ListClustersRequest = request.withMarker(marker)
    }

  implicit def listBootstrapActionsMarker: Marking[ListBootstrapActionsRequest, ListBootstrapActionsResult] =
    new Marking[ListBootstrapActionsRequest, ListBootstrapActionsResult] {
      override def getMarker(result: ListBootstrapActionsResult): String = result.getMarker
      override def withMarker(request: ListBootstrapActionsRequest, marker: String): ListBootstrapActionsRequest = request.withMarker(marker)
    }

  implicit def listInstanceGroupsMarker: Marking[ListInstanceGroupsRequest, ListInstanceGroupsResult] =
    new Marking[ListInstanceGroupsRequest, ListInstanceGroupsResult] {
      override def getMarker(result: ListInstanceGroupsResult): String = result.getMarker
      override def withMarker(request: ListInstanceGroupsRequest, marker: String): ListInstanceGroupsRequest = request.withMarker(marker)
    }

  implicit def listStepsMarker: Marking[ListStepsRequest, ListStepsResult] =
    new Marking[ListStepsRequest, ListStepsResult] {
      override def getMarker(result: ListStepsResult): String = result.getMarker
      override def withMarker(request: ListStepsRequest, marker: String): ListStepsRequest = request.withMarker(marker)
    }
}