From 7bdaac6882dc34cd8df62e584193cd86ede6cfeb Mon Sep 17 00:00:00 2001
From: Qiu Jian <qiujian@yunionyun.com>
Date: Tue, 15 Jan 2019 12:14:07 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=AD=A3=EF=BC=9A=E5=B0=86=E5=8C=85?=
 =?UTF-8?q?=E5=B9=B4=E5=8C=85=E6=9C=88=E8=B5=84=E6=BA=90=E6=B1=A0=E7=9A=84?=
 =?UTF-8?q?=E8=99=9A=E6=8B=9F=E5=AD=98=E5=82=A8=E4=BB=8E=E5=8F=AF=E7=94=A8?=
 =?UTF-8?q?=E5=8C=BA=E5=9F=9F=E7=9A=84=E5=AD=98=E5=82=A8=E5=88=97=E8=A1=A8?=
 =?UTF-8?q?=E4=B8=AD=E6=8E=92=E9=99=A4?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/compute/models/capabilities.go | 51 +++++++++++++++++++++++++++---
 1 file changed, 46 insertions(+), 5 deletions(-)

diff --git a/pkg/compute/models/capabilities.go b/pkg/compute/models/capabilities.go
index d4b0c9d526..352a36435c 100644
--- a/pkg/compute/models/capabilities.go
+++ b/pkg/compute/models/capabilities.go
@@ -12,6 +12,7 @@ import (
 
 type SCapabilities struct {
 	Hypervisors        []string `json:",allowempty"`
+	ResourceTypes      []string `json:",allowempty"`
 	StorageTypes       []string `json:",allowempty"`
 	GPUModels          []string `json:",allowempty"`
 	MinNicCount        int
@@ -26,6 +27,7 @@ type SCapabilities struct {
 func GetCapabilities(ctx context.Context, userCred mcclient.TokenCredential, query jsonutils.JSONObject, zone *SZone) (SCapabilities, error) {
 	capa := SCapabilities{}
 	capa.Hypervisors = getHypervisors(zone)
+	capa.ResourceTypes = getResourceTypes(zone)
 	capa.StorageTypes = getStorageTypes(zone)
 	capa.GPUModels = getGPUs(zone)
 	capa.SchedPolicySupport = isSchedPolicySupported(zone)
@@ -71,14 +73,12 @@ func getHypervisors(zone *SZone) []string {
 	return hypervisors
 }
 
-func getStorageTypes(zone *SZone) []string {
-	q := StorageManager.Query("storage_type", "medium_type")
+func getResourceTypes(zone *SZone) []string {
+	q := HostManager.Query("resource_type")
 	if zone != nil {
 		q = q.Equals("zone_id", zone.Id)
 	}
-	q = q.IsNotEmpty("storage_type").IsNotNull("storage_type")
-	q = q.IsNotEmpty("medium_type").IsNotNull("medium_type")
-	q = q.In("status", []string{STORAGE_ENABLED, STORAGE_ONLINE})
+	q = q.IsNotEmpty("resource_type").IsNotNull("resource_type")
 	q = q.IsTrue("enabled")
 	q = q.Distinct()
 	rows, err := q.Rows()
@@ -86,6 +86,47 @@ func getStorageTypes(zone *SZone) []string {
 		return nil
 	}
 	defer rows.Close()
+	resourceTypes := make([]string, 0)
+	for rows.Next() {
+		var resType string
+		rows.Scan(&resType)
+		if len(resType) > 0 {
+			resourceTypes = append(resourceTypes, resType)
+		}
+	}
+	return resourceTypes
+}
+
+func getStorageTypes(zone *SZone) []string {
+	storages := StorageManager.Query().SubQuery()
+	hostStorages := HoststorageManager.Query().SubQuery()
+	hosts := HostManager.Query().SubQuery()
+
+	q := storages.Query(storages.Field("storage_type"), storages.Field("medium_type"))
+	q = q.Join(hostStorages, sqlchemy.Equals(
+		hostStorages.Field("storage_id"),
+		storages.Field("id"),
+	))
+	q = q.Join(hosts, sqlchemy.Equals(
+		hosts.Field("id"),
+		hostStorages.Field("host_id"),
+	))
+	if zone != nil {
+		q = q.Filter(sqlchemy.Equals(storages.Field("zone_id"), zone.Id))
+	}
+	q = q.Filter(sqlchemy.Equals(hosts.Field("resource_type"), HostResourceTypeShared))
+	q = q.Filter(sqlchemy.IsNotEmpty(storages.Field("storage_type")))
+	q = q.Filter(sqlchemy.IsNotNull(storages.Field("storage_type")))
+	q = q.Filter(sqlchemy.IsNotEmpty(storages.Field("medium_type")))
+	q = q.Filter(sqlchemy.IsNotNull(storages.Field("medium_type")))
+	q = q.Filter(sqlchemy.In(storages.Field("status"), []string{STORAGE_ENABLED, STORAGE_ONLINE}))
+	q = q.Filter(sqlchemy.IsTrue(storages.Field("enabled")))
+	q = q.Distinct()
+	rows, err := q.Rows()
+	if err != nil {
+		return nil
+	}
+	defer rows.Close()
 	storageTypes := make([]string, 0)
 	for rows.Next() {
 		var storageType, mediumType string