virtbuf/helpers.go

package virtbuf

import (
	"fmt"

	"github.com/google/uuid"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/encoding/prototext"
)

// can the json protobuf output use a string and have a type handler
// to convert it back to int64?
func SetGB(gb int) int64 {
	return int64(gb * 1024 * 1024 * 1024)
}

func SetMB(mb int) int64 {
	return int64(mb * 1024 * 1024)
}

func (x *Hypervisor) SetMemoryGB(gb int) {
	x.Memory = int64(gb * 1024 * 1024 * 1024)
}

func (x *Hypervisor) GetMemoryPrintable() string {
	i := x.Memory / (1024 * 1024 * 1024)
	return fmt.Sprintf("%d GB", i)
}

func (c *Cluster) MarshalJSON() ([]byte, error) {
	return protojson.Marshal(c)
}

func (c *Cluster) FormatJSON() string {
	return protojson.Format(c)
}

func (c *Cluster) UnmarshalJSON(data []byte) error {
	return protojson.Unmarshal(data, c)
}

// apparently this isn't supposed to be used?
// https://protobuf.dev/reference/go/faq/#unstable-text
// this is a shame because this is much nicer output than JSON Format()
func (c *Cluster) FormatTEXT() string {
	return prototext.Format(c)
}

func (d *Droplets) FormatTEXT() string {
	return prototext.Format(d)
}

func (c *Cluster) FindDroplet(name string) *Droplet {
	for _, d := range c.Droplets {
		if d.Hostname == name {
			return d
		}
	}
	return nil
}

func (c *Cluster) FindHypervisor(name string) *Hypervisor {
	for _, h := range c.Hypervisors {
		if h.Hostname == name {
			return h
		}
	}
	return nil
}

func (c *Cluster) AddHypervisor(hostname string, cpus int, mem int) *Hypervisor {
	h := c.FindHypervisor(hostname)
	if h != nil {
		return h
	}
	// Generate a new UUID
	id := uuid.New()
	h = &Hypervisor{
		Uuid:     id.String(),
		Hostname: hostname,
		Cpus:     int64(cpus),
		Comment:  "this is a fake hypervisor",
	}
	if cpus < 0 {
		h.Cpus = 1
	}
	h.SetMemoryGB(mem * 32)
	c.Hypervisors = append(c.Hypervisors, h)
	return h
}

func (c *Cluster) AddDroplet(uuid string, hostname string, cpus int, mem int) *Droplet {
	d := c.FindDroplet(hostname)
	if d != nil {
		return d
	}

	d = &Droplet{
		Uuid:     uuid,
		Hostname: hostname,
		Cpus:     int64(cpus),
	}

	if cpus < 0 {
		d.Cpus = 1
	}
	d.Memory = SetGB(mem * 32)
	if d.Humantest == nil {
		var newInfo StorageInfo
		newInfo = StorageInfo{Capacity: 64}
		d.Humantest = &newInfo
	} else {
		d.Humantest.Capacity = SetGB(mem * 32)
	}
	c.Droplets = append(c.Droplets, d)
	return d
}