From c00dec63ec766bcc48941edfc3b02b0999ab5250 Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Tue, 9 Jul 2024 09:54:23 -0500 Subject: [PATCH] update session initializers on neo4j --- .../nosqlbench/adapter/neo4j/Neo4JSpace.java | 16 ++++++++++++ .../Neo4JAsyncAutoCommitOpDispenser.java | 2 +- .../Neo4JAsyncReadTxnOpDispenser.java | 2 +- .../Neo4JAsyncWriteTxnOpDispenser.java | 2 +- .../opdispensers/Neo4JBaseOpDispenser.java | 25 +++++++++++++++++++ .../Neo4JSyncAutoCommitOpDispenser.java | 2 +- .../Neo4JSyncReadTxnOpDispenser.java | 4 +-- .../Neo4JSyncWriteTxnOpDispenser.java | 9 +++++-- 8 files changed, 53 insertions(+), 9 deletions(-) diff --git a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/Neo4JSpace.java b/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/Neo4JSpace.java index 82e06a473..5f81a917d 100644 --- a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/Neo4JSpace.java +++ b/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 usernameOpt = cfg.getOptional("username"); Optional 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){ diff --git a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JAsyncAutoCommitOpDispenser.java b/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JAsyncAutoCommitOpDispenser.java index 4d1079be7..73b5014cb 100644 --- a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JAsyncAutoCommitOpDispenser.java +++ b/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 createOpFunc() { return l -> new Neo4JAsyncAutoCommitOp( - spaceFunc.apply(l).getDriver().session(AsyncSession.class), + asyncSessionFunc.apply(l), queryFunc.apply(l) ); } diff --git a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JAsyncReadTxnOpDispenser.java b/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JAsyncReadTxnOpDispenser.java index a8eb63036..867fc8c62 100644 --- a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JAsyncReadTxnOpDispenser.java +++ b/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 createOpFunc() { return l -> new Neo4JAsyncReadTxnOp( - spaceFunc.apply(l).getDriver().session(AsyncSession.class), + asyncSessionFunc.apply(l), queryFunc.apply(l) ); } diff --git a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JAsyncWriteTxnOpDispenser.java b/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JAsyncWriteTxnOpDispenser.java index 5bc5dc5da..ddf90a253 100644 --- a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JAsyncWriteTxnOpDispenser.java +++ b/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 createOpFunc() { return l -> new Neo4JAsyncWriteTxnOp( - spaceFunc.apply(l).getDriver().session(AsyncSession.class), + asyncSessionFunc.apply(l), queryFunc.apply(l) ); } diff --git a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JBaseOpDispenser.java b/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JBaseOpDispenser.java index ce852c886..c70ea74c9 100644 --- a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JBaseOpDispenser.java +++ b/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 queryFunc; protected final LongFunction paramFunc; protected final LongFunction opFunc; + protected final LongFunction sessionFunc; + protected final LongFunction asyncSessionFunc; public Neo4JBaseOpDispenser(Neo4JDriverAdapter adapter, ParsedOp op, LongFunction spaceFunc, String requiredTemplateKey) { super(adapter, op); @@ -43,6 +48,8 @@ public abstract class Neo4JBaseOpDispenser extends BaseOpDispenser) createOpFunc(); + this.sessionFunc = getSessionFunction(spaceFunc,op); + this.asyncSessionFunc = getAsyncSessionFunction(spaceFunc,op); } private LongFunction createParamFunc(ParsedOp op) { @@ -50,6 +57,24 @@ public abstract class Neo4JBaseOpDispenser extends BaseOpDispenser Collections.emptyMap()); } + private LongFunction getSessionFunction(LongFunction spaceFunc, ParsedOp op) { + LongFunction scbF = (long l) -> SessionConfig.builder(); + scbF = op.enhanceFuncOptionally(scbF,"database",String.class,SessionConfig.Builder::withDatabase); + LongFunction finalScbF = scbF; + LongFunction scF = (long l) -> finalScbF.apply(l).build(); + return (long l) -> spaceFunc.apply(l).getDriver().session(Session.class,scF.apply(l)); + + } + + private LongFunction getAsyncSessionFunction(LongFunction spaceFunc, ParsedOp op) { + LongFunction scbF = (long l) -> SessionConfig.builder(); + scbF = op.enhanceFuncOptionally(scbF,"database",String.class,SessionConfig.Builder::withDatabase); + LongFunction finalScbF = scbF; + LongFunction 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) diff --git a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JSyncAutoCommitOpDispenser.java b/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JSyncAutoCommitOpDispenser.java index bd08aa88d..36e817c27 100644 --- a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JSyncAutoCommitOpDispenser.java +++ b/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 createOpFunc() { return l -> new Neo4JSyncAutoCommitOp( - spaceFunc.apply(l).getDriver().session(Session.class), + sessionFunc.apply(l), queryFunc.apply(l) ); } diff --git a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JSyncReadTxnOpDispenser.java b/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JSyncReadTxnOpDispenser.java index b1b170566..d4939dee4 100644 --- a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JSyncReadTxnOpDispenser.java +++ b/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 createOpFunc() { return l -> new Neo4JSyncReadTxnOp( - spaceFunc.apply(l).getDriver().session(Session.class), + sessionFunc.apply(l), queryFunc.apply(l) ); } diff --git a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JSyncWriteTxnOpDispenser.java b/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JSyncWriteTxnOpDispenser.java index 7c7d095aa..53233b0f6 100644 --- a/nb-adapters/adapter-neo4j/src/main/java/io/nosqlbench/adapter/neo4j/opdispensers/Neo4JSyncWriteTxnOpDispenser.java +++ b/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 spaceFunc, String requiredTemplateKey) { + public Neo4JSyncWriteTxnOpDispenser( + Neo4JDriverAdapter adapter, + ParsedOp op, + LongFunction spaceFunc, + String requiredTemplateKey + ) { super(adapter, op, spaceFunc, requiredTemplateKey); } @Override public LongFunction createOpFunc() { return l -> new Neo4JSyncWriteTxnOp( - spaceFunc.apply(l).getDriver().session(Session.class), + sessionFunc.apply(l), queryFunc.apply(l) ); }