diff --git a/zio-http/src/main/scala/zio/http/Middleware.scala b/zio-http/src/main/scala/zio/http/Middleware.scala index 5f5a9212e7..22a4fa0579 100644 --- a/zio-http/src/main/scala/zio/http/Middleware.scala +++ b/zio-http/src/main/scala/zio/http/Middleware.scala @@ -170,6 +170,68 @@ object Middleware extends HandlerAspects { } } + def logAnnotate(key: => String, value: => String)(implicit trace: Trace): Middleware[Any] = + logAnnotate(LogAnnotation(key, value)) + + def logAnnotate(logAnnotation: => LogAnnotation, logAnnotations: LogAnnotation*)(implicit + trace: Trace, + ): Middleware[Any] = + logAnnotate((logAnnotation +: logAnnotations).toSet) + + def logAnnotate(logAnnotations: => Set[LogAnnotation])(implicit trace: Trace): Middleware[Any] = + new Middleware[Any] { + def apply[Env1 <: Any, Err](routes: Routes[Env1, Err]): Routes[Env1, Err] = + routes.transform[Env1] { h => + handler((req: Request) => ZIO.logAnnotate(logAnnotations)(h(req))) + } + } + + /** + * Creates a middleware that will annotate log messages that are logged while + * a request is handled with log annotations derived from the request. + */ + def logAnnotate(fromRequest: Request => Set[LogAnnotation])(implicit trace: Trace): Middleware[Any] = + new Middleware[Any] { + def apply[Env1 <: Any, Err](routes: Routes[Env1, Err]): Routes[Env1, Err] = + routes.transform[Env1] { h => + handler((req: Request) => ZIO.logAnnotate(fromRequest(req))(h(req))) + } + } + + /** + * Creates a middleware that will annotate log messages that are logged while + * a request is handled with the names and the values of the specified + * headers. + */ + def logAnnotateHeaders(headerName: String, headerNames: String*)(implicit trace: Trace): Middleware[Any] = + new Middleware[Any] { + def apply[Env1 <: Any, Err](routes: Routes[Env1, Err]): Routes[Env1, Err] = { + val headers = headerName +: headerNames + routes.transform[Env1] { h => + handler((req: Request) => { + val annotations = Set.newBuilder[LogAnnotation] + annotations.sizeHint(headers.length) + var i = 0 + while (i < headers.length) { + val name = headers(i) + annotations += LogAnnotation(name, req.headers.get(name).mkString) + i += 1 + } + ZIO.logAnnotate(annotations.result())(h(req)) + }) + } + } + } + + /** + * Creates middleware that will annotate log messages that are logged while a + * request is handled with the names and the values of the specified headers. + */ + def logAnnotateHeaders(header: Header.HeaderType, headers: Header.HeaderType*)(implicit + trace: Trace, + ): Middleware[Any] = + logAnnotateHeaders(header.name, headers.map(_.name): _*) + def timeout(duration: Duration)(implicit trace: Trace): Middleware[Any] = new Middleware[Any] { def apply[Env1 <: Any, Err](routes: Routes[Env1, Err]): Routes[Env1, Err] = diff --git a/zio-http/src/test/scala/zio/http/LogAnnotationMiddlewareSpec.scala b/zio-http/src/test/scala/zio/http/LogAnnotationMiddlewareSpec.scala new file mode 100644 index 0000000000..f9f58b5a1f --- /dev/null +++ b/zio-http/src/test/scala/zio/http/LogAnnotationMiddlewareSpec.scala @@ -0,0 +1,72 @@ +package zio.http + +import zio._ +import zio.test._ + +object LogAnnotationMiddlewareSpec extends ZIOSpecDefault { + override def spec: Spec[TestEnvironment with Scope, Any] = + suite("LogAnnotationMiddlewareSpec")( + test("add static log annotation") { + val response = Routes + .singleton( + handler(ZIO.logWarning("Oh!") *> ZIO.succeed(Response.text("Hey logging!"))), + ) + .@@(Middleware.logAnnotate("label", "value")) + .toHttpApp + .runZIO(Request.get("/")) + + for { + _ <- response + logs <- ZTestLogger.logOutput + log = logs.filter(_.message() == "Oh!").head + } yield assertTrue(log.annotations.get("label").contains("value")) + + }, + test("add request method and path as annotation") { + val response = Routes + .singleton( + handler(ZIO.logWarning("Oh!") *> ZIO.succeed(Response.text("Hey logging!"))), + ) + .@@( + Middleware.logAnnotate(req => + Set(LogAnnotation("method", req.method.name), LogAnnotation("path", req.path.encode)), + ), + ) + .toHttpApp + .runZIO(Request.get("/")) + + for { + _ <- response + logs <- ZTestLogger.logOutput + log = logs.filter(_.message() == "Oh!").head + } yield assertTrue( + log.annotations.get("method").contains("GET"), + log.annotations.get("path").contains("/"), + ) + }, + test("add headers as annotation") { + val response = Routes + .singleton( + handler(ZIO.logWarning("Oh!") *> ZIO.succeed(Response.text("Hey logging!"))), + ) + .@@(Middleware.logAnnotateHeaders("header")) + .@@(Middleware.logAnnotateHeaders(Header.UserAgent.name)) + .toHttpApp + .runZIO { + Request + .get("/") + .addHeader("header", "value") + .addHeader(Header.UserAgent.Product("zio-http", Some("3.0.0"))) + } + + for { + _ <- response + logs <- ZTestLogger.logOutput + log = logs.filter(_.message() == "Oh!").head + } yield assertTrue( + log.annotations.get("header").contains("value"), + log.annotations.get(Header.UserAgent.name).contains("zio-http/3.0.0"), + ) + }, + ) +}