update session initializers on neo4j

nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/Neo4JSpace.java

@@ -24,6 +24,8 @@ import org.apache.logging.log4j.Logger;
 import org.neo4j.driver.AuthTokens;
 import org.neo4j.driver.Driver;
 import org.neo4j.driver.GraphDatabase;
+import org.neo4j.driver.*;
+import org.neo4j.driver.async.AsyncSession;
 
 import java.util.Optional;
 
@@ -32,6 +34,7 @@ public class Neo4JSpace implements AutoCloseable {
     private final static Logger logger = LogManager.getLogger(Neo4JSpace.class);
     private final String space;
     private Driver driver;
+    private SessionConfig sessionConfig;
 
     public Neo4JSpace(String space, NBConfiguration cfg) {
         this.space = space;
@@ -40,6 +43,10 @@ public class Neo4JSpace implements AutoCloseable {
     }
 
     private Driver initializeDriver(NBConfiguration cfg) {
+        SessionConfig.Builder builder = SessionConfig.builder();
+        cfg.getOptional("database").ifPresent(builder::withDatabase);
+        this.sessionConfig = builder.build();
+
         String dbURI = cfg.get("db_uri");
         Optional<String> usernameOpt = cfg.getOptional("username");
         Optional<String> passwordOpt = cfg.getOptional("password");
@@ -84,6 +91,7 @@ public class Neo4JSpace implements AutoCloseable {
             .add(Param.required("db_uri", String.class))
             .add(Param.optional("username", String.class))
             .add(Param.optional("password", String.class))
+            .add(Param.optional("database", String.class))
             .asReadOnly();
     }
 
@@ -91,6 +99,14 @@ public class Neo4JSpace implements AutoCloseable {
         return driver;
     }
 
+    public AsyncSession newAsyncSession() {
+        return driver.session(AsyncSession.class,sessionConfig);
+    }
+
+    public Session newSession() {
+        return driver.session(sessionConfig);
+    }
+
     @Override
     public void close() throws Exception {
         if (driver != null){

nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JAsyncAutoCommitOpDispenser.java

@@ -35,7 +35,7 @@ public class Neo4JAsyncAutoCommitOpDispenser extends Neo4JBaseOpDispenser {
     @Override
     public LongFunction<Neo4JAsyncAutoCommitOp> createOpFunc() {
         return l -> new Neo4JAsyncAutoCommitOp(
-            spaceFunc.apply(l).getDriver().session(AsyncSession.class),
+            asyncSessionFunc.apply(l),
             queryFunc.apply(l)
         );
     }

nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JAsyncReadTxnOpDispenser.java

@@ -34,7 +34,7 @@ public class Neo4JAsyncReadTxnOpDispenser extends Neo4JBaseOpDispenser {
     @Override
     public LongFunction<Neo4JAsyncReadTxnOp> createOpFunc() {
         return l -> new Neo4JAsyncReadTxnOp(
-            spaceFunc.apply(l).getDriver().session(AsyncSession.class),
+            asyncSessionFunc.apply(l),
             queryFunc.apply(l)
         );
     }

nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JAsyncWriteTxnOpDispenser.java

@@ -35,7 +35,7 @@ public class Neo4JAsyncWriteTxnOpDispenser extends Neo4JBaseOpDispenser {
     @Override
     public LongFunction<Neo4JAsyncWriteTxnOp> createOpFunc() {
         return l -> new Neo4JAsyncWriteTxnOp(
-            spaceFunc.apply(l).getDriver().session(AsyncSession.class),
+            asyncSessionFunc.apply(l),
             queryFunc.apply(l)
         );
     }

nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JBaseOpDispenser.java

@@ -23,6 +23,9 @@ import io.nosqlbench.adapters.api.activityimpl.BaseOpDispenser;
 import io.nosqlbench.adapters.api.templating.ParsedOp;
 
 import org.neo4j.driver.Query;
+import org.neo4j.driver.Session;
+import org.neo4j.driver.SessionConfig;
+import org.neo4j.driver.async.AsyncSession;
 
 import java.util.Collections;
 import java.util.function.LongFunction;
@@ -35,6 +38,8 @@ public abstract class Neo4JBaseOpDispenser extends BaseOpDispenser<Neo4JBaseOp,
     protected final LongFunction<Query> queryFunc;
     protected final LongFunction<Map> paramFunc;
     protected final LongFunction<Neo4JBaseOp> opFunc;
+    protected final LongFunction<Session> sessionFunc;
+    protected final LongFunction<AsyncSession> asyncSessionFunc;
 
     public Neo4JBaseOpDispenser(Neo4JDriverAdapter adapter, ParsedOp op, LongFunction<Neo4JSpace> spaceFunc, String requiredTemplateKey) {
         super(adapter, op);
@@ -43,6 +48,8 @@ public abstract class Neo4JBaseOpDispenser extends BaseOpDispenser<Neo4JBaseOp,
         this.paramFunc = createParamFunc(op);
         this.queryFunc = createQueryFunc();
         this.opFunc = (LongFunction<Neo4JBaseOp>) createOpFunc();
+        this.sessionFunc = getSessionFunction(spaceFunc,op);
+        this.asyncSessionFunc = getAsyncSessionFunction(spaceFunc,op);
     }
 
     private LongFunction<Map> createParamFunc(ParsedOp op) {
@@ -50,6 +57,24 @@ public abstract class Neo4JBaseOpDispenser extends BaseOpDispenser<Neo4JBaseOp,
             .orElse(l -> Collections.emptyMap());
     }
 
+    private LongFunction<Session> getSessionFunction(LongFunction<Neo4JSpace> spaceFunc, ParsedOp op) {
+        LongFunction<SessionConfig.Builder> scbF = (long l) -> SessionConfig.builder();
+        scbF = op.enhanceFuncOptionally(scbF,"database",String.class,SessionConfig.Builder::withDatabase);
+        LongFunction<SessionConfig.Builder> finalScbF = scbF;
+        LongFunction<SessionConfig> scF = (long l) -> finalScbF.apply(l).build();
+        return (long l) -> spaceFunc.apply(l).getDriver().session(Session.class,scF.apply(l));
+
+    }
+
+    private LongFunction<AsyncSession> getAsyncSessionFunction(LongFunction<Neo4JSpace> spaceFunc, ParsedOp op) {
+        LongFunction<SessionConfig.Builder> scbF = (long l) -> SessionConfig.builder();
+        scbF = op.enhanceFuncOptionally(scbF,"database",String.class,SessionConfig.Builder::withDatabase);
+        LongFunction<SessionConfig.Builder> finalScbF = scbF;
+        LongFunction<SessionConfig> scF = (long l) -> finalScbF.apply(l).build();
+        return (long l) -> spaceFunc.apply(l).getDriver().session(AsyncSession.class,scF.apply(l));
+
+    }
+
     /**
      * Reference:
      * - https://neo4j.com/docs/api/java-driver/current/org.neo4j.driver/org/neo4j/driver/Query.html#withParameters(java.util.Map)

nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JSyncAutoCommitOpDispenser.java

@@ -35,7 +35,7 @@ public class Neo4JSyncAutoCommitOpDispenser extends Neo4JBaseOpDispenser {
     @Override
     public LongFunction<Neo4JSyncAutoCommitOp> createOpFunc() {
         return l -> new Neo4JSyncAutoCommitOp(
-            spaceFunc.apply(l).getDriver().session(Session.class),
+            sessionFunc.apply(l),
             queryFunc.apply(l)
         );
     }

nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JSyncReadTxnOpDispenser.java

@@ -21,8 +21,6 @@ import io.nosqlbench.adapter.neo4j.ops.Neo4JSyncReadTxnOp;
 import io.nosqlbench.adapter.neo4j.Neo4JSpace;
 import io.nosqlbench.adapters.api.templating.ParsedOp;
 
-import org.neo4j.driver.Session;
-
 import java.util.function.LongFunction;
 
 
@@ -35,7 +33,7 @@ public class Neo4JSyncReadTxnOpDispenser extends Neo4JBaseOpDispenser {
     @Override
     public LongFunction<Neo4JSyncReadTxnOp> createOpFunc() {
         return l -> new Neo4JSyncReadTxnOp(
-            spaceFunc.apply(l).getDriver().session(Session.class),
+            sessionFunc.apply(l),
             queryFunc.apply(l)
         );
     }

nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JSyncWriteTxnOpDispenser.java

@@ -28,14 +28,19 @@ import java.util.function.LongFunction;
 
 public class Neo4JSyncWriteTxnOpDispenser extends Neo4JBaseOpDispenser {
 
-    public Neo4JSyncWriteTxnOpDispenser(Neo4JDriverAdapter adapter, ParsedOp op, LongFunction<Neo4JSpace> spaceFunc, String requiredTemplateKey) {
+    public Neo4JSyncWriteTxnOpDispenser(
+        Neo4JDriverAdapter adapter,
+        ParsedOp op,
+        LongFunction<Neo4JSpace> spaceFunc,
+        String requiredTemplateKey
+    ) {
         super(adapter, op, spaceFunc, requiredTemplateKey);
     }
 
     @Override
     public LongFunction<Neo4JSyncWriteTxnOp> createOpFunc() {
         return l -> new Neo4JSyncWriteTxnOp(
-            spaceFunc.apply(l).getDriver().session(Session.class),
+            sessionFunc.apply(l),
             queryFunc.apply(l)
         );
     }