Skip to content

Commit

Permalink
add stock query commands
Browse files Browse the repository at this point in the history
  • Loading branch information
caguiclajmg committed Jun 27, 2022
1 parent 1f75033 commit 185e8cb
Show file tree
Hide file tree
Showing 2 changed files with 136 additions and 0 deletions.
43 changes: 43 additions & 0 deletions api/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,21 @@ type DeployServerResponse struct {
} `json:"server"`
}

type ListGpuStockResponse struct {
Response
Stock map[string]map[string]struct {
AvailableNow int `json:"available_now"`
AvailableReserve int `json:"available_reserve"`
} `json:"stock"`
}

type ListCpuStockResponse struct {
Response
Stock map[string]map[string]struct {
AvailableNow string `json:"available_now"`
} `json:"stock"`
}

type Client struct {
BaseUrl string
ApiKey string
Expand Down Expand Up @@ -283,3 +298,31 @@ func (client *Client) RestartServer(server string) (*Response, error) {

return &res, nil
}

func (client *Client) ListGpuStock() (*ListGpuStockResponse, error) {
raw, err := client.get("stock/list", nil, false)
if err != nil {
return nil, err
}

var res ListGpuStockResponse
if err := json.Unmarshal(*raw, &res); err != nil {
return nil, err
}

return &res, nil
}

func (client *Client) ListCpuStock() (*ListCpuStockResponse, error) {
raw, err := client.get("stock/cpu/list", nil, false)
if err != nil {
return nil, err
}

var res ListCpuStockResponse
if err := json.Unmarshal(*raw, &res); err != nil {
return nil, err
}

return &res, nil
}
93 changes: 93 additions & 0 deletions commands/stock.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
package commands

import (
"errors"
"os"

"github.com/jedib0t/go-pretty/v6/table"
"github.com/spf13/cobra"
)

var (
stockCmd = &cobra.Command{
Use: "stock",
Short: "Query stock",
}
listStockCmd = &cobra.Command{
Use: "list",
Short: "List stock",
RunE: listStock,
}
)

func init() {
stockCmd.AddCommand(listStockCmd)
listStockCmd.Flags().String("type", "gpu", "Instance type (gpu or cpu)")
listStockCmd.Flags().Bool("all", false, "Include out-of-stock instances")

rootCmd.AddCommand(stockCmd)
}

func listStock(cmd *cobra.Command, args []string) error {
instanceType, err := cmd.Flags().GetString("type")
if err != nil {
return err
}

all, err := cmd.Flags().GetBool("all")
if err != nil {
return err
}

t := table.NewWriter()
t.SetOutputMirror(os.Stdout)

switch instanceType {
case "gpu":
res, err := client.ListGpuStock()
if err != nil {
return nil
}

if !res.Success {
return errors.New(res.Error)
}

t.AppendHeader(table.Row{"GPU", "Region", "Available Now", "Available Reserve"})

for gpuModel, regionStock := range res.Stock {
for region, stock := range regionStock {
if (stock.AvailableNow > 0 || stock.AvailableReserve > 0) || all {
t.AppendRow(table.Row{gpuModel, region, stock.AvailableNow, stock.AvailableReserve})
}
}
}

case "cpu":
res, err := client.ListCpuStock()
if err != nil {
return nil
}

if !res.Success {
return errors.New(res.Error)
}

t.AppendHeader(table.Row{"CPU Model", "Region", "Available Now"})

for cpuModel, regionStock := range res.Stock {
for region, stock := range regionStock {
if stock.AvailableNow != "None" || all {
t.AppendRow(table.Row{cpuModel, region, stock.AvailableNow})
}
}
}

default:
return errors.New("unknown instance type")
}

t.Render()

return nil
}

0 comments on commit 185e8cb

Please sign in to comment.