Merge pull request #627 from vitessio/admin-pr-exec

Allow Admin to add PR to execution queue for benchmarking and profiling
vitessio · Nov 19, 2024 · a687fc2 · a687fc2
2 parents c7dfb8c + 5aaebae
commit a687fc2
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 9 deletions.
diff --git a/go/admin/api.go b/go/admin/api.go
@@ -58,6 +58,7 @@ type (
 		Auth               string   `json:"auth"`
 		Source             string   `json:"source"`
 		SHA                string   `json:"sha"`
+		PR                 string   `json:"pr"`
 		Workloads          []string `json:"workloads"`
 		NumberOfExecutions string   `json:"number_of_executions"`
 		EnableProfile      bool     `json:"enable_profile"`
@@ -227,15 +228,21 @@ func (a *Admin) handleExecutionsAdd(c *gin.Context) {
 	requestPayload := executionRequest{
 		Source:             c.PostForm("source"),
 		SHA:                c.PostForm("sha"),
+		PR:                 c.PostForm("pr"),
 		Workloads:          c.PostFormArray("workloads"),
 		NumberOfExecutions: c.PostForm("numberOfExecutions"),
 		EnableProfile:      c.PostForm("enableProfiling") != "",
 		BinaryToProfile:    c.PostForm("binaryToProfile"),
 		ProfileMode:        c.PostForm("profileMode"),
 	}
 
-	if requestPayload.Source == "" || requestPayload.SHA == "" || len(requestPayload.Workloads) == 0 || requestPayload.NumberOfExecutions == "" {
-		c.JSON(http.StatusBadRequest, gin.H{"error": "Missing required fields: Source, SHA, workflows, numberOfExecutions"})
+	if requestPayload.Source == "" || len(requestPayload.Workloads) == 0 || requestPayload.NumberOfExecutions == "" {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "Missing required fields: Source, workflows, numberOfExecutions"})
+		return
+	}
+
+	if requestPayload.SHA == "" && requestPayload.PR == "" {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "Please provide either a SHA or a PR number"})
 		return
 	}
 

diff --git a/go/admin/templates/add_new_executions.html b/go/admin/templates/add_new_executions.html
@@ -20,6 +20,15 @@
         name="sha"
       />
     </div>
+    <div>
+      <label class="label">PR Number</label>
+      <input
+        class="shadow appearance-none border rounded w-full py-2 px-3 text-gray-700 leading-tight focus:outline-none focus:shadow-outline"
+        type="number"
+        class="form-control"
+        name="pr"
+      />
+    </div>
     <br>
     <label class="text-lg font-semibold mt-4">Workloads</label>
     <div class="flex flex-row gap-4">

diff --git a/go/server/api.go b/go/server/api.go
@@ -558,6 +558,7 @@ type ExecutionRequest struct {
 	Auth               string   `json:"auth"`
 	Source             string   `json:"source"`
 	SHA                string   `json:"sha"`
+	PR                 string   `json:"pr"`
 	Workloads          []string `json:"workloads"`
 	NumberOfExecutions string   `json:"number_of_executions"`
 	EnableProfile      bool     `json:"enable_profile"`
@@ -578,11 +579,31 @@ func (s *Server) addExecutions(c *gin.Context) {
 		return
 	}
 
-	if req.Source == "" || req.SHA == "" || len(req.Workloads) == 0 || req.NumberOfExecutions == "" {
+	if req.Source == "" || len(req.Workloads) == 0 || req.NumberOfExecutions == "" {
 		c.JSON(http.StatusBadRequest, &ErrorAPI{Error: "missing argument"})
 		return
 	}
 
+	var pr int
+	if req.PR != "" {
+		var err error
+		pr, err = strconv.Atoi(req.PR)
+		if err != nil {
+			c.JSON(http.StatusBadRequest, &ErrorAPI{Error: "PR must be an integer"})
+			return
+		}
+		url := "https://api.github.com/repos/vitessio/vitess/pulls/" + req.PR
+		prInfo, err := git.GetPullRequestHeadAndBase(url)
+		if err != nil {
+			c.JSON(http.StatusBadRequest, &ErrorAPI{Error: "unable to get PR information"})
+		}
+		req.SHA = prInfo.SHA
+	}
+
+	if pr == 0 && req.SHA == "" {
+		c.JSON(http.StatusBadRequest, &ErrorAPI{Error: "missing argument PR or SHA"})
+	}
+
 	if len(req.Workloads) == 1 && req.Workloads[0] == "all" {
 		req.Workloads = s.workloads
 	}
@@ -603,7 +624,7 @@ func (s *Server) addExecutions(c *gin.Context) {
 
 	for _, workload := range req.Workloads {
 		for i := 0; i < execs; i++ {
-			elem := s.createSimpleExecutionQueueElement(s.benchmarkConfig[strings.ToLower(workload)], req.Source, req.SHA, workload, string(macrobench.Gen4Planner), false, 0, git.Version{}, profileInformation)
+			elem := s.createSimpleExecutionQueueElement(s.benchmarkConfig[strings.ToLower(workload)], req.Source, req.SHA, workload, string(macrobench.Gen4Planner), false, pr, git.Version{}, profileInformation)
 			elem.identifier.UUID = uuid.NewString()
 			newElements = append(newElements, elem)
 		}

diff --git a/go/tools/git/github.go b/go/tools/git/github.go
@@ -33,7 +33,7 @@ func GetPullRequestsFromGitHub(labels []string, repo string) ([]PRInfo, error) {
 
 	prInfos := []PRInfo{}
 	for _, pull := range pulls {
-		prInfo, err := getPullRequestHeadAndBase(pull)
+		prInfo, err := GetPullRequestHeadAndBase(pull)
 		if err != nil {
 			return nil, err
 		}

diff --git a/go/tools/git/pulls.go b/go/tools/git/pulls.go
@@ -26,11 +26,11 @@ import (
 
 type PRInfo struct {
 	Number int
-	Base string
-	SHA  string
+	Base   string
+	SHA    string
 }
 
-func getPullRequestHeadAndBase(url string) (PRInfo, error) {
+func GetPullRequestHeadAndBase(url string) (PRInfo, error) {
 	client := http.Client{}
 	request, err := http.NewRequest(http.MethodGet, url, nil)
 	if err != nil {
@@ -44,7 +44,7 @@ func getPullRequestHeadAndBase(url string) (PRInfo, error) {
 
 	res := struct {
 		Number int `json:"number"`
-		Head struct {
+		Head   struct {
 			SHA string `json:"sha"`
 		} `json:"head"`
 		Base struct {