feat: custom connection factory when using withSession (#1271)

akka · Nov 26, 2024 · c5a614b · c5a614b
1 parent ba65122
commit c5a614b
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 1 deletion.
diff --git a/akka-projection-r2dbc/src/main/scala/akka/projection/r2dbc/javadsl/R2dbcSession.scala b/akka-projection-r2dbc/src/main/scala/akka/projection/r2dbc/javadsl/R2dbcSession.scala
@@ -17,6 +17,7 @@ import akka.actor.typed.ActorSystem
 import akka.persistence.r2dbc.internal.R2dbcExecutor
 import akka.projection.r2dbc.scaladsl
 import io.r2dbc.spi.Connection
+import io.r2dbc.spi.ConnectionFactory
 import io.r2dbc.spi.Row
 import io.r2dbc.spi.Statement
 
@@ -40,6 +41,19 @@ object R2dbcSession {
     }
   }.asJava
 
+  /**
+   * Provide a custom connectionFactory. The config closeCallsExceeding is loaded from the default path.
+   */
+  def withSession[A](
+      system: ActorSystem[_],
+      connectionFactory: ConnectionFactory,
+      fun: JFunction[R2dbcSession, CompletionStage[A]]): CompletionStage[A] = {
+    scaladsl.R2dbcSession.withSession(system, connectionFactory) { scaladslSession =>
+      val javadslSession = new R2dbcSession(scaladslSession.connection)(system.executionContext, system)
+      fun(javadslSession).asScala
+    }
+  }.asJava
+
 }
 
 final class R2dbcSession(val connection: Connection)(implicit ec: ExecutionContext, system: ActorSystem[_]) {

diff --git a/akka-projection-r2dbc/src/main/scala/akka/projection/r2dbc/scaladsl/R2dbcSession.scala b/akka-projection-r2dbc/src/main/scala/akka/projection/r2dbc/scaladsl/R2dbcSession.scala
@@ -16,6 +16,7 @@ import akka.persistence.r2dbc.ConnectionPoolSettings
 import akka.persistence.r2dbc.internal.R2dbcExecutor
 import akka.projection.r2dbc.R2dbcProjectionSettings
 import io.r2dbc.spi.Connection
+import io.r2dbc.spi.ConnectionFactory
 import io.r2dbc.spi.Row
 import io.r2dbc.spi.Statement
 import org.slf4j.LoggerFactory
@@ -38,8 +39,25 @@ object R2dbcSession {
     val connectionFactoryProvider = ConnectionFactoryProvider(system)
     val connectionFactory = connectionFactoryProvider.connectionFactoryFor(connectionFactoryConfigPath)
     val poolSettings = new ConnectionPoolSettings(system.settings.config.getConfig(connectionFactoryConfigPath))
+    withSession(system, connectionFactory, poolSettings.closeCallsExceeding)(fun)
+  }
+
+  /**
+   * Provide a custom connectionFactory. The config closeCallsExceeding is loaded from the default path.
+   */
+  def withSession[A](system: ActorSystem[_], connectionFactory: ConnectionFactory)(
+      fun: R2dbcSession => Future[A]): Future[A] = {
+    val poolConfig = system.settings.config.getConfig(connectionFactoryConfigPath(system))
+    val poolSettings = new ConnectionPoolSettings(poolConfig)
+    withSession(system, connectionFactory, poolSettings.closeCallsExceeding)(fun)
+  }
+
+  private def withSession[A](
+      system: ActorSystem[_],
+      connectionFactory: ConnectionFactory,
+      closeCallsExceeding: Option[FiniteDuration])(fun: R2dbcSession => Future[A]): Future[A] = {
     val r2dbcExecutor =
-      new R2dbcExecutor(connectionFactory, log, logDbCallsDisabled, poolSettings.closeCallsExceeding)(
+      new R2dbcExecutor(connectionFactory, log, logDbCallsDisabled, closeCallsExceeding)(
         system.executionContext,
         system)
     r2dbcExecutor.withConnection("R2dbcSession") { connection =>