Add functionality to add labels to PRs (#11)

datarobot-oss · Dec 18, 2024 · eca032f · eca032f
1 parent d9f430a
commit eca032f
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/main.go b/main.go
@@ -51,9 +51,16 @@ func main() {
 			os.Getenv("PLUGIN_PR_SOURCE_BRANCH"),
 			os.Getenv("PLUGIN_PR_TARGET_BRANCH"),
 			os.Getenv("PLUGIN_PR_TITLE"),
-			os.Getenv("PLUGIN_PR_BODY"))
+			os.Getenv("PLUGIN_PR_BODY"),
+			os.Getenv("PLUGIN_PR_LABELS"))
 		writeResult(*results, fields)
 	}
+	if strings.Contains(commands, "AddPullRequestLabels") {
+		verifyPluginParameters([]string{"PLUGIN_PR_NUMBER", "PLUGIN_PR_LABELS"})
+		addPullRequestLabels(client, &ctx, repositoryName, repositoryOwner,
+			os.Getenv("PLUGIN_PR_NUMBER"),
+			os.Getenv("PLUGIN_PR_LABELS"))
+	}
 	if strings.Contains(commands, "setStatusCheck") {
 		verifyPluginParameters([]string{"PLUGIN_STATUS_CHECK_SHA",
 			"PLUGIN_STATUS_CHECK_CONTEXT",

diff --git a/pullRequest.go b/pullRequest.go
@@ -10,7 +10,7 @@ import (
 )
 
 func createPullRequest(client *github.Client, ctx *context.Context, repositoryName string, repositoryOwner string,
-	sourceBranch string, targetBranch string, title string, body string) map[string]string {
+	sourceBranch string, targetBranch string, title string, body string, labels string) map[string]string {
 	pr, _, err := client.PullRequests.Create(*ctx, repositoryOwner, repositoryName,
 		&github.NewPullRequest{Title: &title, Body: &body, Head: &sourceBranch, Base: &targetBranch})
 	if err != nil {
@@ -32,6 +32,10 @@ func createPullRequest(client *github.Client, ctx *context.Context, repositoryNa
 			log.Fatal(err)
 		}
 	}
+	if labels != "" {
+		fmt.Println("Adding labels: " + labels)
+		addPullRequestLabels(client, ctx, repositoryName, repositoryOwner, strconv.Itoa(*pr.Number), labels)
+	}
 	fields := map[string]string{
 		"PR_NUMBER":   strconv.Itoa(*pr.Number),
 		"PR_URL":      pr.GetHTMLURL(),
@@ -71,3 +75,13 @@ func getPullRequest(client *github.Client, ctx *context.Context, repositoryName
 	}
 	return fields
 }
+
+func addPullRequestLabels(client *github.Client, ctx *context.Context, repositoryName string, repositoryOwner string, pullRequestNumber string, labels string) {
+	number, _ := strconv.Atoi(pullRequestNumber)
+	labelsStrArray := strings.Split(labels, ",")
+	for i, labelName := range labelsStrArray {
+		labelsStrArray[i] = strings.Trim(labelName, " ")
+	}
+	_, _, err := client.Issues.AddLabelsToIssue(*ctx, repositoryOwner, repositoryName, number, labelsStrArray)
+	failOnErr(err)
+}