From ad4e72a0e7893149e4204763ac95f89b9b7517a7 Mon Sep 17 00:00:00 2001
From: Jules Ivanic <jules.ivanic@gmail.com>
Date: Fri, 20 Oct 2023 13:04:06 +0400
Subject: [PATCH] Fix #2333 and #2325 (#2923)

---
 .../scala/io/getquill/util/QueryLogger.scala  | 50 ++++-------
 .../main/scala/zio/logging/extensions.scala   | 87 +++++++++++++++++++
 2 files changed, 105 insertions(+), 32 deletions(-)
 create mode 100644 quill-core/src/main/scala/zio/logging/extensions.scala

diff --git a/quill-core/src/main/scala/io/getquill/util/QueryLogger.scala b/quill-core/src/main/scala/io/getquill/util/QueryLogger.scala
index ef3dd20f7e..c3b357ec1c 100644
--- a/quill-core/src/main/scala/io/getquill/util/QueryLogger.scala
+++ b/quill-core/src/main/scala/io/getquill/util/QueryLogger.scala
@@ -3,6 +3,7 @@ package io.getquill.util
 import io.getquill.util.Messages.LogToFile
 import zio._
 import zio.logging._
+import zio.logging.extensions.executeWithLogger
 
 import java.nio.file.Paths
 import java.time.ZonedDateTime
@@ -10,41 +11,26 @@ import java.time.format.DateTimeFormatter
 
 class QueryLogger(logToFile: LogToFile) {
 
-  val runtime =
-    Unsafe.unsafe { implicit u =>
-      logToFile match {
-        case LogToFile.Enabled(logFile) =>
-          Some(
-            Runtime.unsafe.fromLayer(
-              fileLogger(
-                FileLoggerConfig(
-                  destination = Paths.get(logFile),
-                  format = LogFormat.line,
-                  filter = LogFilter.logLevel(LogLevel.Info)
-                )
-              )
-            )
+  def apply(queryString: String, sourcePath: String, line: Int, column: Int): Unit =
+    logToFile match {
+      case LogToFile.Enabled(logFile) =>
+        val config =
+          FileLoggerConfig(
+            destination = Paths.get(logFile),
+            format = LogFormat.line,
+            filter = LogFilter.logLevel(LogLevel.Info)
           )
-        case LogToFile.Disabled => None
-      }
-    }
 
-  def apply(queryString: String, sourcePath: String, line: Int, column: Int): Unit =
-    runtime match {
-      case Some(runtimeValue) =>
-        Unsafe.unsafe { implicit u =>
-          runtimeValue.unsafe
-            .run(
-              ZIO.logInfo(
-                s"""
-                   |-- file: $sourcePath:$line:$column
-                   |-- time: ${ZonedDateTime.now().format(DateTimeFormatter.ISO_LOCAL_DATE_TIME)}
-                   |$queryString;
-                   |""".stripMargin
-              )
+        executeWithLogger(config) {
+          ZIO
+            .logInfo(
+              s"""
+                 |-- file: $sourcePath:$line:$column
+                 |-- time: ${ZonedDateTime.now().format(DateTimeFormatter.ISO_LOCAL_DATE_TIME)}
+                 |$queryString;
+                 |""".stripMargin
             )
-            .getOrThrow()
         }
-      case None => // do nothing
+      case LogToFile.Disabled => // do nothing
     }
 }
diff --git a/quill-core/src/main/scala/zio/logging/extensions.scala b/quill-core/src/main/scala/zio/logging/extensions.scala
new file mode 100644
index 0000000000..588dd54063
--- /dev/null
+++ b/quill-core/src/main/scala/zio/logging/extensions.scala
@@ -0,0 +1,87 @@
+package zio.logging
+
+import zio.{FiberRef, Runtime, Scope, UIO, Unsafe, ZIO, ZLogger}
+
+import java.nio.charset.Charset
+import java.nio.file.Path
+
+object extensions {
+
+  /**
+   * This code is basically an inlined version of the [[fileLogger]] code from
+   * which we removed the ZLayer part which is getting in our way. We don't need
+   * layers here anyway.
+   */
+  def executeWithLogger(config: FileLoggerConfig)(zio: UIO[Unit]): Unit =
+    Unsafe.unsafe { implicit u =>
+      Runtime.default.unsafe.run {
+        ZIO.scoped {
+          for {
+            logger <- makeFileLogger(config)
+            _      <- FiberRef.currentLoggers.locallyScopedWith(_ + logger)
+            _      <- zio
+          } yield ()
+        }
+      }.getOrThrow()
+    }
+
+  /**
+   * Copied from zio-logging internals.
+   *
+   * We need to copy the function as it's private and we can't access it
+   *
+   * See:
+   *   - https://github.com/zio/zio-logging/blob/v2.1.14/core/shared/src/main/scala/zio/logging/package.scala#L459-L468
+   */
+  def makeFileLogger(config: FileLoggerConfig): ZIO[Scope, Throwable, ZLogger[String, Any]] =
+    makeFileLogger(
+      config.destination,
+      config.format.toLogger,
+      config.filter,
+      config.charset,
+      config.autoFlushBatchSize,
+      config.bufferedIOSize,
+      config.rollingPolicy
+    )
+
+  /**
+   * Adapted from zio-logging internals
+   *
+   *   1. We need to copy the function as it's private and we can't access it.
+   *      2. We wrapped the code we copied into a ZIO to be sure to correctly
+   *      close the writer.
+   *
+   * See:
+   *   - https://github.com/zio/zio-logging/blob/v2.1.14/core/shared/src/main/scala/zio/logging/package.scala#L470-L491
+   */
+  def makeFileLogger(
+    destination: Path,
+    logger: ZLogger[String, String],
+    logFilter: LogFilter[String],
+    charset: Charset,
+    autoFlushBatchSize: Int,
+    bufferedIOSize: Option[Int],
+    rollingPolicy: Option[FileLoggerConfig.FileRollingPolicy]
+  ): ZIO[Scope, Throwable, ZLogger[String, Any]] =
+    for {
+      logWriter <- ZIO.acquireRelease {
+                     ZIO.attempt(
+                       new zio.logging.internal.FileWriter(
+                         destination,
+                         charset,
+                         autoFlushBatchSize,
+                         bufferedIOSize,
+                         rollingPolicy
+                       )
+                     )
+                   }(writer => ZIO.succeed(writer.close()))
+      stringLogger = logFilter.filter(logger.map { (line: String) =>
+                       try logWriter.writeln(line)
+                       catch {
+                         case t: VirtualMachineError => throw t
+                         case _: Throwable           => ()
+                       }
+                     })
+    } yield stringLogger
+
+}