-
-
Save acjay/e6ab95743f0a6e97312a52ec95646a9e to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import akka.http.scaladsl.server.Directives.{extractRequestContext, mapRouteResult} | |
import akka.http.scaladsl.server.{Directive0, RequestContext, RouteResult} | |
import akka.http.scaladsl.server.RouteResult.{Complete, Rejected} | |
import akka.stream.scaladsl.Flow | |
import akka.util.ByteString | |
import scala.concurrent.ExecutionContext | |
import scala.util.{Success, Try} | |
object AroundDirectives { | |
/** | |
* @param onRequest Called when a request starts; the result (a function | |
* `Try[RouteResult => Unit]`) is applied to when the | |
* request completes either successfully or fails. | |
*/ | |
def aroundRequest[T](onRequest: RequestContext => Try[RouteResult] => Unit)(implicit ec: ExecutionContext): Directive0 = { | |
extractRequestContext.flatMap { ctx => | |
val onDone = onRequest(ctx) | |
mapRouteResult { | |
case c @ Complete(response) => | |
Complete(response.mapEntity { entity => | |
if (entity.isKnownEmpty()) { | |
// On an empty entity, `transformDataBytes` unsets `isKnownEmpty`. | |
// Call onDone right away, since there's no significant amount of | |
// data to send, anyway. | |
onDone(Success(c)) | |
entity | |
} else { | |
entity.transformDataBytes(Flow[ByteString].watchTermination() { | |
case (m, f) => | |
f.map(_ => c).onComplete(onDone) | |
m | |
}) | |
} | |
}) | |
case r: Rejected => | |
onDone(Success(r)) | |
r | |
} | |
} | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment