Add post query cycle script execution hook

benchmarks/test/my_post_query_cycle_script.py

@@ -0,0 +1,11 @@
+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)
+
+    file_path = sys.argv[1]
+    increment_file_value(file_path)

benchmarks/test/stage_4.json

@@ -14,6 +14,10 @@
     "echo \"run this script after each query in this stage is complete\"",
     "python3 my_post_query_script.py count.txt"
   ],
+  "post_query_cycle_scripts": [
+    "echo \"execute this script after all runs of the same query in this stage have completed\"",
+    "python3 my_post_query_cycle_script.py count.txt"
+  ],
   "next": [
     "stage_5.json"
   ],

stage/stage.go

@@ -44,6 +44,8 @@ type Stage struct {
 	PostStageShellScripts []string `json:"post_stage_scripts,omitempty"`
 	// Run shell scripts after executing each query.
 	PostQueryShellScripts []string `json:"post_query_scripts,omitempty"`
+	// Run shell scripts after finishing full query cycle runs each query.
+	PostQueryCycleShellScripts []string `json:"post_query_cycle_scripts,omitempty"`
 	// A map from [catalog.schema] to arrays of integers as expected row counts for all the queries we run
 	// under different schemas. This includes the queries from both Queries and QueryFiles. Queries first and QueryFiles follows.
 	// Can use regexp as key to match multiple [catalog.schema] pairs.
@@ -438,6 +440,11 @@ func (s *Stage) runQueries(ctx context.Context, queries []string, queryFile *str
 			}
 			log.Info().EmbedObject(result).Msgf("query finished")
 		}
+		// run post query cycle shell scripts
+		postQueryCycleErr := s.runShellScripts(ctx, s.PostQueryCycleShellScripts)
+		if retErr == nil {
+			retErr = postQueryCycleErr
+		}
 	}
 	return nil
 }

stage/stage_test.go

@@ -83,12 +83,12 @@ func testParseAndExecute(t *testing.T, abortOnError bool, totalQueryCount int, e
 func TestParseStageGraph(t *testing.T) {
 	t.Run("abortOnError = true", func(t *testing.T) {
 		testParseAndExecute(t, true, 10, 16, []string{
-			"SYNTAX_ERROR: Table tpch.sf1.foo does not exist"}, 3)
+			"SYNTAX_ERROR: Table tpch.sf1.foo does not exist"}, 5)
 	})
 	t.Run("abortOnError = false", func(t *testing.T) {
 		testParseAndExecute(t, false, 15, 24, []string{
 			"SYNTAX_ERROR: Table tpch.sf1.foo does not exist",
-			"SYNTAX_ERROR: line 1:11: Function sum1 not registered"}, 4)
+			"SYNTAX_ERROR: line 1:11: Function sum1 not registered"}, 8)
 	})
 }
 

stage/stage_utils.go

@@ -90,6 +90,7 @@ func (s *Stage) MergeWith(other *Stage) *Stage {
 
 	s.PostQueryShellScripts = append(s.PostQueryShellScripts, other.PostQueryShellScripts...)
 	s.PostStageShellScripts = append(s.PostStageShellScripts, other.PostStageShellScripts...)
+	s.PostQueryCycleShellScripts = append(s.PostQueryCycleShellScripts, other.PostQueryCycleShellScripts...)
 
 	return s
 }