From b59b1c4519f1ea59ac39a0278b182c0adf28f5da Mon Sep 17 00:00:00 2001
From: Zexi Li <lizexi@yunionyun.com>
Date: Mon, 14 Jan 2019 12:28:34 +0800
Subject: [PATCH] =?UTF-8?q?region:=20host=20=E6=94=AF=E6=8C=81=E9=80=9A?=
 =?UTF-8?q?=E8=BF=87=20set-schedtag=20=E6=8E=A5=E5=8F=A3=E6=B7=BB=E5=8A=A0?=
 =?UTF-8?q?=E5=88=A0=E9=99=A4=E8=B0=83=E5=BA=A6=E6=A0=87=E7=AD=BE?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/climc/shell/hosts.go    | 17 +++++++++++
 pkg/compute/models/hosts.go | 60 +++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+)

diff --git a/cmd/climc/shell/hosts.go b/cmd/climc/shell/hosts.go
index 3b7783bdbd..5ae2b87798 100644
--- a/cmd/climc/shell/hosts.go
+++ b/cmd/climc/shell/hosts.go
@@ -504,4 +504,21 @@ func init() {
 		printObject(result)
 		return nil
 	})
+
+	type HostSetSchedtagOptions struct {
+		ID       string   `help:"ID or Name of host"`
+		Schedtag []string `help:"Ids of schedtag"`
+	}
+	R(&HostSetSchedtagOptions{}, "host-set-schedtag", "Set schedtags to a host", func(s *mcclient.ClientSession, args *HostSetSchedtagOptions) error {
+		params := jsonutils.NewDict()
+		for idx, tag := range args.Schedtag {
+			params.Add(jsonutils.NewString(tag), fmt.Sprintf("schedtag.%d", idx))
+		}
+		result, err := modules.Hosts.PerformAction(s, args.ID, "set-schedtag", params)
+		if err != nil {
+			return err
+		}
+		printObject(result)
+		return nil
+	})
 }
diff --git a/pkg/compute/models/hosts.go b/pkg/compute/models/hosts.go
index 50becad1cd..634d13f7c8 100644
--- a/pkg/compute/models/hosts.go
+++ b/pkg/compute/models/hosts.go
@@ -3643,3 +3643,63 @@ func (manager *SHostManager) PingDetectionTask(ctx context.Context, userCred mcc
 func (self *SHost) IsPrepaidRecycleResource() bool {
 	return self.ResourceType == HostResourceTypePrepaidRecycle
 }
+
+func (host *SHost) AllowPerformSetSchedtag(ctx context.Context, userCred mcclient.TokenCredential, query jsonutils.JSONObject, data jsonutils.JSONObject) bool {
+	return db.IsAdminAllowPerform(userCred, host, "set-schedtag")
+}
+
+func (host *SHost) getHostschedtags() []SHostschedtag {
+	tags := make([]SHostschedtag, 0)
+	hostschedtags := HostschedtagManager.Query().SubQuery()
+	q := hostschedtags.Query().Equals("host_id", host.Id)
+	err := db.FetchModelObjects(HostschedtagManager, q, &tags)
+	if err != nil {
+		log.Errorf("Get hostschedtags: %v", err)
+		return nil
+	}
+	return tags
+}
+
+func (host *SHost) PerformSetSchedtag(ctx context.Context, userCred mcclient.TokenCredential, query jsonutils.JSONObject, data jsonutils.JSONObject) (jsonutils.JSONObject, error) {
+	schedtags := jsonutils.GetArrayOfPrefix(data, "schedtag")
+	setTagsId := []string{}
+	for idx := 0; idx < len(schedtags); idx++ {
+		schedtagIdent, _ := schedtags[idx].GetString()
+		tag, err := SchedtagManager.FetchByIdOrName(userCred, schedtagIdent)
+		if err != nil {
+			if err == sql.ErrNoRows {
+				return nil, httperrors.NewNotFoundError("Schedtag %s not found", schedtagIdent)
+			}
+			return nil, httperrors.NewGeneralError(err)
+		}
+		setTagsId = append(setTagsId, tag.GetId())
+	}
+	oldTags := host.getHostschedtags()
+	for _, oldTag := range oldTags {
+		if !utils.IsInStringArray(oldTag.SchedtagId, setTagsId) {
+			if err := oldTag.Detach(ctx, userCred); err != nil {
+				return nil, httperrors.NewGeneralError(err)
+			}
+		}
+	}
+	var oldTagIds []string
+	for _, tag := range oldTags {
+		oldTagIds = append(oldTagIds, tag.SchedtagId)
+	}
+	for _, setTagId := range setTagsId {
+		if !utils.IsInStringArray(setTagId, oldTagIds) {
+			if newTagObj, err := db.NewModelObject(HostschedtagManager); err != nil {
+				return nil, httperrors.NewGeneralError(err)
+			} else {
+				newTag := newTagObj.(*SHostschedtag)
+				newTag.HostId = host.Id
+				newTag.SchedtagId = setTagId
+				if err := newTag.GetModelManager().TableSpec().Insert(newTag); err != nil {
+					return nil, httperrors.NewGeneralError(err)
+				}
+			}
+		}
+	}
+	host.ClearSchedDescCache()
+	return nil, nil
+}