provide request in context transforms

.gitignore

@@ -13,3 +13,5 @@ examples/project/metals.sbt
 website/node_modules
 website/build
 website/i18n/en.json
+
+.DS_Store

build.sbt

@@ -13,6 +13,8 @@ ThisBuild / resolvers ++= Resolver.sonatypeOssRepos("snapshots")
 
 ThisBuild / versionScheme := Some("early-semver")
 
+ThisBuild / version := "0.6.1"
+
 publish / skip := true
 
 sonatypeProfileName := "com.thesamet"

code-gen/src/main/scala/scalapb/zio_grpc/ZioCodeGenerator.scala

@@ -233,9 +233,9 @@ class ZioFilePrinter(
       val delegate = s"self.${method.name}"
       val newImpl  = method.streamType match {
         case StreamType.Unary | StreamType.ClientStreaming         =>
-          s"f.effect($delegate(request, _))(context)"
+          s"f.effect($delegate)(request, context)"
         case StreamType.ServerStreaming | StreamType.Bidirectional =>
-          s"f.stream($delegate(request, _))(context)"
+          s"f.stream($delegate)(request, context)"
       }
       fp.add(
         methodSignature(
@@ -630,7 +630,7 @@ class ZioFilePrinter(
       ).indent
         .add(s"$makeMetadata.flatMap { metadata =>")
         .indented(
-          _.add(s"transforms.$transformMethod { context => ")
+          _.add(s"transforms.$transformMethod { (_: Any, context) => ")
             .indented(
               _.add(
                 s"$clientCall("
@@ -644,7 +644,7 @@ class ZioFilePrinter(
                 )
                 .add(")")
             )
-            .add(s"}($ClientCallContext(${method.grpcDescriptor.fullName}, $CallOptions.DEFAULT, metadata))")
+            .add(s"}(request, $ClientCallContext(${method.grpcDescriptor.fullName}, $CallOptions.DEFAULT, metadata))")
         )
         .add("}")
         .outdent

core/src/main/scala/scalapb/zio_grpc/transforms.scala

@@ -11,31 +11,31 @@ import io.grpc.StatusException
   */
 trait Transform {
   self =>
-  def effect[A](io: ZIO[Any, StatusException, A]): ZIO[Any, StatusException, A]
-  def stream[A](io: ZStream[Any, StatusException, A]): ZStream[Any, StatusException, A]
+  def effect[A, B](io: A => ZIO[Any, StatusException, B]): A => ZIO[Any, StatusException, B]
+  def stream[A, B](io: A => ZStream[Any, StatusException, B]): A => ZStream[Any, StatusException, B]
 
   // Converts this Transform to GTransform that transforms the effects like this, but
   // leaves the Context unchanged.
   def toGTransform[Context]: GTransform[Context, StatusException, Context, StatusException] =
     new GTransform[Context, StatusException, Context, StatusException] {
-      def effect[A](
-          io: Context => ZIO[Any, StatusException, A]
-      ): Context => ZIO[Any, StatusException, A] = { c =>
-        self.effect(io(c))
+      def effect[A, B](
+          io: (A, Context) => ZIO[Any, StatusException, B]
+      ): (A, Context) => ZIO[Any, StatusException, B] = { (a, c) =>
+        self.effect[A, B](io(_, c))(a)
       }
 
-      def stream[A](
-          io: Context => ZStream[Any, StatusException, A]
-      ): Context => ZStream[Any, StatusException, A] = { c =>
-        self.stream(io(c))
+      def stream[A, B](
+          io: (A, Context) => ZStream[Any, StatusException, B]
+      ): (A, Context) => ZStream[Any, StatusException, B] = { (a, c) =>
+        self.stream(io(_, c))(a)
       }
     }
 
   def andThen(other: Transform): Transform = new Transform {
-    def effect[A](io: ZIO[Any, StatusException, A]): ZIO[Any, StatusException, A] =
+    def effect[A, B](io: A => ZIO[Any, StatusException, B]): A => ZIO[Any, StatusException, B] =
       other.effect(self.effect(io))
 
-    def stream[A](io: ZStream[Any, StatusException, A]): ZStream[Any, StatusException, A] =
+    def stream[A, B](io: A => ZStream[Any, StatusException, B]): A => ZStream[Any, StatusException, B] =
       other.stream(self.stream(io))
   }
 
@@ -44,10 +44,13 @@ trait Transform {
 
 object Transform {
   def fromGTransform(ct: GTransform[Any, StatusException, Any, StatusException]) = new Transform {
-    def effect[A](io: ZIO[Any, StatusException, A]): ZIO[Any, StatusException, A] = ct.effect(_ => io)(())
+    def effect[A, B](io: A => ZIO[Any, StatusException, B]): A => ZIO[Any, StatusException, B] = { a =>
+      ct.effect((a, _) => io(a))(a, ())
+    }
 
-    def stream[A](io: ZStream[Any, StatusException, A]): ZStream[Any, StatusException, A] =
-      ct.stream(_ => io)(())
+    def stream[A, B](io: A => ZStream[Any, StatusException, B]): A => ZStream[Any, StatusException, B] = { a =>
+      ct.stream((a, _) => io(a))(a, ())
+    }
   }
 }
 
@@ -58,25 +61,25 @@ object Transform {
   */
 trait GTransform[+ContextIn, -ErrorIn, -ContextOut, +ErrorOut] {
   self =>
-  def effect[A](
-      io: ContextIn => ZIO[Any, ErrorIn, A]
-  ): (ContextOut => ZIO[Any, ErrorOut, A])
-  def stream[A](
-      io: ContextIn => ZStream[Any, ErrorIn, A]
-  ): (ContextOut => ZStream[Any, ErrorOut, A])
+  def effect[A, B](
+      io: (A, ContextIn) => ZIO[Any, ErrorIn, B]
+  ): ((A, ContextOut) => ZIO[Any, ErrorOut, B])
+  def stream[A, B](
+      io: (A, ContextIn) => ZStream[Any, ErrorIn, B]
+  ): ((A, ContextOut) => ZStream[Any, ErrorOut, B])
 
   def andThen[ContextIn2 <: ContextOut, ErrorIn2 >: ErrorOut, ContextOut2, ErrorOut2](
       other: GTransform[ContextIn2, ErrorIn2, ContextOut2, ErrorOut2]
   ): GTransform[ContextIn, ErrorIn, ContextOut2, ErrorOut2] =
     new GTransform[ContextIn, ErrorIn, ContextOut2, ErrorOut2] {
-      def effect[A](
-          io: ContextIn => ZIO[Any, ErrorIn, A]
-      ): ContextOut2 => ZIO[Any, ErrorOut2, A] =
+      def effect[A, B](
+          io: (A, ContextIn) => ZIO[Any, ErrorIn, B]
+      ): (A, ContextOut2) => ZIO[Any, ErrorOut2, B] =
         other.effect(self.effect(io))
 
-      def stream[A](
-          io: ContextIn => ZStream[Any, ErrorIn, A]
-      ): ContextOut2 => ZStream[Any, ErrorOut2, A] =
+      def stream[A, B](
+          io: (A, ContextIn) => ZStream[Any, ErrorIn, B]
+      ): (A, ContextOut2) => ZStream[Any, ErrorOut2, B] =
         other.stream(self.stream(io))
     }
 
@@ -88,45 +91,47 @@ trait GTransform[+ContextIn, -ErrorIn, -ContextOut, +ErrorOut] {
 object GTransform {
 
   def identity[C, E]: GTransform[C, E, C, E] = new GTransform[C, E, C, E] {
-    def effect[A](io: C => ZIO[Any, E, A]): C => ZIO[Any, E, A]         = io
-    def stream[A](io: C => ZStream[Any, E, A]): C => ZStream[Any, E, A] = io
+    def effect[A, B](io: (A, C) => ZIO[Any, E, B]): (A, C) => ZIO[Any, E, B]         = io
+    def stream[A, B](io: (A, C) => ZStream[Any, E, B]): (A, C) => ZStream[Any, E, B] = io
   }
 
   // Returns a GTransform that effectfully transforms the context parameter
   def apply[ContextIn, Error, ContextOut](
       f: ContextOut => ZIO[Any, Error, ContextIn]
   ): GTransform[ContextIn, Error, ContextOut, Error] =
     new GTransform[ContextIn, Error, ContextOut, Error] {
-      def effect[A](
-          io: ContextIn => ZIO[Any, Error, A]
-      ): ContextOut => ZIO[Any, Error, A] = { (context: ContextOut) =>
-        f(context).flatMap(io)
+      def effect[A, B](
+          io: (A, ContextIn) => ZIO[Any, Error, B]
+      ): (A, ContextOut) => ZIO[Any, Error, B] = { (a: A, context: ContextOut) =>
+        f(context).flatMap(io(a, _))
       }
 
-      def stream[A](
-          io: ContextIn => ZStream[Any, Error, A]
-      ): ContextOut => ZStream[Any, Error, A] = { (context: ContextOut) =>
-        ZStream.fromZIO(f(context)).flatMap(io)
+      def stream[A, B](
+          io: (A, ContextIn) => ZStream[Any, Error, B]
+      ): (A, ContextOut) => ZStream[Any, Error, B] = { (a: A, context: ContextOut) =>
+        ZStream.fromZIO(f(context)).flatMap(io(a, _))
       }
     }
 
   // Returns a GTransform that maps the error parameter.
   def mapError[C, E1, E2](f: E1 => E2): GTransform[C, E1, C, E2] = new GTransform[C, E1, C, E2] {
-    def effect[A](io: C => zio.ZIO[Any, E1, A]): C => zio.ZIO[Any, E2, A]                       = { (context: C) =>
-      io(context).mapError(f)
+    def effect[A, B](io: (A, C) => zio.ZIO[Any, E1, B]): (A, C) => zio.ZIO[Any, E2, B]                       = { (a: A, context: C) =>
+      io(a, context).mapError(f)
     }
-    def stream[A](io: C => zio.stream.ZStream[Any, E1, A]): C => zio.stream.ZStream[Any, E2, A] = { (context: C) =>
-      io(context).mapError(f)
+    def stream[A, B](io: (A, C) => zio.stream.ZStream[Any, E1, B]): (A, C) => zio.stream.ZStream[Any, E2, B] = {
+      (a: A, context: C) =>
+        io(a, context).mapError(f)
     }
   }
 
   // Returns a GTransform that effectfully maps the error parameter.
   def mapErrorZIO[C, E1, E2](f: E1 => zio.UIO[E2]): GTransform[C, E1, C, E2] = new GTransform[C, E1, C, E2] {
-    def effect[A](io: C => zio.ZIO[Any, E1, A]): C => zio.ZIO[Any, E2, A]                       = { (context: C) =>
-      io(context).flatMapError(f)
+    def effect[A, B](io: (A, C) => zio.ZIO[Any, E1, B]): (A, C) => zio.ZIO[Any, E2, B]                       = { (a: A, context: C) =>
+      io(a, context).flatMapError(f)
     }
-    def stream[A](io: C => zio.stream.ZStream[Any, E1, A]): C => zio.stream.ZStream[Any, E2, A] = { (context: C) =>
-      io(context).catchAll(e => ZStream.fromZIO(f(e).flatMap(ZIO.fail(_))))
+    def stream[A, B](io: (A, C) => zio.stream.ZStream[Any, E1, B]): (A, C) => zio.stream.ZStream[Any, E2, B] = {
+      (a: A, context: C) =>
+        io(a, context).catchAll(e => ZStream.fromZIO(f(e).flatMap(ZIO.fail(_))))
     }
   }
 }

e2e/src/test/scalajvm/scalapb/zio_grpc/ClientTransformSpec.scala

@@ -19,7 +19,10 @@ object ClientTransformSpec extends ZIOSpecDefault {
     for {
       metadata <- SafeMetadata.make
       context  <-
-        transform.effect(ZIO.succeed(_))(ClientCallContext(TestServiceGrpc.METHOD_UNARY, CallOptions.DEFAULT, metadata))
+        transform.effect((_: Any, c) => ZIO.succeed(c))(
+          (),
+          ClientCallContext(TestServiceGrpc.METHOD_UNARY, CallOptions.DEFAULT, metadata)
+        )
     } yield context
 
   def spec = suite("ClientTransformSpec")(

examples/fullapp/build.sbt

@@ -6,7 +6,7 @@ ThisBuild / cancelable := true
 
 ThisBuild / connectInput := true
 
-val grpcVersion = "1.50.1"
+val grpcVersion = "1.59.0"
 
 lazy val protos = crossProject(JSPlatform, JVMPlatform)
   .in(file("protos"))

examples/fullapp/project/plugins.sbt

@@ -4,7 +4,7 @@ addSbtPlugin("org.scalameta" % "sbt-scalafmt" % "2.4.3")
 
 addSbtPlugin("com.thesamet" % "sbt-protoc" % "1.0.6")
 
-val zioGrpcVersion = "0.6.0-rc6"
+val zioGrpcVersion = "0.6.1"
 
 libraryDependencies ++= Seq(
   "com.thesamet.scalapb.zio-grpc" %% "zio-grpc-codegen" % zioGrpcVersion,

examples/fullapp/server/src/main/scala/ExampleServer.scala

@@ -75,6 +75,6 @@ object GreeterService extends ZGreeter[RequestContext] {
 
 object ExampleServer extends ServerMain {
   def services = ServiceList.add(
-    GreeterService.transformContextZIO(RequestContext.fromMetadata(_))
+    GreeterService.transformContextZIO(RequestContext.fromMetadata).transform(LoggingTransform)
   )
 }

examples/fullapp/server/src/main/scala/ExampleServerWithMetadata.scala

@@ -92,7 +92,7 @@ object GreeterServiceWithMetadata {
   val layer
       : ZLayer[UserRepo with GreetingsRepo, Nothing, ZGreeter[RequestContext]] =
     ZLayer.fromFunction((userRepo: UserRepo, greetingsRepo: GreetingsRepo) =>
-      GreeterImpl(greetingsRepo).transformContextZIO(findUser(userRepo, _))
+      GreeterImpl(greetingsRepo).transformContextZIO(findUser(userRepo, _)).transform(LoggingTransform)
     )
 }
 

examples/fullapp/server/src/main/scala/LoggingTransform.scala

@@ -0,0 +1,30 @@
+package examples
+
+import io.grpc.StatusException
+import scalapb.zio_grpc.Transform
+import zio._
+import zio.stream._
+
+object LoggingTransform extends Transform {
+  override def effect[A, B](io: A => IO[StatusException, B]): A => IO[StatusException, B] = { a =>
+    for {
+      _ <- ZIO.log(s"Received request: $a")
+      b <- io(a)
+      _ <- ZIO.log(s"Responding with: $b")
+    } yield b
+  }
+
+  override def stream[A, B](io: A => Stream[StatusException, B]): A => Stream[StatusException, B] = { a =>
+    val logOutput = (b: B) => ZIO.log(s"Responding with: $b")
+
+    a match {
+      case s: Stream[StatusException, Any] @unchecked =>
+        val loggedInput = s.tap(r => ZIO.log(s"Received request: $r")).asInstanceOf[A]
+        io(loggedInput).tap(logOutput)
+
+      case _ =>
+        val logInput = ZIO.log(s"Received request: $a")
+        ZStream.fromZIO(logInput).drain ++ io(a).tap(logOutput)
+    }
+  }
+}