From 6d67e7fe6ec99ca5a9406722c8c62384e9fd5670 Mon Sep 17 00:00:00 2001
From: ioito <qu_xuan@icloud.com>
Date: Wed, 18 May 2022 21:21:59 +0800
Subject: [PATCH] fix(region): check is gpu by server sku

---
 pkg/compute/models/guest_queries.go | 27 ++++++++++++++++++++++++++-
 pkg/compute/models/guests.go        |  8 +++++++-
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/pkg/compute/models/guest_queries.go b/pkg/compute/models/guest_queries.go
index 775ffb3506..ac707df0ec 100644
--- a/pkg/compute/models/guest_queries.go
+++ b/pkg/compute/models/guest_queries.go
@@ -24,6 +24,7 @@ import (
 	"yunion.io/x/pkg/errors"
 	"yunion.io/x/pkg/tristate"
 	"yunion.io/x/pkg/util/compare"
+	"yunion.io/x/pkg/utils"
 	"yunion.io/x/sqlchemy"
 
 	"yunion.io/x/onecloud/pkg/apis"
@@ -49,6 +50,7 @@ func (manager *SGuestManager) FetchCustomizeColumns(
 	hostRows := manager.SHostResourceBaseManager.FetchCustomizeColumns(ctx, userCred, query, objs, fields, isList)
 	encRows := manager.SEncryptedResourceManager.FetchCustomizeColumns(ctx, userCred, query, objs, fields, isList)
 	guestIds := make([]string, len(objs))
+	guests := make([]SGuest, len(objs))
 	for i := range objs {
 		rows[i] = api.ServerDetails{
 			VirtualResourceDetails: virtRows[i],
@@ -56,7 +58,9 @@ func (manager *SGuestManager) FetchCustomizeColumns(
 
 			EncryptedResourceDetails: encRows[i],
 		}
-		guestIds[i] = objs[i].(*SGuest).GetId()
+		guest := objs[i].(*SGuest)
+		guestIds[i] = guest.GetId()
+		guests[i] = *guest
 	}
 
 	if len(fields) == 0 || fields.Contains("disk") {
@@ -205,6 +209,14 @@ func (manager *SGuestManager) FetchCustomizeColumns(
 				}
 			}
 		}
+		instanceTypes := fetchGuestGpuInstanceTypes(guestIds)
+		if len(instanceTypes) > 0 {
+			for i := range rows {
+				if utils.IsInStringArray(guests[i].InstanceType, instanceTypes) {
+					rows[i].IsGpu = true
+				}
+			}
+		}
 	}
 	if len(fields) == 0 || fields.Contains("cdrom") {
 		gcds := fetchGuestCdroms(guestIds)
@@ -639,6 +651,19 @@ func fetchGuestKeypairs(guestIds []string) map[string]sGuestKeypair {
 	return ret
 }
 
+func fetchGuestGpuInstanceTypes(guestIds []string) []string {
+	sq := GuestManager.Query("instance_type").In("id", guestIds).SubQuery()
+	q := ServerSkuManager.Query("name").In("name", sq).GT("gpu_count", 0).Distinct()
+	instanceTypes := []string{}
+	skus, _ := q.AllStringMap()
+	for i := range skus {
+		for _, v := range skus[i] {
+			instanceTypes = append(instanceTypes, v)
+		}
+	}
+	return instanceTypes
+}
+
 func fetchGuestIsolatedDevices(guestIds []string) map[string][]api.SIsolatedDevice {
 	q := IsolatedDeviceManager.Query().In("guest_id", guestIds)
 	devs := make([]SIsolatedDevice, 0)
diff --git a/pkg/compute/models/guests.go b/pkg/compute/models/guests.go
index 2ca0962770..5a0d848f76 100644
--- a/pkg/compute/models/guests.go
+++ b/pkg/compute/models/guests.go
@@ -481,6 +481,7 @@ func (manager *SGuestManager) ListItemFilter(
 
 	devTypeQ := func(q *sqlchemy.SQuery, checkType *bool, dType string) *sqlchemy.SQuery {
 		if checkType != nil {
+			conditions := []sqlchemy.ICondition{}
 			isodev := IsolatedDeviceManager.Query().SubQuery()
 			sgq := isodev.Query(isodev.Field("guest_id")).
 				Filter(sqlchemy.AND(
@@ -490,7 +491,12 @@ func (manager *SGuestManager) ListItemFilter(
 			if *checkType {
 				cond = sqlchemy.In
 			}
-			return q.Filter(cond(q.Field("id"), sgq))
+			if dType == "GPU" {
+				sq := ServerSkuManager.Query("name").GT("gpu_count", 0).Distinct().SubQuery()
+				conditions = append(conditions, cond(q.Field("instance_type"), sq))
+			}
+			conditions = append(conditions, cond(q.Field("id"), sgq))
+			return q.Filter(sqlchemy.OR(conditions...))
 		}
 		return q
 	}