From 7d2675ca403169101f9eff3047701756f1ac611e Mon Sep 17 00:00:00 2001
From: Siddharth Seth <sseth@cloudera.com>
Date: Tue, 25 Aug 2020 16:47:29 -0700
Subject: [PATCH] CB-7443. Populate the initial set roles (in grains) with
 additional prewarmed roles which may be set up in prewarmed images.

---
 Makefile         |  2 +-
 saltboot/salt.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/Makefile b/Makefile
index f65d7b5..ea73c7d 100644
--- a/Makefile
+++ b/Makefile
@@ -1,6 +1,6 @@
 BINARY=salt-bootstrap
 
-VERSION=0.13.2
+VERSION=0.13.3
 BUILD_TIME=$(shell date +%FT%T)
 LDFLAGS=-ldflags "-X github.com/hortonworks/salt-bootstrap/saltboot.Version=${VERSION} -X github.com/hortonworks/salt-bootstrap/saltboot.BuildTime=${BUILD_TIME}"
 GOFILES_NOVENDOR = $(shell find . -type f -name '*.go' -not -path "./vendor/*" -not -path "./.git/*")
diff --git a/saltboot/salt.go b/saltboot/salt.go
index 775aaac..f7aa1fe 100644
--- a/saltboot/salt.go
+++ b/saltboot/salt.go
@@ -1,6 +1,7 @@
 package saltboot
 
 import (
+	"bufio"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -193,7 +194,26 @@ func SaltMinionRunRequestHandler(w http.ResponseWriter, req *http.Request) {
 	}
 
 	grainConfigPath := baseDir + "/etc/salt/grains"
+	prewarmedRolesPath := baseDir + "/etc/salt/prewarmed_roles"
 	if isGrainsConfigNeeded(grainConfigPath) {
+		// Check if prewarmed roles exist, and add them to the roles before generating the file
+		if shouldAppendPrewarmedRoles(prewarmedRolesPath) {
+			file, err := os.Open(prewarmedRolesPath)
+			if err != nil {
+				resp = model.Response{ErrorText: err.Error(), StatusCode: http.StatusInternalServerError}
+				resp.WriteHttp(w)
+				return
+			}
+			defer file.Close()
+			scanner := bufio.NewScanner(file)
+			for scanner.Scan() {
+				s := strings.TrimSpace(scanner.Text())
+				if len(s) > 0 {
+					saltMinion.Roles = append(saltMinion.Roles, s)
+				}
+			}
+		}
+
 		grainConfig := GrainConfig{Roles: saltMinion.Roles, HostGroup: saltMinion.HostGroup}
 		grainYaml, err := yaml.Marshal(grainConfig)
 		if err != nil {
@@ -476,6 +496,16 @@ func isGrainsConfigNeeded(grainConfigLocation string) bool {
 	return true
 }
 
+func shouldAppendPrewarmedRoles(prewarmRoleLocation string) bool {
+	// Essentially a file exists check.
+	log.Println("[shouldAppendPrewarmedRoles] check whether prewarm roles exist, file location: " + prewarmRoleLocation)
+	_, err := os.Stat(prewarmRoleLocation)
+	if os.IsNotExist(err) {
+		return false
+	}
+	return true
+}
+
 func isSaltMinionRestartNeeded(servers []string) bool {
 	log.Println("[isSaltMinionRestartNeeded] check whether salt-minion requires restart")
 	masterConfFile := "/etc/salt/minion.d/master.conf"