Change vm db structure

Allow remove tags from VM

Author: mhewedy

cmd/tag.go

@@ -26,15 +26,17 @@ import (
 // tagCmd represents the tag command
 var tagCmd = &cobra.Command{
 	Use:   "tag",
-	Short: "Tag a VM",
-	Long: `Tag a VM
+	Short: "Add or remove tag to a VM",
+	Long: `Add or remove tag to a VM
 You can tag a VM as many times as you want
 `,
 	Run: func(cmd *cobra.Command, args []string) {
 		vmName := args[0]
 		tag := args[1]
 
-		err := vms.Tag(vmName, tag)
+		remove, _ := cmd.Flags().GetBool("remove")
+
+		err := vms.Tag(vmName, tag, remove)
 		if err != nil {
 			fmt.Println(err)
 			os.Exit(1)
@@ -63,5 +65,5 @@ func init() {
 
 	// Cobra supports local flags which will only run when this command
 	// is called directly, e.g.:
-	//tagCmd.Flags().BoolP("purge", "p", false, "Purge the IP cache")
+	tagCmd.Flags().BoolP("remove", "r", false, "remove tag")
 }

db/db.go

@@ -1,15 +1,11 @@
 package db
 
 import (
-	"io/ioutil"
 	"log"
 	"os"
-	"strings"
 )
 
 const (
-	image    = "image"
-	tags     = "tags"
 	Username = "vermin"
 )
 
@@ -32,33 +28,6 @@ func GetVMPath(vm string) string {
 	return VMsBaseDir + string(os.PathSeparator) + vm
 }
 
-// ----
-func WriteTag(vmName string, tag string) error {
-	return appendToFile(GetVMPath(vmName)+"/"+tags, []byte(tag+"\n"), 0755)
-}
-
-func ReadTags(vmName string) (string, error) {
-	content, err := readFromFile(vmName, tags)
-	if err != nil {
-		return "", err
-	}
-	return strings.TrimSuffix(strings.ReplaceAll(content, "\n", ", "), ", "), nil
-}
-
-// ----
-
-func WriteImageData(vmName string, imageName string) error {
-	return ioutil.WriteFile(GetVMPath(vmName)+"/"+image, []byte(imageName), 0755)
-}
-
-func ReadImageData(vmName string) (string, error) {
-	content, err := readFromFile(vmName, image)
-	if err != nil {
-		return "", err
-	}
-	return strings.ReplaceAll(content, "\n", ""), nil
-}
-
 func getVerminDir() string {
 	dir, err := os.UserHomeDir()
 	if err != nil {

db/vmdb.go

@@ -0,0 +1,115 @@
+package db
+
+import (
+	"encoding/gob"
+	"io"
+	"io/ioutil"
+	"os"
+	"strings"
+)
+
+// Decided to use gob package instead of protobuf after reading this https://blog.golang.org/gob
+
+const vmdbFilename = "vmdb"
+
+type VMDB struct {
+	Image string
+	Tags  []string
+}
+
+func Load(vm string) (VMDB, error) {
+	migrate(vm)
+	file, err := os.Open(GetVMPath(vm) + "/" + vmdbFilename)
+	if err != nil {
+		return VMDB{}, err
+	}
+	defer file.Close()
+
+	return load(file)
+}
+
+func SetImage(vm string, image string) error {
+	migrate(vm)
+	return override(vm, func(vmdb *VMDB) {
+		vmdb.Image = image
+	})
+}
+
+func AddTag(vm string, tag string) error {
+	migrate(vm)
+	return override(vm, func(vmdb *VMDB) {
+		vmdb.Tags = append(vmdb.Tags, tag)
+	})
+}
+
+func RemoveTag(vm string, tag string) error {
+	migrate(vm)
+	return override(vm, func(vmdb *VMDB) {
+		for i, v := range vmdb.Tags {
+			if v == tag {
+				vmdb.Tags = append(vmdb.Tags[:i], vmdb.Tags[i+1:]...)
+				break
+			}
+		}
+	})
+}
+
+// Override the VMDB for the vm passed as the first parameter
+// by applying the fn function on the data loaded from disk
+// and then write back to the disk
+func override(vm string, fn func(vmdb *VMDB)) error {
+	file, err := os.OpenFile(GetVMPath(vm)+"/"+vmdbFilename, os.O_RDWR|os.O_CREATE, 0755)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	vmdb, err := load(file)
+	if err != nil {
+		return err
+	}
+
+	fn(&vmdb)
+
+	// clear the file before save
+	if err := file.Truncate(0); err != nil {
+		return err
+	}
+	if _, err := file.Seek(0, 0); err != nil {
+		return err
+	}
+	// serialize
+	if err := gob.NewEncoder(file).Encode(vmdb); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func load(file *os.File) (VMDB, error) {
+	var vmdb VMDB
+	if err := gob.NewDecoder(file).Decode(&vmdb); err != nil && io.EOF != err {
+		return VMDB{}, err
+	}
+	return vmdb, nil
+}
+
+// --- migrate old db structure
+func migrate(vm string) {
+	if _, err := os.Stat(GetVMPath(vm) + "/" + vmdbFilename); err == nil {
+		return // already migrated
+	}
+	_ = override(vm, func(vmdb *VMDB) {
+		b, err := ioutil.ReadFile(GetVMPath(vm) + "/image")
+		if err == nil {
+			vmdb.Image = string(b)
+		}
+
+		b, err = ioutil.ReadFile(GetVMPath(vm) + "/tags")
+		if err == nil {
+			tags := string(b)
+			split := strings.Split(tags, "\n")
+			vmdb.Tags = split[:len(split)-1]
+		}
+	})
+}

vms/create.go

@@ -34,7 +34,7 @@ func Create(imageName string, script string, cpus int, mem int) (string, error)
 		return "", err
 	}
 
-	if err = db.WriteImageData(vmName, imageName); err != nil {
+	if err = db.SetImage(vmName, imageName); err != nil {
 		return "", err
 	}
 

vms/list.go

@@ -173,15 +173,14 @@ func getVMInfo(vm string) vmInfo {
 
 	box, _ := db.GetBoxInfo(vm)
 	disk := getDiskSizeInGB(vm, box.HDLocation)
-	image, _ := db.ReadImageData(vm)
-	tags, _ := db.ReadTags(vm)
+	vmdb, _ := db.Load(vm)
 
 	return vmInfo{
 		name:  vm,
-		image: image,
+		image: vmdb.Image,
 		box:   *box,
 		disk:  disk,
-		tags:  tags,
+		tags:  strings.Join(vmdb.Tags, ", "),
 	}
 }
 

vms/mount.go

@@ -15,12 +15,12 @@ func Mount(vmName string, hostPath string) error {
 		return err
 	}
 
-	image, err := db.ReadImageData(vmName)
+	vmdb, err := db.Load(vmName)
 	if err != nil {
 		return err
 	}
 
-	if err = images.CheckCanMount(image); err != nil {
+	if err = images.CheckCanMount(vmdb.Image); err != nil {
 		return err
 	}
 

vms/vms.go

@@ -13,11 +13,15 @@ import (
 	"strings"
 )
 
-func Tag(vmName string, tag string) error {
+func Tag(vmName string, tag string, remove bool) error {
 	if err := checkVM(vmName); err != nil {
 		return err
 	}
-	return db.WriteTag(vmName, tag)
+	if remove {
+		return db.RemoveTag(vmName, tag)
+	} else {
+		return db.AddTag(vmName, tag)
+	}
 }
 
 func Start(vmName string) error {