nebula-contrib · jxnu-liguobin · Jan 15, 2024 · Jan 13, 2024 · Jan 15, 2024
diff --git a/README.md b/README.md
@@ -50,7 +50,7 @@ import zio.nebula._
 final class NebulaSessionClientExample(sessionClient: NebulaSessionClient) {
 
   def execute(stmt: String): ZIO[Any, Throwable, NebulaResultSet] = {
-    // Your business logic
+    // your custom logic
     sessionClient.execute(stmt)
   }
 }
@@ -62,14 +62,15 @@ object NebulaSessionClientExample {
 object NebulaSessionClientMain extends ZIOAppDefault {
 
   override def run = (for {
-    _ <- ZIO.serviceWithZIO[NebulaSessionClient](_.init()) // since 0.1.1, no need to call it manually. 
+    // since 0.1.1, no need to call it manually.
+    _ <- ZIO.serviceWithZIO[NebulaSessionClient](_.init()) 
     _ <- ZIO.serviceWithZIO[NebulaSessionClientExample](
              _.execute("""
                          |INSERT VERTEX person(name, age) VALUES 
                          |'Bob':('Bob', 10), 
                          |'Lily':('Lily', 9),'Tom':('Tom', 10),
                          |'Jerry':('Jerry', 13),
-                         |'John':('John', 11);""".stripMargin).flatMap(r => ZIO.logInfo(r.toString))
+                         |'John':('John', 11);""".stripMargin)
            )
     _ <- ZIO.serviceWithZIO[NebulaSessionClientExample](
              _.execute("""
@@ -78,14 +79,13 @@ object NebulaSessionClientMain extends ZIOAppDefault {
                          |'Bob'->'Tom':(70.0),
                          |'Lily'->'Jerry':(84.0),
                          |'Tom'->'Jerry':(68.3),
-                         |'Bob'->'John':(97.2);""".stripMargin).flatMap(r => ZIO.logInfo(r.toString))
+                         |'Bob'->'John':(97.2);""".stripMargin)
            )
     _ <- ZIO.serviceWithZIO[NebulaSessionClientExample](
              _.execute("""
                          |USE test;
                          |MATCH (p:person) RETURN p LIMIT 4;
                          |""".stripMargin)
-               .flatMap(r => ZIO.logInfo(r.rows.toString()))
            )
   } yield ())
     .provide(
@@ -100,10 +100,10 @@ object NebulaSessionClientMain extends ZIOAppDefault {
 ## Configuration
 
 Introduction for configuring keys:
-  - key `graph` for `NebulaSessionClient`
-  - key `meta` for `NebulaMetaClient`
-  - key `storage` for `NebulaStorageClient`
-  - key `pool` for `NebulaClient`
+  - key `graph` for `NebulaSessionClient`, For structure, please refer to `zio.nebula.NebulaSessionPoolConfig`
+  - key `meta` for `NebulaMetaClient`, For structure, please refer to `zio.nebula.NebulaMetaConfig`
+  - key `storage` for `NebulaStorageClient`, For structure, please refer to `zio.nebula.NebulaStorageConfig`
+  - key `pool` for `NebulaClient`, For structure, please refer to `zio.nebula.NebulaPoolConfig`
 
 Sample Configuration:
 ```hocon

diff --git a/core/src/main/scala/zio/nebula/NebulaConfig.scala b/core/src/main/scala/zio/nebula/NebulaConfig.scala
@@ -28,6 +28,9 @@ final case class NebulaStorageConfig(
 )
 
 final case class NebulaPoolConfig(
+  address: List[NebulaHostAddress],
+  auth: NebulaAuth,
+  spaceName: Option[String],
   minConnsSize: Int = 0,
   maxConnsSize: Int = 10,
   timeoutMills: Int = 0,
@@ -36,7 +39,8 @@ final case class NebulaPoolConfig(
   waitTimeMills: Int = 0,
   minClusterHealthRate: Double = 1d,
   enableSsl: Boolean = false,
-  sslParam: Option[SSLParam]
+  sslParam: Option[SSLParam],
+  reconnect: Boolean = false
 ) {
 
   def toJava: PoolConfig = {
@@ -55,9 +59,9 @@ final case class NebulaPoolConfig(
 }
 
 final case class NebulaSessionPoolConfig(
-  address: List[NebulaHostAddress], // both for NebulaClient and NebulaSessionClient
-  auth: NebulaAuth,                 // both for NebulaClient and NebulaSessionClient
-  spaceName: String,                // both for NebulaClient and NebulaSessionClient
+  address: List[NebulaHostAddress],
+  auth: NebulaAuth,
+  spaceName: String,
   maxSessionSize: Int = 10,
   minSessionSize: Int = 1,
   waitTimeMills: Int = 0,
@@ -66,7 +70,7 @@ final case class NebulaSessionPoolConfig(
   intervalTimeMills: Int = 0,
   healthCheckTimeSeconds: Int = 600,
   cleanTimeSeconds: Int = 3600,
-  reconnect: Boolean = false,       // both for NebulaClient and NebulaSessionClient
+  reconnect: Boolean = false,
   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
@@ -2,7 +2,6 @@ package zio.nebula.net
 
 import zio._
 import zio.nebula._
-import zio.nebula.NebulaPoolConfig
 
 import com.vesoft.nebula.client.graph.{ NebulaPoolConfig => _ }
 import com.vesoft.nebula.client.graph.net.{ NebulaPool => Pool }
@@ -14,13 +13,22 @@ import com.vesoft.nebula.client.graph.net.{ NebulaPool => Pool }
  */
 trait NebulaClient {
 
-  def init(): ZIO[NebulaSessionPoolConfig & NebulaPoolConfig, Throwable, Boolean]
+  def init(): ZIO[NebulaPoolConfig, Throwable, Boolean]
 
+  /**
+   * close the client
+   */
   def close(): Task[Unit]
 
-  def openSession(): ZIO[NebulaSessionPoolConfig, Throwable, NebulaSession]
+  /**
+   * init the client and execute `USE spaceName` if exists
+   */
+  def openSession(useSpace: Boolean): ZIO[NebulaPoolConfig, Throwable, NebulaSession]
 
-  def openSession(sessionPoolConfig: NebulaSessionPoolConfig): ZIO[Any, Throwable, NebulaSession]
+  /**
+   * init the client by using poolConfig
+   */
+  def openSession(poolConfig: NebulaPoolConfig): 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
@@ -4,7 +4,6 @@ import scala.jdk.CollectionConverters._
 
 import zio._
 import zio.nebula._
-import zio.nebula.NebulaPoolConfig
 
 import com.vesoft.nebula.client.graph.{ NebulaPoolConfig => _ }
 import com.vesoft.nebula.client.graph.data.HostAddress
@@ -17,46 +16,45 @@ import com.vesoft.nebula.client.graph.net.{ NebulaPool => NebulaPl }
  */
 private[nebula] final class NebulaClientLive(underlying: NebulaPl) extends NebulaClient {
 
-  def init(): ZIO[NebulaSessionPoolConfig & NebulaPoolConfig, Throwable, Boolean] =
+  def init(): ZIO[NebulaPoolConfig, Throwable, Boolean] =
     for {
       config <- ZIO.service[NebulaPoolConfig]
       status <-
-        ZIO.serviceWithZIO[NebulaSessionPoolConfig](sessionConfig =>
-          ZIO.attempt(
-            underlying.init(sessionConfig.address.map(d => new HostAddress(d.host, d.port)).asJava, config.toJava)
-          )
+        ZIO.attempt(
+          underlying.init(config.address.map(d => new HostAddress(d.host, d.port)).asJava, config.toJava)
         )
     } yield status
 
   def close(): Task[Unit] = ZIO.attempt(underlying.close())
 
-  def openSession(sessionPoolConfig: NebulaSessionPoolConfig): ZIO[Any, Throwable, NebulaSession] =
+  def openSession(poolConfig: NebulaPoolConfig): ZIO[Any, Throwable, NebulaSession] =
     for {
       session <- ZIO.attempt(
                    new NebulaSession(
                      underlying.getSession(
-                       sessionPoolConfig.auth.username,
-                       sessionPoolConfig.auth.password,
-                       sessionPoolConfig.reconnect
+                       poolConfig.auth.username,
+                       poolConfig.auth.password,
+                       poolConfig.reconnect
                      )
                    )
                  )
-      _       <- session.execute(Stmt.str(s"USE `${sessionPoolConfig.spaceName}`"))
     } yield session
 
-  def openSession(): ZIO[NebulaSessionPoolConfig, Throwable, NebulaSession] =
+  def openSession(useSpace: Boolean): ZIO[NebulaPoolConfig, Throwable, NebulaSession] =
     for {
-      sessionConfig <- ZIO.service[NebulaSessionPoolConfig]
-      session       <- ZIO.attempt(
-                         new NebulaSession(
-                           underlying.getSession(
-                             sessionConfig.auth.username,
-                             sessionConfig.auth.password,
-                             sessionConfig.reconnect
-                           )
-                         )
-                       )
-      _             <- session.execute(Stmt.str(s"USE `${sessionConfig.spaceName}`"))
+      poolConfig <- ZIO.service[NebulaPoolConfig]
+      session    <- ZIO.attempt(
+                      new NebulaSession(
+                        underlying.getSession(
+                          poolConfig.auth.username,
+                          poolConfig.auth.password,
+                          poolConfig.reconnect
+                        )
+                      )
+                    )
+      _          <- ZIO.when(useSpace && poolConfig.spaceName.nonEmpty) {
+                      session.execute(Stmt.str(s"USE `${poolConfig.spaceName.orNull}`"))
+                    }
 
     } yield session
 

diff --git a/core/src/main/scala/zio/nebula/package.scala b/core/src/main/scala/zio/nebula/package.scala
@@ -12,7 +12,7 @@ import zio.nebula.storage._
 package object nebula {
 
   type SessionClient = NebulaSessionClient
-  type Client        = NebulaClient with NebulaSessionPoolConfig with NebulaPoolConfig
+  type Client        = NebulaClient with NebulaPoolConfig
   type Storage       = NebulaStorageClient
   type Meta          = NebulaMetaClient
 

diff --git a/core/src/test/scala/zio/nebula/NebulaSpec.scala b/core/src/test/scala/zio/nebula/NebulaSpec.scala
@@ -22,12 +22,12 @@ trait NebulaSpec extends ZIOSpecDefault {
     (specLayered @@ beforeAll(
       ZIO.serviceWithZIO[NebulaClient](_.init())
         *> ZIO.serviceWithZIO[NebulaClient](
-          _.openSession().flatMap(_.execute(Stmt.str("""
-                                                       |CREATE SPACE IF NOT EXISTS test(vid_type=fixed_string(20));
-                                                       |USE test;
-                                                       |CREATE TAG IF NOT EXISTS person(name string, age int);
-                                                       |CREATE EDGE IF NOT EXISTS like(likeness double)
-                                                       |""".stripMargin)))
+          _.openSession(false).flatMap(_.execute(Stmt.str("""
+                                                            |CREATE SPACE IF NOT EXISTS test(vid_type=fixed_string(20));
+                                                            |USE test;
+                                                            |CREATE TAG IF NOT EXISTS person(name string, age int);
+                                                            |CREATE EDGE IF NOT EXISTS like(likeness double)
+                                                            |""".stripMargin)))
         )
     ) @@ sequential @@ eventually)
       .provideShared(

diff --git a/core/src/test/scala/zio/nebula/ZioNebulaEnvironment.scala b/core/src/test/scala/zio/nebula/ZioNebulaEnvironment.scala
@@ -39,15 +39,10 @@ object ZioNebulaEnvironment {
         maxConnsSize = 10,
         intervalIdleMills = 100,
         waitTimeMills = 100,
-        sslParam = None
-      )
-    ) ++ ZLayer.fromZIO(
-      ZIO.attempt(
-        NebulaSessionPoolConfig(
-          List(NebulaHostAddress(host, port)),
-          NebulaAuth(defaultUser, defaultPwd),
-          defaultSpace
-        )
+        sslParam = None,
+        address = List(NebulaHostAddress(host, port)),
+        auth = NebulaAuth(defaultUser, defaultPwd),
+        spaceName = Some(defaultSpace)
       )
     )
 

diff --git a/examples/src/main/scala/zio/nebula/example/NebulaClientExample.scala b/examples/src/main/scala/zio/nebula/example/NebulaClientExample.scala
@@ -6,8 +6,8 @@ import zio.nebula.net.{ NebulaClient, Stmt }
 
 final class NebulaClientExample(nebulaClient: NebulaClient) {
 
-  def execute(stmt: String): ZIO[Scope with NebulaSessionPoolConfig, Throwable, NebulaResultSet] =
-    nebulaClient.openSession().flatMap(_.execute(Stmt.str(stmt)))
+  def execute(stmt: String): ZIO[Scope with NebulaPoolConfig, Throwable, NebulaResultSet] =
+    nebulaClient.openSession(false).flatMap(_.execute(Stmt.str(stmt)))
 }
 
 object NebulaClientExample {