From 425faf28b8c5ac886c91638fe3b69f11a07b9f45 Mon Sep 17 00:00:00 2001
From: lenhattan86 <lenhattan86@users.noreply.github.com>
Date: Thu, 16 Sep 2021 16:29:25 -0700
Subject: [PATCH] Adds priority for aurora-scheduler (#13)

Adds priority for task config
---
 job.go       | 5 +++++
 jobUpdate.go | 5 +++++
 task.go      | 9 ++++++++-
 task_test.go | 1 +
 4 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/job.go b/job.go
index b553520..698ecfa 100644
--- a/job.go
+++ b/job.go
@@ -156,6 +156,11 @@ func (j *AuroraJob) IsService(isService bool) *AuroraJob {
 	return j
 }
 
+func (j *AuroraJob) Priority(priority int32) *AuroraJob {
+	j.task.Priority(priority)
+	return j
+}
+
 func (j *AuroraJob) TaskConfig() *aurora.TaskConfig {
 	return j.task.TaskConfig()
 }
diff --git a/jobUpdate.go b/jobUpdate.go
index 3ef2cb4..0f231f2 100644
--- a/jobUpdate.go
+++ b/jobUpdate.go
@@ -221,6 +221,11 @@ func (j *JobUpdate) IsService(isService bool) *JobUpdate {
 	return j
 }
 
+func (j *JobUpdate) Priority(priority int32) *JobUpdate {
+	j.task.Priority(priority)
+	return j
+}
+
 func (j *JobUpdate) TaskConfig() *aurora.TaskConfig {
 	return j.task.TaskConfig()
 }
diff --git a/task.go b/task.go
index e4037c3..e055190 100644
--- a/task.go
+++ b/task.go
@@ -78,7 +78,8 @@ func TaskFromThrift(config *aurora.TaskConfig) *AuroraTask {
 		Role(config.Job.Role).
 		Name(config.Job.Name).
 		MaxFailure(config.MaxTaskFailures).
-		IsService(config.IsService)
+		IsService(config.IsService).
+		Priority(config.Priority)
 
 	if config.Tier != nil {
 		newTask.Tier(*config.Tier)
@@ -287,6 +288,12 @@ func (t *AuroraTask) IsService(isService bool) *AuroraTask {
 	return t
 }
 
+//set priority for preemption or priority-queueing
+func (t *AuroraTask) Priority(priority int32) *AuroraTask {
+	t.task.Priority = priority
+	return t
+}
+
 // Add a list of URIs with the same extract and cache configuration. Scheduler must have
 // --enable_mesos_fetcher flag enabled. Currently there is no duplicate detection.
 func (t *AuroraTask) AddURIs(extract bool, cache bool, values ...string) *AuroraTask {
diff --git a/task_test.go b/task_test.go
index 14d06c8..deaf7b1 100644
--- a/task_test.go
+++ b/task_test.go
@@ -34,6 +34,7 @@ func TestAuroraTask_Clone(t *testing.T) {
 		RAM(643).
 		Disk(1000).
 		IsService(true).
+		Priority(1).
 		AddPorts(10).
 		Tier("preferred").
 		MaxFailure(23).