Bolt Sampler - Neo4j 4.x and cluster features

src/protocol/bolt/src/main/java/org/apache/jmeter/protocol/bolt/config/BoltConnectionElement.java

@@ -26,11 +26,13 @@
 import org.apache.jmeter.threads.JMeterContextService;
 import org.apache.jmeter.threads.JMeterVariables;
 import org.neo4j.driver.AuthTokens;
+import org.neo4j.driver.Config;
 import org.neo4j.driver.Driver;
 import org.neo4j.driver.GraphDatabase;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+
 @TestElementMetadata(labelResource = "displayName")
 public class BoltConnectionElement extends AbstractTestElement
         implements ConfigElement, TestStateListener, TestBean {
@@ -39,6 +41,7 @@ public class BoltConnectionElement extends AbstractTestElement
     private String boltUri;
     private String username;
     private String password;
+    private int maxConnectionPoolSize;
     private Driver driver;
 
     public static final String BOLT_CONNECTION = "boltConnection";
@@ -65,7 +68,10 @@ public void testStarted() {
             log.error("Bolt connection already exists");
         } else {
             synchronized (this) {
-                driver = GraphDatabase.driver(getBoltUri(), AuthTokens.basic(getUsername(), getPassword()));
+                Config config = Config.builder()
+                        .withMaxConnectionPoolSize( getMaxConnectionPoolSize() )
+                        .build();
+                driver = GraphDatabase.driver(getBoltUri(), AuthTokens.basic(getUsername(), getPassword()), config);
                 variables.putObject(BOLT_CONNECTION, driver);
             }
         }
@@ -100,6 +106,14 @@ public void setBoltUri(String boltUri) {
         this.boltUri = boltUri;
     }
 
+    public int getMaxConnectionPoolSize() {
+        return maxConnectionPoolSize;
+    }
+
+    public void setMaxConnectionPoolSize(int maxConnectionPoolSize) {
+        this.maxConnectionPoolSize = maxConnectionPoolSize;
+    }
+
     public String getUsername() {
         return username;
     }

src/protocol/bolt/src/main/java/org/apache/jmeter/protocol/bolt/config/BoltConnectionElementBeanInfo.java

@@ -33,7 +33,7 @@ public class BoltConnectionElementBeanInfo extends BeanInfoSupport {
     public BoltConnectionElementBeanInfo() {
         super(BoltConnectionElement.class);
 
-        createPropertyGroup("connection", new String[] { "boltUri", "username", "password" });
+        createPropertyGroup("connection", new String[] { "boltUri", "username", "password", "maxConnectionPoolSize" });
 
         PropertyDescriptor propertyDescriptor =  property("boltUri");
         propertyDescriptor.setValue(NOT_UNDEFINED, Boolean.TRUE);
@@ -44,6 +44,9 @@ public BoltConnectionElementBeanInfo() {
         propertyDescriptor = property("password", TypeEditor.PasswordEditor);
         propertyDescriptor.setValue(NOT_UNDEFINED, Boolean.TRUE);
         propertyDescriptor.setValue(DEFAULT, "");
+        propertyDescriptor = property("maxConnectionPoolSize");
+        propertyDescriptor.setValue(NOT_UNDEFINED, Boolean.TRUE);
+        propertyDescriptor.setValue(DEFAULT, 100);
 
         if(log.isDebugEnabled()) {
             String descriptorsAsString = Arrays.stream(getPropertyDescriptors())

src/protocol/bolt/src/main/java/org/apache/jmeter/protocol/bolt/sampler/AbstractBoltTestElement.java

@@ -17,13 +17,52 @@
 
 package org.apache.jmeter.protocol.bolt.sampler;
 
+import java.time.Duration;
+
+import org.apache.commons.lang3.EnumUtils;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.jmeter.testelement.AbstractTestElement;
+import org.neo4j.driver.AccessMode;
+import org.neo4j.driver.SessionConfig;
+import org.neo4j.driver.TransactionConfig;
 
 public abstract class AbstractBoltTestElement extends AbstractTestElement {
 
     private String cypher;
     private String params;
+    private String database;
+    private String accessMode;
     private boolean recordQueryResults;
+    private int txTimeout;
+
+    public int getTxTimeout() {
+        return txTimeout;
+    }
+
+    public void setTxTimeout(int txTimeout) {
+        this.txTimeout = txTimeout;
+    }
+
+    public String getAccessMode() {
+        if (accessMode == null) {
+            return AccessMode.WRITE.toString();
+        }
+        return accessMode;
+    }
+
+    public void setAccessMode(String accessMode) {
+        if (EnumUtils.isValidEnum(AccessMode.class, accessMode)) {
+            this.accessMode = accessMode;
+        }
+    }
+
+    public String getDatabase() {
+        return database;
+    }
+
+    public void setDatabase(String database) {
+        this.database = database;
+    }
 
     public String getCypher() {
         return cypher;
@@ -48,4 +87,27 @@ public boolean isRecordQueryResults() {
     public void setRecordQueryResults(boolean recordQueryResults) {
         this.recordQueryResults = recordQueryResults;
     }
+
+    //returns a SessionConfig object that can be passed to the driver session
+    public SessionConfig getSessionConfig() {
+        SessionConfig.Builder sessionConfigBuilder = SessionConfig.builder()
+                .withDefaultAccessMode(Enum.valueOf(AccessMode.class, getAccessMode()));
+
+        if (StringUtils.isNotBlank(database)) {
+            sessionConfigBuilder.withDatabase(database);
+        }
+
+        return sessionConfigBuilder.build();
+    }
+
+    //returns a TransactionConfig object that can be passed to the driver transaction
+    public TransactionConfig getTransactionConfig() {
+        TransactionConfig.Builder txConfigBuilder = TransactionConfig.builder();
+
+        if (txTimeout > 0) {
+            txConfigBuilder.withTimeout(Duration.ofSeconds(txTimeout));
+        }
+
+        return txConfigBuilder.build();
+    }
 }

src/protocol/bolt/src/main/java/org/apache/jmeter/protocol/bolt/sampler/BoltSampler.java

@@ -22,6 +22,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
@@ -39,6 +40,8 @@
 import org.neo4j.driver.Record;
 import org.neo4j.driver.Result;
 import org.neo4j.driver.Session;
+import org.neo4j.driver.SessionConfig;
+import org.neo4j.driver.TransactionConfig;
 import org.neo4j.driver.exceptions.Neo4jException;
 import org.neo4j.driver.summary.ResultSummary;
 
@@ -82,7 +85,14 @@ public SampleResult sample(Entry e) {
 
         try {
             res.setResponseHeaders("Cypher request: " + getCypher());
-            res.setResponseData(execute(BoltConnectionElement.getDriver(), getCypher(), params), StandardCharsets.UTF_8.name());
+            res.setResponseData(
+                    execute(
+                        BoltConnectionElement.getDriver(),
+                        getCypher(),
+                        params,
+                        getSessionConfig(),
+                        getTransactionConfig()),
+                    StandardCharsets.UTF_8.name());
         } catch (Exception ex) {
             res = handleException(res, ex);
         } finally {
@@ -100,9 +110,10 @@ public boolean applies(ConfigTestElement configElement) {
         return APPLICABLE_CONFIG_CLASSES.contains(guiClass);
     }
 
-    private String execute(Driver driver, String cypher, Map<String, Object> params) {
-        try (Session session = driver.session()) {
-            Result statementResult = session.run(cypher, params);
+    private String execute(Driver driver, String cypher, Map<String, Object> params,
+                           SessionConfig sessionConfig, TransactionConfig txConfig) {
+        try (Session session = driver.session(sessionConfig)) {
+            Result statementResult = session.run(cypher, params, txConfig);
             return response(statementResult);
         }
     }
@@ -135,12 +146,25 @@ private String request() {
                 .append(getCypher())
                 .append("\n")
                 .append("Parameters: \n")
-                .append(getParams());
+                .append(getParams())
+                .append("\n")
+                .append("Database: \n")
+                .append(getDatabase())
+                .append("\n")
+                .append("Access Mode: \n")
+                .append(getAccessMode());
         return request.toString();
     }
 
     private String response(Result result) {
         StringBuilder response = new StringBuilder();
+        List<Record> records;
+        if (isRecordQueryResults()) {
+            //get records already as consume() will exhaust the stream
+            records = result.list();
+        } else {
+            records = Collections.emptyList();
+        }
         response.append("\nSummary:");
         ResultSummary summary = result.consume();
         response.append("\nConstraints Added: ")
@@ -167,7 +191,7 @@ private String response(Result result) {
                 .append(summary.counters().relationshipsDeleted());
         response.append("\n\nRecords: ");
         if (isRecordQueryResults()) {
-            for (Record record : result.list()) {
+            for (Record record : records) {
                 response.append("\n").append(record);
             }
         } else {

src/protocol/bolt/src/main/java/org/apache/jmeter/protocol/bolt/sampler/BoltTestElementBeanInfoSupport.java

@@ -18,10 +18,12 @@
 package org.apache.jmeter.protocol.bolt.sampler;
 
 import java.beans.PropertyDescriptor;
+import java.util.Arrays;
 
 import org.apache.jmeter.testbeans.BeanInfoSupport;
 import org.apache.jmeter.testbeans.TestBean;
 import org.apache.jmeter.testbeans.gui.TypeEditor;
+import org.neo4j.driver.AccessMode;
 
 public abstract class BoltTestElementBeanInfoSupport extends BeanInfoSupport {
     /**
@@ -33,17 +35,35 @@ protected BoltTestElementBeanInfoSupport(Class<? extends TestBean> beanClass) {
         super(beanClass);
 
         createPropertyGroup("query", new String[] { "cypher","params","recordQueryResults"});
+        createPropertyGroup("options", new String[] { "accessMode","database", "txTimeout"});
 
-        PropertyDescriptor propertyDescriptor =  property("cypher", TypeEditor.TextAreaEditor);
+        PropertyDescriptor propertyDescriptor = property("cypher", TypeEditor.TextAreaEditor);
         propertyDescriptor.setValue(NOT_UNDEFINED, Boolean.TRUE);
         propertyDescriptor.setValue(DEFAULT, "");
 
-        propertyDescriptor =  property("params", TypeEditor.TextAreaEditor);
+        propertyDescriptor = property("params", TypeEditor.TextAreaEditor);
         propertyDescriptor.setValue(NOT_UNDEFINED, Boolean.TRUE);
         propertyDescriptor.setValue(DEFAULT, "{\"paramName\":\"paramValue\"}");
 
-        propertyDescriptor =  property("recordQueryResults");
+        propertyDescriptor = property("recordQueryResults");
         propertyDescriptor.setValue(NOT_UNDEFINED, Boolean.TRUE);
         propertyDescriptor.setValue(DEFAULT, Boolean.FALSE);
+
+        propertyDescriptor = property("accessMode", TypeEditor.ComboStringEditor);
+        propertyDescriptor.setValue(NOT_UNDEFINED, Boolean.TRUE);
+        propertyDescriptor.setValue(NOT_EXPRESSION, Boolean.TRUE);
+        propertyDescriptor.setValue(DEFAULT, AccessMode.WRITE.toString());
+        propertyDescriptor.setValue(TAGS, getListAccessModes());
+
+        propertyDescriptor = property("database", TypeEditor.ComboStringEditor);
+        propertyDescriptor.setValue(DEFAULT, "neo4j");
+
+        propertyDescriptor = property("txTimeout");
+        propertyDescriptor.setValue(NOT_UNDEFINED, Boolean.TRUE);
+        propertyDescriptor.setValue(DEFAULT, 60);
+    }
+
+    private String[] getListAccessModes() {
+        return Arrays.stream(AccessMode.values()).map(Enum::toString).toArray(String[]::new);
     }
 }

src/protocol/bolt/src/main/resources/org/apache/jmeter/protocol/bolt/config/BoltConnectionElementResources.properties

@@ -18,8 +18,10 @@
 displayName=Bolt Connection Configuration
 connection.displayName=Bolt Configuration
 boltUri.displayName=Bolt URI
-boltUri.shortDescription=Bolt URI
+boltUri.shortDescription=Bolt URI<br>for a direct connection: bolt://&lt;HOST&gt;:&lt;PORT&gt;<br>for a cluster: neo4j://&lt;HOST&gt;:&lt;PORT&gt;[?&lt;ROUTING_CONTEXT&gt;])
 username.displayName=Username
 username.shortDescription=Username
 password.displayName=Password
 password.shortDescription=Password
+maxConnectionPoolSize.displayName=Connection Pool Max Size
+maxConnectionPoolSize.shortDescription=Size limit for the pool of Bolt connections

src/protocol/bolt/src/main/resources/org/apache/jmeter/protocol/bolt/sampler/BoltSamplerResources.properties

@@ -17,9 +17,16 @@
 
 displayName=Bolt Request
 query.displayName=Query
+options.displayName=Options
 cypher.displayName=Cypher Statement
 cypher.shortDescription=Cypher Statement
 params.displayName=Params
 params.shortDescription=Params
 recordQueryResults.displayName=Record Query Results
 recordQueryResults.shortDescription=Records the results of queries and displays in listeners such as View Results Tree, this iterates through the entire resultset. Use to debug only.
+accessMode.displayName=Access Mode
+accessMode.shortDescription=Whether it's a READ or WRITE query (affects query routing in clusters)
+database.displayName=Database
+database.shortDescription=Neo4j 4.x: database to query (leave empty for 3.5)
+txTimeout.displayName=Transaction timeout
+txTimeout.shortDescription=Transaction timeout in seconds