Update zio-query to 0.7.1 (#170)

yoohaemin · Jun 9, 2024 · 5514540 · 5514540
1 parent 1a0ed9e
commit 5514540
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 17 deletions.
diff --git a/build.sbt b/build.sbt
@@ -285,7 +285,7 @@ lazy val V = new {
 
   val cats         = "2.12.0"
   val zio          = "2.0.22"
-  val zioQuery     = "0.6.1"
+  val zioQuery     = "0.7.1"
   val fetch        = "3.1.2"
   val izumiReflect = "2.3.8"
   val scalacheck   = "1.18.0"

diff --git a/zquery/src/main/scala/decrel/reify/zquery.scala b/zquery/src/main/scala/decrel/reify/zquery.scala
@@ -13,7 +13,7 @@ import izumi.reflect.Tag
 import zio.*
 import zio.query.{ CompletedRequestMap, DataSource, ZQuery }
 
-import scala.collection.{ BuildFrom, IterableOps }
+import scala.collection.{ mutable, BuildFrom, IterableOps }
 
 trait zquery[R] extends bifunctor.module[ZQuery[R, +*, +*]] {
 
@@ -62,22 +62,27 @@ trait zquery[R] extends bifunctor.module[ZQuery[R, +*, +*]] {
         val deduplicated = requests.distinctBy(_.id)
 
         batchExecute(deduplicated.map(_.id)).flatMap { results =>
-          val resultsMap = results.toMap
-          ZIO.foldLeft(requests)(CompletedRequestMap.empty) { (crm, request) =>
-            resultsMap.get(request.id) match {
-              case Some(result) =>
-                ZIO.succeed(crm.insert(request, Exit.succeed(result)))
-              case None =>
-                ZIO.die(
-                  new NoSuchElementException(s"Response for request not found: ${request.id}")
-                )
-            }
-          }
+          val mapBuilder = mutable.Map.newBuilder[In, Exit[E, Out]]
+          mapBuilder.sizeHint(results)
+          val resultsMap = mapBuilder
+            .addAll(results.view.map(pair => pair._1 -> Exit.succeed(pair._2)))
+            .result()
+            .withDefault(in =>
+              Exit.die(new NoSuchElementException(s"Response for request not found: $in"))
+            )
+
+          ZIO.succeed(
+            CompletedRequestMap.fromIterableWith[E, RelationRequest[Rel, In, E, Out], Out](
+              requests
+            )(
+              (a: zio.query.Request[E, Out]) => a,
+              request => resultsMap(request.id)
+            )
+          )
         }.catchAll { e =>
+          val failure = Exit.fail(e)
           ZIO.succeed(
-            requests.foldLeft(CompletedRequestMap.empty) { (crm, request) =>
-              crm.insert(request, Exit.fail(e))
-            }
+            CompletedRequestMap.fromIterableWith(requests)(identity, _ => failure)
           )
         }
       }

diff --git a/zquery/src/test/scala/decrel/reify/zqueryProofSpec.scala b/zquery/src/test/scala/decrel/reify/zqueryProofSpec.scala
@@ -419,7 +419,7 @@ object zqueryProofSpec extends ZIOSpecDefault {
                     User.fetch          -> Chunk(rental2.userId),
                     User.currentRentals -> Chunk(user2),
                     // The following two calls have been batched, avoiding N+1 problem
-                    Book.fetch -> Chunk(rental3.bookId, rental2.bookId)
+                    Book.fetch -> Chunk(rental2.bookId, rental3.bookId)
                   )
                 )
               )