Env access in handle error (#2870)

zio · Sep 9, 2024 · f064483 · f064483
1 parent 23a771b
commit f064483
Show file tree

Hide file tree

Showing 5 changed files with 68 additions and 25 deletions.
diff --git a/docs/reference/aop/handler_aspect.md b/docs/reference/aop/handler_aspect.md
@@ -303,11 +303,11 @@ object UserRepository {
 
 ```scala mdoc:silent
 Routes(
-  Method.GET / "user" / int("userId") -> sessionMiddleware -> handler { 
-    (userId: Int, session: Session, request: Request) =>
-      UserRepository.getUser(session.organizationId, userId)
+  Method.GET / "user" / int("userId") -> handler { 
+    (userId: Int, request: Request) =>
+      withContext((session: Session) => UserRepository.getUser(session.organizationId, userId))
   }
-)
+) @@ sessionMiddleware
 ```
 
 The `HandlerAspect` companion object provides a number of helpful constructors for these middlewares.

diff --git a/zio-http/jvm/src/test/scala/zio/http/internal/middlewares/AuthSpec.scala b/zio-http/jvm/src/test/scala/zio/http/internal/middlewares/AuthSpec.scala
@@ -83,7 +83,7 @@ object AuthSpec extends ZIOHttpSpec with TestExtensions {
         val app = {
           Routes(
             Method.GET / "context" ->
-              handler { (_: Request) => withContext((c: AuthContext) => Response.text(c.value))} @@ basicAuthContextM,
+              handler { (_: Request) => withContext((c: AuthContext) => Response.text(c.value)) } @@ basicAuthContextM,
           )
         }
         assertZIO(

diff --git a/zio-http/shared/src/main/scala/zio/http/Route.scala b/zio-http/shared/src/main/scala/zio/http/Route.scala
@@ -53,7 +53,9 @@ sealed trait Route[-Env, +Err] { self =>
   final def handleError(f: Err => Response)(implicit trace: Trace): Route[Env, Nothing] =
     self.handleErrorCauseZIO(c => ErrorResponseConfig.configRef.get.map(Response.fromCauseWith(c, _)(f)))
 
-  final def handleErrorZIO(f: Err => ZIO[Any, Nothing, Response])(implicit trace: Trace): Route[Env, Nothing] =
+  final def handleErrorZIO[Env1 <: Env](
+    f: Err => ZIO[Env1, Nothing, Response],
+  )(implicit trace: Trace): Route[Env1, Nothing] =
     self.handleErrorCauseZIO { cause =>
       cause.failureOrCause match {
         case Left(err)    => f(err)
@@ -80,7 +82,10 @@ sealed trait Route[-Env, +Err] { self =>
               handler { (request: Request) =>
                 pattern.decode(request.method, request.path) match {
                   case Left(error)   => ZIO.dieMessage(error)
-                  case Right(params) => handler0(zippable.zip(params, request))
+                  case Right(params) =>
+                    handler0.asInstanceOf[Handler[Any, Err, Any, Response]](
+                      zippable.asInstanceOf[Zippable[Any, Any]].zip(params, request),
+                    )
                 }
               }
 
@@ -97,24 +102,37 @@ sealed trait Route[-Env, +Err] { self =>
    * can be used to convert a route that does not handle its errors into one
    * that does handle its errors.
    */
-  final def handleErrorCauseZIO(
-    f: Cause[Err] => ZIO[Any, Nothing, Response],
-  )(implicit trace: Trace): Route[Env, Nothing] =
+  final def handleErrorCauseZIO[Env1 <: Env](
+    f: Cause[Err] => ZIO[Env1, Nothing, Response],
+  )(implicit trace: Trace): Route[Env1, Nothing] =
     self match {
-      case Provided(route, env)                     => Provided(route.handleErrorCauseZIO(f), env)
-      case Augmented(route, aspect)                 => Augmented(route.handleErrorCauseZIO(f), aspect)
+      case Provided(route, env)                     =>
+        Route.handledIgnoreParams(route.routePattern)(
+          Handler.fromZIO(ZIO.environment[Env1]).flatMap { (env1: ZEnvironment[Env1]) =>
+            val env0 = env.asInstanceOf[ZEnvironment[Any]] ++ env1.asInstanceOf[ZEnvironment[Any]]
+            route
+              .handleErrorCauseZIO(f)
+              .toHandler
+              .asInstanceOf[Handler[Any, Response, Request, Response]]
+              .provideEnvironment(env0)
+          },
+        )
+      case Augmented(route, aspect)                 =>
+        Augmented(route.handleErrorCauseZIO(f).asInstanceOf[Route[Any, Nothing]], aspect)
       case Handled(routePattern, handler, location) =>
         Handled(routePattern, handler.map(_.mapErrorCauseZIO(c => f(c.asInstanceOf[Cause[Nothing]]))), location)
 
       case Unhandled(pattern, handler0, zippable, location) =>
-        val handler2: Handler[Any, Nothing, RoutePattern[_], Handler[Env, Response, Request, Response]] = {
+        val handler2: Handler[Any, Nothing, RoutePattern[_], Handler[Env1, Response, Request, Response]] = {
           handler { (pattern: RoutePattern[_]) =>
             val paramHandler =
               handler { (request: Request) =>
                 pattern.decode(request.method, request.path) match {
                   case Left(error)   => ZIO.dieMessage(error)
                   case Right(params) =>
-                    handler0(zippable.zip(params, request))
+                    handler0.asInstanceOf[Handler[Any, Err, Any, Response]](
+                      zippable.asInstanceOf[Zippable[Any, Any]].zip(params, request),
+                    )
                 }
               }
             paramHandler.mapErrorCauseZIO(f)
@@ -190,7 +208,10 @@ sealed trait Route[-Env, +Err] { self =>
               handler { (request: Request) =>
                 pattern.decode(request.method, request.path) match {
                   case Left(error)   => ZIO.dieMessage(error)
-                  case Right(params) => handler0(zippable.zip(params, request))
+                  case Right(params) =>
+                    handler0.asInstanceOf[Handler[Any, Err, Any, Response]](
+                      zippable.asInstanceOf[Zippable[Any, Any]].zip(params, request),
+                    )
                 }
               }
 
@@ -209,12 +230,23 @@ sealed trait Route[-Env, +Err] { self =>
    * convert a route that does not handle its errors into one that does handle
    * its errors.
    */
-  final def handleErrorRequestCauseZIO(
-    f: (Request, Cause[Err]) => ZIO[Any, Nothing, Response],
-  )(implicit trace: Trace): Route[Env, Nothing] =
+  final def handleErrorRequestCauseZIO[Env1 <: Env](
+    f: (Request, Cause[Err]) => ZIO[Env1, Nothing, Response],
+  )(implicit trace: Trace): Route[Env1, Nothing] =
     self match {
-      case Provided(route, env)                     => Provided(route.handleErrorRequestCauseZIO(f), env)
-      case Augmented(route, aspect)                 => Augmented(route.handleErrorRequestCauseZIO(f), aspect)
+      case Provided(route, env)                     =>
+        Route.handledIgnoreParams(route.routePattern)(
+          Handler.fromZIO(ZIO.environment[Env1]).flatMap { (env1: ZEnvironment[Env1]) =>
+            val env0 = env.asInstanceOf[ZEnvironment[Any]] ++ env1.asInstanceOf[ZEnvironment[Any]]
+            route
+              .handleErrorRequestCauseZIO(f)
+              .toHandler
+              .asInstanceOf[Handler[Any, Response, Request, Response]]
+              .provideEnvironment(env0)
+          },
+        )
+      case Augmented(route, aspect)                 =>
+        Augmented(route.handleErrorRequestCauseZIO(f).asInstanceOf[Route[Any, Nothing]], aspect)
       case Handled(routePattern, handler, location) =>
         Handled(
           routePattern,
@@ -227,13 +259,16 @@ sealed trait Route[-Env, +Err] { self =>
         )
 
       case Unhandled(routePattern, handler0, zippable, location) =>
-        val handler2: Handler[Any, Nothing, RoutePattern[_], Handler[Env, Response, Request, Response]] =
+        val handler2: Handler[Any, Nothing, RoutePattern[_], Handler[Env1, Response, Request, Response]] =
           handler { (pattern: RoutePattern[_]) =>
             val paramHandler =
               handler { (request: Request) =>
                 pattern.decode(request.method, request.path) match {
                   case Left(error)   => ZIO.dieMessage(error)
-                  case Right(params) => handler0(zippable.zip(params, request))
+                  case Right(params) =>
+                    handler0.asInstanceOf[Handler[Any, Err, Any, Response]](
+                      zippable.asInstanceOf[Zippable[Any, Any]].zip(params, request),
+                    )
                 }
               }
             Handler.fromFunctionHandler((req: Request) => paramHandler.mapErrorCauseZIO(f(req, _)))

diff --git a/zio-http/shared/src/main/scala/zio/http/Routes.scala b/zio-http/shared/src/main/scala/zio/http/Routes.scala
@@ -82,7 +82,9 @@ final case class Routes[-Env, +Err](routes: Chunk[zio.http.Route[Env, Err]]) { s
   def handleError(f: Err => Response)(implicit trace: Trace): Routes[Env, Nothing] =
     new Routes(routes.map(_.handleError(f)))
 
-  def handleErrorZIO(f: Err => ZIO[Any, Nothing, Response])(implicit trace: Trace): Routes[Env, Nothing] =
+  def handleErrorZIO[Env1 <: Env](f: Err => ZIO[Env1, Nothing, Response])(implicit
+    trace: Trace,
+  ): Routes[Env1, Nothing] =
     new Routes(routes.map(_.handleErrorZIO(f)))
 
   /**

diff --git a/zio-http/shared/src/main/scala/zio/http/codec/QueryCodecs.scala b/zio-http/shared/src/main/scala/zio/http/codec/QueryCodecs.scala
@@ -15,6 +15,8 @@
  */
 
 package zio.http.codec
+import scala.annotation.tailrec
+
 import zio.stacktracer.TracingImplicits.disableAutoTrace
 
 import zio.schema.Schema
@@ -118,10 +120,14 @@ private[codec] trait QueryCodecs {
         )
     }
 
-  private def supportedElementSchema(elementSchema: Schema[Any]) =
-    elementSchema.isInstanceOf[Schema.Primitive[_]] ||
+  @tailrec
+  private def supportedElementSchema(elementSchema: Schema[Any]): Boolean = elementSchema match {
+    case Schema.Lazy(schema0) => supportedElementSchema(schema0())
+    case _                    =>
+      elementSchema.isInstanceOf[Schema.Primitive[_]] ||
       elementSchema.isInstanceOf[Schema.Enum[_]] && elementSchema.annotations.exists(_.isInstanceOf[simpleEnum]) ||
       elementSchema.isInstanceOf[Schema.Record[_]] && elementSchema.asInstanceOf[Schema.Record[_]].fields.size == 1
+  }
 
   def queryAll[A](implicit schema: Schema[A]): QueryCodec[A] =
     schema match {