From 7c3d2f90b77e4d0daa2e0580e48aa37af79c2418 Mon Sep 17 00:00:00 2001
From: Jules Ivanic <jules.ivanic@gmail.com>
Date: Wed, 9 Oct 2024 12:55:53 +1100
Subject: [PATCH] Add ZIO module providing `connect` and `transact` mechanism
 adapted to ZIO

---
 build.sbt                                     |  13 ++-
 .../com/augustnagro/magnum/zio/zio.scala      | 100 ++++++++++++++++++
 2 files changed, 112 insertions(+), 1 deletion(-)
 create mode 100644 magnum-zio/src/main/scala/com/augustnagro/magnum/zio/zio.scala

diff --git a/build.sbt b/build.sbt
index a04f4ce..5c93f8b 100644
--- a/build.sbt
+++ b/build.sbt
@@ -37,7 +37,7 @@ val testcontainersVersion = "0.40.12"
 
 lazy val root = project
   .in(file("."))
-  .aggregate(magnum, magnumPg)
+  .aggregate(magnum, magnumPg, magnumZio)
 
 lazy val magnum = project
   .in(file("magnum"))
@@ -73,3 +73,14 @@ lazy val magnumPg = project
       "com.dimafeng" %% "testcontainers-scala-postgresql" % testcontainersVersion % Test
     )
   )
+
+lazy val magnumZio = project
+  .in(file("magnum-zio"))
+  .dependsOn(magnum)
+  .settings(
+    Test / fork := true,
+    publish / skip := false,
+    libraryDependencies ++= Seq(
+      "dev.zio" %% "zio" % "2.1.9" % Provided
+    )
+  )
diff --git a/magnum-zio/src/main/scala/com/augustnagro/magnum/zio/zio.scala b/magnum-zio/src/main/scala/com/augustnagro/magnum/zio/zio.scala
new file mode 100644
index 0000000..16b3668
--- /dev/null
+++ b/magnum-zio/src/main/scala/com/augustnagro/magnum/zio/zio.scala
@@ -0,0 +1,100 @@
+package com.augustnagro.magnum.zio
+
+import _root_.zio.{Trace, ZIO}
+import com.augustnagro.magnum.*
+
+import javax.sql.DataSource
+import scala.util.control.NonFatal
+
+/** Executes a given query on a given DataSource
+  *
+  * Re-implementation of
+  * [[com.augustnagro.magnum.connect(dataSource: DataSource)]] for ZIO
+  *
+  * Usage:
+  * {{{
+  *   connect(datasource) { cn ?=> repo.findById(id) }
+  * }}}
+  */
+def connect[A](dataSource: DataSource)(q: DbCon ?=> A)(implicit
+    trace: Trace
+): ZIO[Any, Throwable, A] =
+  connect(Transactor(dataSource))(cn ?=> q(using cn))
+
+/** Executes a given query on a given Transactor
+  *
+  * Re-implementation of
+  * [[com.augustnagro.magnum.connect(dataSource: DataSource)]] for ZIO
+  *
+  * Usage:
+  * {{{
+  *   connect(transactor) { cn ?=> repo.findById(id) }
+  * }}}
+  */
+def connect[A](
+    transactor: Transactor
+)(q: DbCon ?=> A)(implicit trace: Trace): ZIO[Any, Throwable, A] =
+  ZIO.blocking {
+    ZIO.acquireReleaseWith(
+      acquire = ZIO.attempt(transactor.dataSource.getConnection())
+    )(release =
+      conn => ZIO.attempt(if (conn ne null) conn.close() else ()).orDie
+    ) { cn =>
+      ZIO.attempt {
+        transactor.connectionConfig(cn)
+        q(using DbCon(cn, transactor.sqlLogger))
+      }
+    }
+  }
+
+/** Executes a given transaction on a given DataSource
+  *
+  * Re-implementation of
+  * [[com.augustnagro.magnum.transact(transactor: Transactor)]] for ZIO
+  *
+  * Usage:
+  * {{{
+  *   transact(dataSource) { tx ?=> repo.insertReturning(creator) }
+  * }}}
+  */
+def transact[A](dataSource: DataSource)(q: DbTx ?=> A)(implicit
+    trace: Trace
+): ZIO[Any, Throwable, A] =
+  transact(Transactor(dataSource))(tx ?=> q(using tx))
+
+/** Executes a given transaction on a given Transactor
+  *
+  * Re-implementation of
+  * [[com.augustnagro.magnum.transact(transactor: Transactor)]] for ZIO
+  *
+  * Usage:
+  * {{{
+  *   transact(transactor) { tx ?=> repo.insertReturning(creator) }
+  * }}}
+  */
+def transact[A](
+    transactor: Transactor
+)(q: DbTx ?=> A)(implicit trace: Trace): ZIO[Any, Throwable, A] =
+  ZIO.blocking {
+    ZIO.acquireReleaseWith(
+      acquire = ZIO.attempt(transactor.dataSource.getConnection())
+    )(release =
+      conn => ZIO.attempt(if (conn ne null) conn.close() else ()).orDie
+    ) { cn =>
+      ZIO.uninterruptible {
+        ZIO.attempt {
+          transactor.connectionConfig(cn)
+          cn.setAutoCommit(false)
+          try {
+            val res = q(using DbTx(cn, transactor.sqlLogger))
+            cn.commit()
+            res
+          } catch {
+            case NonFatal(t) =>
+              cn.rollback()
+              throw t
+          }
+        }
+      }
+    }
+  }