From eca032fc765b7bf2de3ad1599f3763913a9de00e Mon Sep 17 00:00:00 2001
From: Andrey Hura <andriy.hura@datarobot.com>
Date: Wed, 18 Dec 2024 19:22:50 +0200
Subject: [PATCH] Add functionality to add labels to PRs (#11)

---
 main.go        |  9 ++++++++-
 pullRequest.go | 16 +++++++++++++++-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/main.go b/main.go
index de6b1b6..d233ef9 100644
--- 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
index 89108f2..9368367 100644
--- 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)
+}