Skip to content

Commit

Permalink
Add pre query cycle script execution hook
Browse files Browse the repository at this point in the history
  • Loading branch information
agrawalreetika committed Dec 30, 2024
1 parent 8e0fac7 commit 4c9dc6f
Show file tree
Hide file tree
Showing 5 changed files with 25 additions and 2 deletions.
11 changes: 11 additions & 0 deletions benchmarks/test/my_pre_query_cycle_script.py
Original file line number Diff line number Diff line change
@@ -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)
4 changes: 4 additions & 0 deletions benchmarks/test/stage_4.json
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,10 @@
"echo \"run this script after this stage is complete\"",
"python3 my_post_stage_script.py count.txt"
],
"pre_query_cycle_scripts": [
"echo \"execute this script before starting all runs of the same query in this stage\"",
"python3 my_pre_query_cycle_script.py count.txt"
],
"post_query_scripts": [
"echo \"run this script after each query in this stage is complete\"",
"python3 my_post_query_script.py count.txt"
Expand Down
7 changes: 7 additions & 0 deletions stage/stage.go
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,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 before starting query cycle runs of each query.
PreQueryCycleShellScripts []string `json:"pre_query_cycle_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
Expand Down Expand Up @@ -408,6 +410,11 @@ func (s *Stage) runShellScripts(ctx context.Context, shellScripts []string) erro
func (s *Stage) runQueries(ctx context.Context, queries []string, queryFile *string, expectedRowCountStartIndex int) (retErr error) {
batchSize := len(queries)
for i, queryText := range queries {
// run pre query cycle shell scripts
preQueryCycleErr := s.runShellScripts(ctx, s.PreQueryCycleShellScripts)
if preQueryCycleErr != nil {
return fmt.Errorf("pre-query script execution failed: %w", preQueryCycleErr)
}
for j := 0; j < s.ColdRuns+s.WarmRuns; j++ {
query := &Query{
Text: queryText,
Expand Down
4 changes: 2 additions & 2 deletions stage/stage_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"}, 6)
"SYNTAX_ERROR: Table tpch.sf1.foo does not exist"}, 9)
})
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"}, 9)
"SYNTAX_ERROR: line 1:11: Function sum1 not registered"}, 13)
})
}

Expand Down
1 change: 1 addition & 0 deletions stage/stage_utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ func (s *Stage) MergeWith(other *Stage) *Stage {
s.PreStageShellScripts = append(s.PreStageShellScripts, other.PreStageShellScripts...)
s.PostQueryShellScripts = append(s.PostQueryShellScripts, other.PostQueryShellScripts...)
s.PostStageShellScripts = append(s.PostStageShellScripts, other.PostStageShellScripts...)
s.PreQueryCycleShellScripts = append(s.PreQueryCycleShellScripts, other.PreQueryCycleShellScripts...)
s.PostQueryCycleShellScripts = append(s.PostQueryCycleShellScripts, other.PostQueryCycleShellScripts...)

return s
Expand Down

0 comments on commit 4c9dc6f

Please sign in to comment.