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())