Add pre stage script execution hook

benchmarks/test/my_pre_stage_script.py

@@ -0,0 +1,19 @@
+import sys
+from utils import increment_file_value
+
+# Main function to handle the command-line argument
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print("Missing <file_path>")
+        sys.exit(-1)
+
+    value = -1
+    file_path = sys.argv[1]
+    if os.path.exists(file_path):
+        # If the file exists, read its contents
+        try:
+            with open(file_path, 'r') as file:
+                value = int(file.read().strip())  # Read the integer value from the file
+        except:
+            pass
+    sys.exit(value)

benchmarks/test/stage_4.json

@@ -6,6 +6,10 @@
   "query_files": [
     "stage_4.sql"
   ],
+  "pre_stage_scripts": [
+    "echo \"run this script before this stage is started\"",
+    "python3 my_pre_stage_script.py count.txt"
+  ],
   "post_stage_scripts": [
     "echo \"run this script after this stage is complete\"",
     "python3 my_post_stage_script.py count.txt"

stage/stage.go

@@ -40,6 +40,8 @@ type Stage struct {
 	// If a stage has both Queries and QueryFiles, the queries in the Queries array will be executed first then
 	// the QueryFiles will be read and executed.
 	QueryFiles []string `json:"query_files,omitempty"`
+	// Run shell scripts before starting the execution of queries in a stage.
+	PreStageShellScripts []string `json:"pre_stage_scripts,omitempty"`
 	// Run shell scripts after executing all the queries in a stage.
 	PostStageShellScripts []string `json:"post_stage_scripts,omitempty"`
 	// Run shell scripts after executing each query.
@@ -229,6 +231,10 @@ func (s *Stage) run(ctx context.Context) (returnErr error) {
 	s.setDefaults()
 	s.prepareClient()
 	s.propagateStates()
+	preStageErr := s.runShellScripts(ctx, s.PreStageShellScripts)
+	if returnErr == nil {
+		returnErr = preStageErr
+	}
 	if len(s.Queries)+len(s.QueryFiles) > 0 {
 		if *s.RandomExecution {
 			returnErr = s.runRandomly(ctx)

stage/stage_utils.go

@@ -88,6 +88,7 @@ func (s *Stage) MergeWith(other *Stage) *Stage {
 	s.NextStagePaths = append(s.NextStagePaths, other.NextStagePaths...)
 	s.BaseDir = other.BaseDir
 
+	s.PreStageShellScripts = append(s.PreStageShellScripts, other.PreStageShellScripts...)
 	s.PostQueryShellScripts = append(s.PostQueryShellScripts, other.PostQueryShellScripts...)
 	s.PostStageShellScripts = append(s.PostStageShellScripts, other.PostStageShellScripts...)
 	s.PostQueryCycleShellScripts = append(s.PostQueryCycleShellScripts, other.PostQueryCycleShellScripts...)