revise GTransform interface

portone-io · Mar 13, 2024 · 781dfbe · 781dfbe
1 parent 2bc4f0d
commit 781dfbe
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 101 deletions.
diff --git a/build.sbt b/build.sbt
@@ -23,9 +23,9 @@ val commonPublishSettings = Seq(
     )
   ),
   publishTo := Some(
-    "GitHub Package Registry (iamport/zio-grpc)" at "https://maven.pkg.github.com/iamport/zio-grpc"
+    "GitHub Package Registry (portone-io/zio-grpc)" at "https://maven.pkg.github.com/portone-io/zio-grpc"
   ),
-  resolvers += "GitHub Package Registry (iamport/zio-grpc)" at "https://maven.pkg.github.com/iamport/zio-grpc"
+  resolvers += "GitHub Package Registry (portone-io/zio-grpc)" at "https://maven.pkg.github.com/portone-io/zio-grpc"
 )
 
 publish / skip := true
@@ -61,7 +61,7 @@ lazy val core = projectMatrix
       "dev.zio" %%% "zio-test-sbt" % Version.zio % "test"
     )
   )
-  .settings(commonPublishSettings)
+//  .settings(commonPublishSettings)
   .jvmPlatform(
     ScalaVersions,
     Seq(
@@ -97,7 +97,7 @@ lazy val codeGen = projectMatrix
       "com.thesamet.scalapb" %% "compilerplugin" % scalapb.compiler.Version.scalapbVersion
     )
   )
-  .settings(commonPublishSettings)
+//  .settings(commonPublishSettings)
   .jvmPlatform(scalaVersions = ScalaVersions)
 
 lazy val codeGenJVM212 = codeGen.jvm(Scala212)
@@ -113,7 +113,7 @@ lazy val protocGenZio = protocGenProject("protoc-gen-zio", codeGenJVM212)
         (assembly / assemblyMergeStrategy).value.apply(x)
     }
   )
-  .settings(commonPublishSettings)
+//  .settings(commonPublishSettings)
 
 lazy val e2eProtos =
   projectMatrix

diff --git a/code-gen/src/main/scala/scalapb/zio_grpc/ZioCodeGenerator.scala b/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((req, ctx) => $delegate(req, ctx))(request, context)"
         case StreamType.ServerStreaming | StreamType.Bidirectional =>
-          s"f.stream($delegate(request, _))(context)"
+          s"f.stream((req, ctx) => $delegate(req, ctx))(request, context)"
       }
       fp.add(
         methodSignature(
@@ -246,52 +246,6 @@ class ZioFilePrinter(
       )
     }
 
-    def printClientWithResponseMetadataTransform(
-        fp: FunctionalPrinter,
-        method: MethodDescriptor
-    ): FunctionalPrinter = {
-      val delegate = s"self.${method.name}"
-      val newImpl  = method.streamType match {
-        case StreamType.Unary           =>
-          s"f.effect($delegate(request))"
-        case StreamType.ServerStreaming =>
-          s"f.stream($delegate(request))"
-        case StreamType.ClientStreaming =>
-          s"f.effect($delegate(request))"
-        case StreamType.Bidirectional   =>
-          s"f.stream($delegate(request))"
-      }
-      fp.add(
-        clientWithResponseMetadataSignature(
-          method,
-          "Any"
-        ) + " = " + newImpl
-      )
-    }
-
-    def printClientTransform(
-        fp: FunctionalPrinter,
-        method: MethodDescriptor
-    ): FunctionalPrinter = {
-      val delegate = s"self.${method.name}"
-      val newImpl  = method.streamType match {
-        case StreamType.Unary           =>
-          s"f.effect($delegate(request))"
-        case StreamType.ServerStreaming =>
-          s"f.stream($delegate(request))"
-        case StreamType.ClientStreaming =>
-          s"f.effect($delegate(request))"
-        case StreamType.Bidirectional   =>
-          s"f.stream($delegate(request))"
-      }
-      fp.add(
-        clientMethodSignature(
-          method,
-          contextType = "Any"
-        ) + " = " + newImpl
-      )
-    }
-
     def printGTransformMethod(fp: FunctionalPrinter): FunctionalPrinter =
       fp.add(
         s"def transform[Context1, Error1](f: $GTransform[Context, Error, Context1, Error1]): ${gtraitName.fullName}[Context1, Error1] = new ${gtraitName.fullName}[Context1, Error1] {"
@@ -622,6 +576,9 @@ class ZioFilePrinter(
         case StreamType.ServerStreaming | StreamType.Bidirectional =>
           s"zio.stream.ZStream.fromZIO($SafeMetadata.make)"
       }
+
+      val reqType = methodInType(method, StatusException)
+
       fp.add(
         clientWithResponseMetadataSignature(
           method,
@@ -630,21 +587,21 @@ class ZioFilePrinter(
       ).indent
         .add(s"$makeMetadata.flatMap { metadata =>")
         .indented(
-          _.add(s"transforms.$transformMethod { context => ")
+          _.add(s"transforms.$transformMethod { (req: $reqType, ctx) => ")
             .indented(
               _.add(
                 s"$clientCall("
               )
                 .indented(
                   _.add(s"channel,")
-                    .add(s"context.method.asInstanceOf[$methodDescriptor[$Req, $Res]],")
-                    .add("context.options,")
-                    .add("context.metadata,")
-                    .add("request")
+                    .add(s"ctx.method.asInstanceOf[$methodDescriptor[$Req, $Res]],")
+                    .add("ctx.options,")
+                    .add("ctx.metadata,")
+                    .add("req")
                 )
                 .add(")")
             )
-            .add(s"}($ClientCallContext(${method.grpcDescriptor.fullName}, $CallOptions.DEFAULT, metadata))")
+            .add(s"}(request, $ClientCallContext(${method.grpcDescriptor.fullName}, $CallOptions.DEFAULT, metadata))")
         )
         .add("}")
         .outdent

diff --git a/core/src/main/scala/scalapb/zio_grpc/transforms.scala b/core/src/main/scala/scalapb/zio_grpc/transforms.scala
@@ -18,16 +18,16 @@ trait Transform {
   // 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[Req, Resp](
+          io: (Req, Context) => ZIO[Any, StatusException, Resp]
+      ): (Req, Context) => ZIO[Any, StatusException, Resp] = { (req, ctx) =>
+        self.effect(io(req, ctx))
       }
 
-      def stream[A](
-          io: Context => ZStream[Any, StatusException, A]
-      ): Context => ZStream[Any, StatusException, A] = { c =>
-        self.stream(io(c))
+      def stream[Req, Resp](
+          io: (Req, Context) => ZStream[Any, StatusException, Resp]
+      ): (Req, Context) => ZStream[Any, StatusException, Resp] = { (req, ctx) =>
+        self.stream(io(req, ctx))
       }
     }
 
@@ -44,10 +44,11 @@ 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](io: ZIO[Any, StatusException, A]): ZIO[Any, StatusException, A] =
+      ct.effect((req: Any, ctx) => io)((), ())
 
     def stream[A](io: ZStream[Any, StatusException, A]): ZStream[Any, StatusException, A] =
-      ct.stream(_ => io)(())
+      ct.stream((req: Any, ctx) => io)((), ())
   }
 }
 
@@ -58,25 +59,26 @@ 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[Req, Resp](
+      io: (Req, ContextIn) => ZIO[Any, ErrorIn, Resp]
+  ): (Req, ContextOut) => ZIO[Any, ErrorOut, Resp]
+
+  def stream[Req, Resp](
+      io: (Req, ContextIn) => ZStream[Any, ErrorIn, Resp]
+  ): (Req, ContextOut) => ZStream[Any, ErrorOut, Resp]
 
   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[Req, Resp](
+          io: (Req, ContextIn) => ZIO[Any, ErrorIn, Resp]
+      ): (Req, ContextOut2) => ZIO[Any, ErrorOut2, Resp] =
         other.effect(self.effect(io))
 
-      def stream[A](
-          io: ContextIn => ZStream[Any, ErrorIn, A]
-      ): ContextOut2 => ZStream[Any, ErrorOut2, A] =
+      def stream[Req, Resp](
+          io: (Req, ContextIn) => ZStream[Any, ErrorIn, Resp]
+      ): (Req, ContextOut2) => ZStream[Any, ErrorOut2, Resp] =
         other.stream(self.stream(io))
     }
 
@@ -88,45 +90,51 @@ 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[Req, Resp](io: (Req, C) => ZIO[Any, E, Resp]): (Req, C) => ZIO[Any, E, Resp]         = io
+    def stream[Req, Resp](io: (Req, C) => ZStream[Any, E, Resp]): (Req, C) => ZStream[Any, E, Resp] = 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[Req, Resp](
+          io: (Req, ContextIn) => ZIO[Any, Error, Resp]
+      ): (Req, ContextOut) => ZIO[Any, Error, Resp] = { (req, ctx) =>
+        f(ctx).flatMap(newCtx => io(req, newCtx))
       }
 
-      def stream[A](
-          io: ContextIn => ZStream[Any, Error, A]
-      ): ContextOut => ZStream[Any, Error, A] = { (context: ContextOut) =>
-        ZStream.fromZIO(f(context)).flatMap(io)
+      def stream[Req, Resp](
+          io: (Req, ContextIn) => ZStream[Any, Error, Resp]
+      ): (Req, ContextOut) => ZStream[Any, Error, Resp] = { (req, ctx) =>
+        ZStream.fromZIO(f(ctx)).flatMap(newCtx => io(req, newCtx))
       }
     }
 
   // 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[Req, Resp](io: (Req, C) => zio.ZIO[Any, E1, Resp]): (Req, C) => zio.ZIO[Any, E2, Resp] = { (req, ctx) =>
+      io(req, ctx).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[Req, Resp](
+        io: (Req, C) => zio.stream.ZStream[Any, E1, Resp]
+    ): (Req, C) => zio.stream.ZStream[Any, E2, Resp] = { (req, ctx) =>
+      io(req, ctx).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[Req, Resp](io: (Req, C) => zio.ZIO[Any, E1, Resp]): (Req, C) => zio.ZIO[Any, E2, Resp] = { (req, ctx) =>
+      io(req, ctx).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[Req, Resp](
+        io: (Req, C) => zio.stream.ZStream[Any, E1, Resp]
+    ): (Req, C) => zio.stream.ZStream[Any, E2, Resp] = { (req, ctx) =>
+      io(req, ctx).catchAll(e => ZStream.fromZIO(f(e).flatMap(ZIO.fail(_))))
     }
   }
 }

diff --git a/e2e/src/test/scalajvm/scalapb/zio_grpc/ClientTransformSpec.scala b/e2e/src/test/scalajvm/scalapb/zio_grpc/ClientTransformSpec.scala
@@ -18,8 +18,10 @@ object ClientTransformSpec extends ZIOSpecDefault {
   def run(transform: ClientTransform): ZIO[Any, StatusException, ClientCallContext] =
     for {
       metadata <- SafeMetadata.make
-      context  <-
-        transform.effect(ZIO.succeed(_))(ClientCallContext(TestServiceGrpc.METHOD_UNARY, CallOptions.DEFAULT, metadata))
+      context  <- transform.effect((_: Any, ctx) => ZIO.succeed(ctx))(
+                    (),
+                    ClientCallContext(TestServiceGrpc.METHOD_UNARY, CallOptions.DEFAULT, metadata)
+                  )
     } yield context
 
   def spec = suite("ClientTransformSpec")(