Merge pull request #31 from nebula-contrib/opt

opt usage
nebula-contrib · Dec 25, 2023 · b2badd4 · b2badd4
2 parents e428ea9 + e4d79cd
commit b2badd4
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 34 deletions.
diff --git a/core/src/main/scala/zio/nebula/NebulaConfig.scala b/core/src/main/scala/zio/nebula/NebulaConfig.scala
@@ -55,9 +55,9 @@ final case class NebulaPoolConfig(
 }
 
 final case class NebulaSessionPoolConfig(
-  address: List[NebulaHostAddress],
-  auth: NebulaAuth,
-  spaceName: String,
+  address: List[NebulaHostAddress], // both for NebulaClient and NebulaSessionClient
+  auth: NebulaAuth,                 // both for NebulaClient and NebulaSessionClient
+  spaceName: String,                // both for NebulaClient and NebulaSessionClient
   maxSessionSize: Int = 10,
   minSessionSize: Int = 1,
   waitTimeMills: Int = 0,
@@ -66,7 +66,7 @@ final case class NebulaSessionPoolConfig(
   intervalTimeMills: Int = 0,
   healthCheckTimeSeconds: Int = 600,
   cleanTimeSeconds: Int = 3600,
-  reconnect: Boolean = false,
+  reconnect: Boolean = false,       // both for NebulaClient and NebulaSessionClient
   useHttp2: Boolean = false
 )
 

diff --git a/core/src/main/scala/zio/nebula/net/NebulaClient.scala b/core/src/main/scala/zio/nebula/net/NebulaClient.scala
@@ -18,7 +18,9 @@ trait NebulaClient {
 
   def close(): Task[Unit]
 
-  def openSession(): ZIO[Scope & NebulaSessionPoolConfig, Throwable, NebulaSession]
+  def openSession(): ZIO[NebulaSessionPoolConfig, Throwable, NebulaSession]
+
+  def openSession(sessionPoolConfig: NebulaSessionPoolConfig): ZIO[Any, Throwable, NebulaSession]
 
   def activeConnNum: Task[Int]
 

diff --git a/core/src/main/scala/zio/nebula/net/NebulaClientLive.scala b/core/src/main/scala/zio/nebula/net/NebulaClientLive.scala
@@ -30,21 +30,34 @@ private[nebula] final class NebulaClientLive(underlying: NebulaPl) extends Nebul
 
   def close(): Task[Unit] = ZIO.attempt(underlying.close())
 
-  def openSession(): ZIO[Scope & NebulaSessionPoolConfig, Throwable, NebulaSession] =
+  def openSession(sessionPoolConfig: NebulaSessionPoolConfig): ZIO[Any, Throwable, NebulaSession] =
+    for {
+      session <- ZIO.attempt(
+                   new NebulaSession(
+                     underlying.getSession(
+                       sessionPoolConfig.auth.username,
+                       sessionPoolConfig.auth.password,
+                       sessionPoolConfig.reconnect
+                     )
+                   )
+                 )
+      _       <- session.execute(Stmt.str(s"USE `${sessionPoolConfig.spaceName}`"))
+    } yield session
+
+  def openSession(): ZIO[NebulaSessionPoolConfig, Throwable, NebulaSession] =
     for {
       sessionConfig <- ZIO.service[NebulaSessionPoolConfig]
-      session       <-
-        ZIO.acquireRelease(
-          ZIO.attempt(
-            new NebulaSession(
-              underlying.getSession(
-                sessionConfig.auth.username,
-                sessionConfig.auth.password,
-                sessionConfig.reconnect
-              )
-            )
-          )
-        )(_.close().onError(e => ZIO.logErrorCause(e)).ignoreLogged)
+      session       <- ZIO.attempt(
+                         new NebulaSession(
+                           underlying.getSession(
+                             sessionConfig.auth.username,
+                             sessionConfig.auth.password,
+                             sessionConfig.reconnect
+                           )
+                         )
+                       )
+      _             <- session.execute(Stmt.str(s"USE `${sessionConfig.spaceName}`"))
+
     } yield session
 
   def activeConnNum: Task[Int] = ZIO.attempt(underlying.getActiveConnNum)

diff --git a/core/src/main/scala/zio/nebula/net/NebulaSession.scala b/core/src/main/scala/zio/nebula/net/NebulaSession.scala
@@ -3,7 +3,7 @@ package zio.nebula.net
 import scala.jdk.CollectionConverters._
 
 import zio._
-import zio.nebula.NebulaResultSet
+import zio.nebula.{ GlobalSettings, NebulaResultSet }
 
 import com.vesoft.nebula.client.graph.data.HostAddress
 import com.vesoft.nebula.client.graph.net.Session
@@ -15,22 +15,23 @@ import com.vesoft.nebula.client.graph.net.Session
  */
 final class NebulaSession(private val underlying: Session) {
 
-  def execute(stmt: Stmt): Task[stmt.T] = ZIO.attempt {
-    stmt match {
-      case StringStmt(_stmt)                       =>
-        new NebulaResultSet(underlying.execute(_stmt)).asInstanceOf[stmt.T]
-      case StringStmtWithMap(_stmt, parameterMap)  =>
-        new NebulaResultSet(underlying.executeWithParameter(_stmt, parameterMap.asJava)).asInstanceOf[stmt.T]
-      case JsonStmt(jsonStmt)                      =>
-        underlying
-          .executeJson(jsonStmt)
-          .asInstanceOf[stmt.T]
-      case JsonStmtWithMap(jsonStmt, parameterMap) =>
-        underlying
-          .executeJsonWithParameter(jsonStmt, parameterMap.asJava)
-          .asInstanceOf[stmt.T]
+  def execute(stmt: Stmt): Task[stmt.T] =
+    GlobalSettings.printLog(stmt.toString) *> ZIO.attempt {
+      stmt match {
+        case StringStmt(_stmt)                       =>
+          new NebulaResultSet(underlying.execute(_stmt)).asInstanceOf[stmt.T]
+        case StringStmtWithMap(_stmt, parameterMap)  =>
+          new NebulaResultSet(underlying.executeWithParameter(_stmt, parameterMap.asJava)).asInstanceOf[stmt.T]
+        case JsonStmt(jsonStmt)                      =>
+          underlying
+            .executeJson(jsonStmt)
+            .asInstanceOf[stmt.T]
+        case JsonStmtWithMap(jsonStmt, parameterMap) =>
+          underlying
+            .executeJsonWithParameter(jsonStmt, parameterMap.asJava)
+            .asInstanceOf[stmt.T]
+      }
     }
-  }
 
   def ping(): Task[Boolean] = ZIO.attempt(underlying.ping())