From a1acb1f18f29fd5aef759c4f4622b809b9137c9e Mon Sep 17 00:00:00 2001
From: sachin shingare <35319163+sachin120@users.noreply.github.com>
Date: Sun, 7 Jan 2024 01:44:43 +0530
Subject: [PATCH] feat: support for adding, removing, updating machines/devices
 Workload Annotations (#45)

---
 api/device.go     |  1 +
 api/machine.go    |  1 +
 client/device.go  | 16 +++++++++++++++-
 client/machine.go | 15 +++++++++++++++
 entity/device.go  | 33 +++++++++++++++++----------------
 entity/machine.go |  1 +
 6 files changed, 50 insertions(+), 17 deletions(-)

diff --git a/api/device.go b/api/device.go
index 044296c..cdf048e 100644
--- a/api/device.go
+++ b/api/device.go
@@ -10,4 +10,5 @@ type Device interface {
 	Get(systemID string) (*entity.Device, error)
 	Update(systemID string, deviceParams *entity.DeviceUpdateParams) (*entity.Device, error)
 	Delete(systemID string) error
+	SetWorkloadAnnotations(systemID string, params map[string]string) (*entity.Device, error)
 }
diff --git a/api/machine.go b/api/machine.go
index b75a513..4cd3930 100644
--- a/api/machine.go
+++ b/api/machine.go
@@ -18,4 +18,5 @@ type Machine interface {
 	PowerOn(systemID string, params *entity.MachinePowerOnParams) (*entity.Machine, error)
 	PowerOff(systemID string, params *entity.MachinePowerOffParams) (*entity.Machine, error)
 	GetPowerState(systemID string) (*entity.MachinePowerState, error)
+	SetWorkloadAnnotations(systemID string, params map[string]string) (*entity.Machine, error)
 }
diff --git a/client/device.go b/client/device.go
index 02aae7f..07b9b3a 100644
--- a/client/device.go
+++ b/client/device.go
@@ -1,4 +1,3 @@
-//nolint:dupl // disable dupl check on client for now
 package client
 
 import (
@@ -48,3 +47,18 @@ func (d *Device) Update(systemID string, deviceParams *entity.DeviceUpdateParams
 func (d *Device) Delete(systemID string) error {
 	return d.client(systemID).Delete()
 }
+
+// SetWorkloadAnnotations add, modify or remove workload annotations for given Device
+func (d *Device) SetWorkloadAnnotations(systemID string, params map[string]string) (*entity.Device, error) {
+	qsp := url.Values{}
+	for k, v := range params {
+		qsp.Add(k, v)
+	}
+
+	device := new(entity.Device)
+	err := d.client(systemID).Post("set_workload_annotations", qsp, func(data []byte) error {
+		return json.Unmarshal(data, &device)
+	})
+
+	return device, err
+}
diff --git a/client/machine.go b/client/machine.go
index ef99fae..190b669 100644
--- a/client/machine.go
+++ b/client/machine.go
@@ -171,3 +171,18 @@ func (m *Machine) GetPowerState(systemID string) (*entity.MachinePowerState, err
 
 	return ps, err
 }
+
+// SetWorkloadAnnotations add, modify or remove workload annotations for given Machine
+func (m *Machine) SetWorkloadAnnotations(systemID string, params map[string]string) (*entity.Machine, error) {
+	qsp := url.Values{}
+	for k, v := range params {
+		qsp.Add(k, v)
+	}
+
+	machine := new(entity.Machine)
+	err := m.client(systemID).Post("set_workload_annotations", qsp, func(data []byte) error {
+		return json.Unmarshal(data, &machine)
+	})
+
+	return machine, err
+}
diff --git a/entity/device.go b/entity/device.go
index a003b5c..9c64b5d 100644
--- a/entity/device.go
+++ b/entity/device.go
@@ -3,22 +3,23 @@ package entity
 import "net"
 
 type Device struct {
-	OwnerData    interface{}        `json:"owner_data,omitempty"`
-	Zone         Zone               `json:"zone,omitempty"`
-	FQDN         string             `json:"fqdn,omitempty"`
-	Hostname     string             `json:"hostname,omitempty"`
-	SystemID     string             `json:"system_id,omitempty"`
-	Owner        string             `json:"owner,omitempty"`
-	Description  string             `json:"description,omitempty"`
-	Parent       string             `json:"parent,omitempty"`
-	ResourceURI  string             `json:"resource_uri,omitempty"`
-	NodeTypeName string             `json:"node_type_name,omitempty"`
-	IPAddresses  []net.IP           `json:"ip_addresses,omitempty"`
-	InterfaceSet []NetworkInterface `json:"interface_set,omitempty"`
-	TagNames     []string           `json:"tag_names,omitempty"`
-	Domain       Domain             `json:"domain,omitempty"`
-	NodeType     int                `json:"node_type,omitempty"`
-	AddressTTL   int                `json:"address_ttl,omitempty"`
+	OwnerData           interface{}        `json:"owner_data,omitempty"`
+	Zone                Zone               `json:"zone,omitempty"`
+	FQDN                string             `json:"fqdn,omitempty"`
+	Hostname            string             `json:"hostname,omitempty"`
+	SystemID            string             `json:"system_id,omitempty"`
+	Owner               string             `json:"owner,omitempty"`
+	Description         string             `json:"description,omitempty"`
+	Parent              string             `json:"parent,omitempty"`
+	ResourceURI         string             `json:"resource_uri,omitempty"`
+	NodeTypeName        string             `json:"node_type_name,omitempty"`
+	WorkloadAnnotations map[string]string  `json:"workload_annotations,omitempty"`
+	IPAddresses         []net.IP           `json:"ip_addresses,omitempty"`
+	InterfaceSet        []NetworkInterface `json:"interface_set,omitempty"`
+	TagNames            []string           `json:"tag_names,omitempty"`
+	Domain              Domain             `json:"domain,omitempty"`
+	NodeType            int                `json:"node_type,omitempty"`
+	AddressTTL          int                `json:"address_ttl,omitempty"`
 }
 
 type DeviceCreateParams struct {
diff --git a/entity/machine.go b/entity/machine.go
index d320e11..dd8c5d2 100644
--- a/entity/machine.go
+++ b/entity/machine.go
@@ -16,6 +16,7 @@ type Machine struct {
 	HwLastSync              MAASTime                   `json:"last_sync,omitempty"`
 	OwnerData               interface{}                `json:"owner_data,omitempty"`
 	HardwareInfo            map[string]string          `json:"hardware_info,omitempty"`
+	WorkloadAnnotations     map[string]string          `json:"workload_annotations,omitempty"`
 	OSystem                 string                     `json:"osystem,omitempty"`
 	FQDN                    string                     `json:"fqdn,omitempty"`
 	StatusName              string                     `json:"status_name,omitempty"`