IntenseFinance/nosqlbench
3
0
Fork 0
mirror of https://github.com/nosqlbench/nosqlbench.git synced 2024-12-26 08:41:05 -06:00
Code Issues Packages Projects Releases Wiki Activity

update session initializers on neo4j

Browse Source
This commit is contained in:
Jonathan Shook 2024-07-09 09:54:23 -05:00
parent b7e0f35cd2
commit c00dec63ec
8 changed files with 53 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -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){

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

@ -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)
);
}

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

@ -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)
);
}

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

@ -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)
);
}

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

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

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

@ -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)
);
}

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

@ -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)
);
}

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

@ -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)
);
}