11const terminateAfter = require ( './termination-condition' ) ;
22const ProbabilityTable = require ( './stats/ProbabilityTable' ) ;
3+ const _ = require ( 'lodash' ) ;
34
45const usageStr = `
56Usage: run-workload.js -p password
@@ -8,11 +9,13 @@ Usage: run-workload.js -p password
89 [-n hits] how many total queries to run
910 [--ms milliseconds] how many milliseconds to test for
1011 [--workload /path/to/workload.json] probability table spec
12+ [--query CYPHER_QUERY] single cypher query to run
1113 [--concurrency c] how many concurrent queries to run (default: 10)
1214 [--checkpoint cn] how often to print results in milliseconds (default: 5000)
1315 [--fail-fast] if specified, the work will stop after encountering one failure.
1416
1517You may only specify one of the options --n or --ms.
18+ You may only specify one of the options --workload or --query
1619` ;
1720
1821const defaultProbabilityTable = [
@@ -28,21 +31,11 @@ const defaultProbabilityTable = [
2831 [ 1 , 'rawWrite' ] ,
2932] ;
3033
31- const generateFromArgs = ( args ) => {
32- const badlyConfigured = (
33- // User is being inconsistent about when to stop.
34- ( args . n && args . ms ) ||
35- // We don't know where to connect...
36- ( ! process . env . NEO4J_URI && ! args . a ) ||
37- // Don't know what password to use...
38- ( ! process . env . NEO4J_PASSWORD && ! args . p )
39- ) ;
40-
41- if ( badlyConfigured ) {
42- usage ( ) ;
43- }
44-
45- const probabilityTable = args . workload ? require ( args . workload ) : defaultProbabilityTable ;
34+ /**
35+ * @param {* } args a yargs object
36+ * @returns { iterateUntil, runType } of when to stop.
37+ */
38+ const chooseTerminationType = ( args ) => {
4639 let iterateUntil ;
4740 let runType ;
4841
@@ -58,25 +51,81 @@ const generateFromArgs = (args) => {
5851 runType = 'counted' ;
5952 }
6053
61- const p = Number ( args . concurrency ) || Number ( process . env . CONCURRENCY ) ;
62-
63- const failFast = ( 'fail-fast' in args ) ? args [ 'fail-fast' ] : false ;
54+ return { iterateUntil, runType } ;
55+ } ;
6456
57+ /**
58+ * @param {* } args a yargs object
59+ * @returns { username, password, address } of where to connect
60+ */
61+ const chooseConnectionDetails = ( args ) => {
6562 const addressify = str =>
6663 str . indexOf ( '://' ) === - 1 ? `bolt://${ str } ` : str ;
6764
68- const obj = {
65+ return {
6966 username : args . u || process . env . NEO4J_USER || 'neo4j' ,
7067 password : args . p || process . env . NEO4J_PASSWORD ,
7168 address : addressify ( args . a || process . env . NEO4J_URI ) ,
72- probabilityTable : new ProbabilityTable ( probabilityTable ) ,
73- runType,
74- checkpointFreq : args . checkpoint || process . env . CHECKPOINT_FREQUENCY || 5000 ,
69+ } ;
70+ } ;
71+
72+ const chooseConcurrency = ( args ) => {
73+ const p = Number ( args . concurrency ) || Number ( process . env . CONCURRENCY ) ;
74+ return {
7575 concurrency : ( ! Number . isNaN ( p ) && p > 0 ) ? p : 10 ,
76- iterateUntil,
76+ } ;
77+ } ;
78+
79+ const chooseProbabilityTable = ( args ) => {
80+ let ptData = args . workload ? require ( args . workload ) : defaultProbabilityTable ;
81+
82+ if ( args . query ) {
83+ // Always run the custom query.
84+ ptData = [
85+ [ 1.0 , 'custom' ] ,
86+ ] ;
87+ }
88+
89+ const result = {
90+ probabilityTable : new ProbabilityTable ( ptData )
91+ } ;
92+
93+ if ( args . query ) {
94+ result . query = args . query ;
95+ }
96+
97+ return result ;
98+ } ;
99+
100+ const generateFromArgs = ( args ) => {
101+ const badlyConfigured = (
102+ // User is being inconsistent about when to stop.
103+ ( args . n && args . ms ) ||
104+ // Trying to specify to both run a single query and a different workload...
105+ ( args . query && args . workload ) ||
106+ // We don't know where to connect...
107+ ( ! process . env . NEO4J_URI && ! args . a ) ||
108+ // Don't know what password to use...
109+ ( ! process . env . NEO4J_PASSWORD && ! args . p )
110+ ) ;
111+
112+ if ( badlyConfigured ) {
113+ usage ( ) ;
114+ }
115+
116+ const terminationType = chooseTerminationType ( args ) ;
117+ const connectionDetails = chooseConnectionDetails ( args ) ;
118+ const concurrency = chooseConcurrency ( args ) ;
119+ const probabilityTable = chooseProbabilityTable ( args ) ;
120+
121+ const failFast = ( 'fail-fast' in args ) ? args [ 'fail-fast' ] : false ;
122+
123+ // Merge sub-objects.
124+ const obj = _ . merge ( {
125+ checkpointFreq : args . checkpoint || process . env . CHECKPOINT_FREQUENCY || 5000 ,
77126 failFast,
78127 phase : 'NOT_STARTED' ,
79- } ;
128+ } , terminationType , probabilityTable , connectionDetails , concurrency ) ;
80129
81130 if ( obj . runType === 'counted' ) {
82131 obj . n = args . n || 10000 ;
@@ -101,6 +150,7 @@ module.exports = {
101150 . describe ( 'n' , 'number of hits on the database' )
102151 . describe ( 'ms' , 'number of milliseconds to execute' )
103152 . describe ( 'workload' , 'absolute path to JSON probability table/workload' )
153+ . describe ( 'query' , 'Cypher query to run' )
104154 . default ( 'concurrency' , 10 )
105155 . default ( 'checkpoint' , 5000 )
106156 . demandOption ( [ 'p' ] )
0 commit comments