From 47f3313ffd6ddfc92fc16250b095b4483bbe2b23 Mon Sep 17 00:00:00 2001
From: Piotr Kazmierczak <470696+pkazmierczak@users.noreply.github.com>
Date: Thu, 12 Sep 2024 16:51:53 +0200
Subject: [PATCH] cli: quota status extension for devices (#23899)

quota status CLI now displays device limits (if present in the quota spec)
---
 command/quota_status.go | 60 ++++++++++++++++++++++++++++++++---------
 1 file changed, 47 insertions(+), 13 deletions(-)

diff --git a/command/quota_status.go b/command/quota_status.go
index 7aa6d6a47..30d365bfb 100644
--- a/command/quota_status.go
+++ b/command/quota_status.go
@@ -6,6 +6,7 @@ package command
 import (
 	"encoding/base64"
 	"fmt"
+	"slices"
 	"sort"
 	"strconv"
 	"strings"
@@ -123,6 +124,12 @@ func (c *QuotaStatusCommand) Run(args []string) int {
 	c.Ui.Output(c.Colorize().Color("\n[bold]Quota Limits[reset]"))
 	c.Ui.Output(formatQuotaLimits(spec, usages))
 
+	// If quota has limits on devices, format them separately
+	if slices.ContainsFunc(spec.Limits, func(l *api.QuotaLimit) bool { return l.RegionLimit.Devices != nil }) {
+		c.Ui.Output(c.Colorize().Color("\n[bold]Quota Device Limits[reset]"))
+		c.Ui.Output(formatQuotaDevices(spec, usages))
+	}
+
 	// Display any failures
 	if len(failures) != 0 {
 		c.Ui.Error(c.Colorize().Color("\n[bold][red]Lookup Failures[reset]"))
@@ -176,6 +183,17 @@ func formatQuotaSpecBasics(spec *api.QuotaSpec) string {
 	return formatKV(basic)
 }
 
+// lookupUsage returns the regions quota usage for the limit
+func lookupUsage(usages map[string]*api.QuotaUsage, specLimit *api.QuotaLimit) (*api.QuotaLimit, bool) {
+	usage, ok := usages[specLimit.Region]
+	if !ok {
+		return nil, false
+	}
+
+	used, ok := usage.Used[base64.StdEncoding.EncodeToString(specLimit.Hash)]
+	return used, ok
+}
+
 // formatQuotaLimits formats the limits to display the quota usage versus the
 // limit per quota limit. It takes as input the specification as well as quota
 // usage by region. The formatter handles missing usages.
@@ -193,24 +211,12 @@ func formatQuotaLimits(spec *api.QuotaSpec, usages map[string]*api.QuotaUsage) s
 	for _, specLimit := range spec.Limits {
 		i++
 
-		// lookupUsage returns the regions quota usage for the limit
-		lookupUsage := func() (*api.QuotaLimit, bool) {
-			usage, ok := usages[specLimit.Region]
-			if !ok {
-				return nil, false
-			}
-
-			used, ok := usage.Used[base64.StdEncoding.EncodeToString(specLimit.Hash)]
-			return used, ok
-		}
-
-		used, ok := lookupUsage()
+		used, ok := lookupUsage(usages, specLimit)
 		if !ok {
 			cores := fmt.Sprintf("- / %s", formatQuotaLimitInt(specLimit.RegionLimit.Cores))
 			cpu := fmt.Sprintf("- / %s", formatQuotaLimitInt(specLimit.RegionLimit.CPU))
 			memory := fmt.Sprintf("- / %s", formatQuotaLimitInt(specLimit.RegionLimit.MemoryMB))
 			memoryMax := fmt.Sprintf("- / %s", formatQuotaLimitInt(specLimit.RegionLimit.MemoryMaxMB))
-
 			vars := fmt.Sprintf("- / %s", formatQuotaLimitInt(specLimit.VariablesLimit))
 			limits[i] = fmt.Sprintf("%s|%s|%s|%s|%s|%s", specLimit.Region, cpu, cores, memory, memoryMax, vars)
 			continue
@@ -252,6 +258,34 @@ func formatQuotaLimitInt(value *int) string {
 	return strconv.Itoa(v)
 }
 
+func formatQuotaDevices(spec *api.QuotaSpec, usages map[string]*api.QuotaUsage) string {
+	devices := []string{"Region|Device Name|Device Usage"}
+	i := 0
+
+	for _, specLimit := range spec.Limits {
+		i++
+
+		usage := "-"
+		used, ok := lookupUsage(usages, specLimit)
+		if !ok {
+			for _, d := range specLimit.RegionLimit.Devices {
+				devices = append(devices, fmt.Sprintf("%s|%s|%s / %d", specLimit.Region, d.Name, usage, *d.Count))
+			}
+			continue
+		}
+
+		for _, d := range specLimit.RegionLimit.Devices {
+			idx := slices.IndexFunc(used.RegionLimit.Devices, func(dd *api.RequestedDevice) bool { return dd.Name == d.Name })
+			if idx >= 0 {
+				usage = fmt.Sprintf("%d", int(*used.RegionLimit.Devices[idx].Count))
+			}
+
+			devices = append(devices, fmt.Sprintf("%s|%s|%s / %d", specLimit.Region, d.Name, usage, *d.Count))
+		}
+	}
+	return formatList(devices)
+}
+
 func getQuotaByPrefix(client *api.Quotas, quota string) (match *api.QuotaSpec, possible []*api.QuotaSpec, err error) {
 	// Do a prefix lookup
 	quotas, _, err := client.PrefixList(quota, nil)