Skip to content

Commit

Permalink
Change vm db structure
Browse files Browse the repository at this point in the history
Allow remove tags from VM
  • Loading branch information
mhewedy committed May 30, 2020
1 parent 7206689 commit 15722d3
Show file tree
Hide file tree
Showing 7 changed files with 133 additions and 44 deletions.
10 changes: 6 additions & 4 deletions cmd/tag.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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")
}
31 changes: 0 additions & 31 deletions db/db.go
Original file line number Diff line number Diff line change
@@ -1,15 +1,11 @@
package db

import (
"io/ioutil"
"log"
"os"
"strings"
)

const (
image = "image"
tags = "tags"
Username = "vermin"
)

Expand All @@ -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 {
Expand Down
115 changes: 115 additions & 0 deletions db/vmdb.go
Original file line number Diff line number Diff line change
@@ -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]
}
})
}
2 changes: 1 addition & 1 deletion vms/create.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down
7 changes: 3 additions & 4 deletions vms/list.go
Original file line number Diff line number Diff line change
Expand Up @@ -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, ", "),
}
}

Expand Down
4 changes: 2 additions & 2 deletions vms/mount.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down
8 changes: 6 additions & 2 deletions vms/vms.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down

0 comments on commit 15722d3

Please sign in to comment.