Add `completeWithPagination` directive (#180)

* Extracts pagination from the request

* Adds ResourceCount/PageCount headers to the response

3 files changed, 55 insertions, 8 deletions

diff --git a/src/main/scala/xyz/driver/core/rest/package.scala b/src/main/scala/xyz/driver/core/rest/package.scala
index f85c39a..4858fa7 100644
--- a/src/main/scala/xyz/driver/core/rest/package.scala
+++ b/src/main/scala/xyz/driver/core/rest/package.scala
@@ -3,21 +3,21 @@ package xyz.driver.core.rest
 import java.net.InetAddress
 
 import akka.http.scaladsl.marshalling.{ToEntityMarshaller, ToResponseMarshallable}
-import akka.http.scaladsl.model.headers._
 import akka.http.scaladsl.model._
+import akka.http.scaladsl.model.headers._
 import akka.http.scaladsl.server.Directives._
 import akka.http.scaladsl.server._
 import akka.http.scaladsl.unmarshalling.Unmarshal
 import akka.stream.Materializer
 import akka.stream.scaladsl.Flow
 import akka.util.ByteString
+import scalaz.Scalaz.{intInstance, stringInstance}
+import scalaz.syntax.equal._
+import scalaz.{Functor, OptionT}
 import xyz.driver.tracing.TracingDirectives
 
 import scala.concurrent.Future
 import scala.util.Try
-import scalaz.{Functor, OptionT}
-import scalaz.Scalaz.{intInstance, stringInstance}
-import scalaz.syntax.equal._
 
 trait Service
 
@@ -253,6 +253,21 @@ object `package` {
   def paginationQuery(pagination: Pagination) =
     Seq("pageNumber" -> pagination.pageNumber.toString, "pageSize" -> pagination.pageSize.toString)
 
+  def completeWithPagination[T](handler: Option[Pagination] => Future[ListResponse[T]])(
+      implicit marshaller: ToEntityMarshaller[Seq[T]]): Route = {
+    optionalPagination { pagination =>
+      onSuccess(handler(pagination)) {
+        case ListResponse(resultPart, ListResponse.Meta(count, _, pageSize)) =>
+          val pageCount = (count / pageSize) + (if (count % pageSize == 0) 0 else 1)
+          val headers = List(
+            RawHeader(ContextHeaders.ResourceCount, count.toString),
+            RawHeader(ContextHeaders.PageCount, pageCount.toString))
+
+          respondWithHeaders(headers)(complete(ToResponseMarshallable(resultPart)))
+      }
+    }
+  }
+
   private def extractSorting(sortingString: Option[String]): Sorting = {
     val sortingFields = sortingString.fold(Seq.empty[SortingField])(
       _.split(",")
diff --git a/src/test/scala/xyz/driver/core/rest/DriverRouteTest.scala b/src/test/scala/xyz/driver/core/rest/DriverRouteTest.scala
index d32fefd..86cf8b5 100644
--- a/src/test/scala/xyz/driver/core/rest/DriverRouteTest.scala
+++ b/src/test/scala/xyz/driver/core/rest/DriverRouteTest.scala
@@ -4,16 +4,15 @@ import akka.http.scaladsl.marshallers.sprayjson.SprayJsonSupport
 import akka.http.scaladsl.model.StatusCodes
 import akka.http.scaladsl.model.headers.Connection
 import akka.http.scaladsl.server.Directives.{complete => akkaComplete}
-import akka.http.scaladsl.server.{Directives, Rejection, RejectionHandler, Route}
+import akka.http.scaladsl.server.{Directives, RejectionHandler, Route}
 import akka.http.scaladsl.testkit.ScalatestRouteTest
 import com.typesafe.scalalogging.Logger
 import org.scalatest.{AsyncFlatSpec, Matchers}
-import xyz.driver.core.logging.NoLogger
-import xyz.driver.core.json.serviceExceptionFormat
 import xyz.driver.core.FutureExtensions
+import xyz.driver.core.json.serviceExceptionFormat
+import xyz.driver.core.logging.NoLogger
 import xyz.driver.core.rest.errors._
 
-import scala.collection.immutable
 import scala.concurrent.Future
 
 class DriverRouteTest
diff --git a/src/test/scala/xyz/driver/core/rest/RestTest.scala b/src/test/scala/xyz/driver/core/rest/RestTest.scala
index 68fe419..5403765 100644
--- a/src/test/scala/xyz/driver/core/rest/RestTest.scala
+++ b/src/test/scala/xyz/driver/core/rest/RestTest.scala
@@ -1,5 +1,6 @@
 package xyz.driver.core.rest
 
+import akka.http.scaladsl.marshalling.ToResponseMarshallable
 import akka.http.scaladsl.model.StatusCodes
 import akka.http.scaladsl.server.{Directives, Route, ValidationRejection}
 import akka.http.scaladsl.testkit.ScalatestRouteTest
@@ -7,6 +8,9 @@ import akka.util.ByteString
 import org.scalatest.{Matchers, WordSpec}
 import xyz.driver.core.rest
 
+import scala.concurrent.Future
+import scala.util.Random
+
 class RestTest extends WordSpec with Matchers with ScalatestRouteTest with Directives {
   "`escapeScriptTags` function" should {
     "escape script tags properly" in {
@@ -70,4 +74,33 @@ class RestTest extends WordSpec with Matchers with ScalatestRouteTest with Direc
       }
     }
   }
+
+  "completeWithPagination directive" should {
+    import akka.http.scaladsl.marshallers.sprayjson.SprayJsonSupport._
+    import spray.json.DefaultJsonProtocol._
+
+    val data = Seq.fill(103)(Random.alphanumeric.take(10).mkString)
+    val route: Route = rest.completeWithPagination[String] {
+      case Some(pagination) =>
+        val filtered = data.slice(pagination.offset, pagination.offset + pagination.pageSize)
+        Future.successful(ListResponse(filtered, data.size, Some(pagination)))
+      case None => Future.successful(ListResponse(data, data.size, None))
+    }
+
+    "return a response with pagination headers when pagination has been passed" in {
+      Get("/?pageNumber=2&pageSize=10") ~> route ~> check {
+        responseAs[Seq[String]] shouldBe data.slice(10, 20)
+        header(ContextHeaders.ResourceCount).map(_.value) should contain("103")
+        header(ContextHeaders.PageCount).map(_.value) should contain("11")
+      }
+    }
+
+    "return a response with pagination headers when no pagination has been passed" in {
+      Get("/") ~> route ~> check {
+        responseAs[Seq[String]] shouldBe data
+        header(ContextHeaders.ResourceCount).map(_.value) should contain("103")
+        header(ContextHeaders.PageCount).map(_.value) should contain("1")
+      }
+    }
+  }
 }