Add annotate subcommand

app.go → bebida-shaker.go

@@ -1,11 +1,14 @@
 package main
 
 import (
+	"flag"
+	"fmt"
 	"os"
 	"strconv"
 
 	connectors "github.com/RyaxTech/bebida-shaker/connectors"
 	"github.com/RyaxTech/bebida-shaker/events"
+	"github.com/RyaxTech/bebida-shaker/utils"
 	"github.com/apex/log"
 )
 
@@ -84,11 +87,34 @@ func getStrEnv(envName string, defaultValue string) string {
 }
 
 func main() {
-	log.Info("Starting Bebida Shaker")
-	params = Parameters{
-		maxPendingPunchJob: getIntEnv("BEBIDA_MAX_PENDING_PUNCH_JOB", 2),
-		HPCSchedulerType:   getStrEnv("BEBIDA_HPC_SCHEDULER_TYPE", "OAR"),
+	annotateCmd := flag.NewFlagSet("annotate", flag.ExitOnError)
+	deadline := annotateCmd.String("deadline", "", "App deadline date")
+	duration := annotateCmd.Int("duration", 900, "App duration in seconds")
+	cores := annotateCmd.Int("cores", 1, "Number of cores reqired")
+	memory := annotateCmd.Int("memory", 1024, "Amount of memory reqired in Bytes")
+
+	flag.Parse()
+	if len(os.Args) < 2 {
+		fmt.Println("expected 'shaker' or 'annotate' subcommands")
+		os.Exit(1)
+	}
+	switch os.Args[1] {
+	case "run":
+		log.Info("Starting Bebida Shaker")
+		params = Parameters{
+			maxPendingPunchJob: getIntEnv("BEBIDA_MAX_PENDING_PUNCH_JOB", 2),
+			HPCSchedulerType:   getStrEnv("BEBIDA_HPC_SCHEDULER_TYPE", "OAR"),
+		}
+		log.Infof("Parameters: %+v\n", params)
+		run()
+	case "annotate":
+		annotateCmd.Parse(os.Args[2:])
+		err := utils.Annotate(annotateCmd.Arg(0), *deadline, *duration, *cores, *memory)
+		if err != nil {
+			panic(err)
+		}
+	default:
+		fmt.Println("expected 'annotate' or 'run' subcommands")
+		os.Exit(1)
 	}
-	log.Infof("Parameters: %+v\n", params)
-	run()
 }

flake.nix

@@ -12,7 +12,7 @@
         "x86_64-linux"
         # "aarch64-linux"
       ];
-      inherit (flake-utils.lib) eachSystem filterPackages;
+      inherit (flake-utils.lib) eachSystem;
       bebidaShakerPackage =
         { pkgs }:
         pkgs.buildGoModule {
@@ -27,7 +27,6 @@
           #}
 
           checkPhase = "";
-          # vendorHash = null;
           # vendorHash = pkgs.lib.fakeHash;
           vendorHash = "sha256-n+Pe2nVWlwDLPbzaWTSYtMyYLzMpC1H+oilg7YJhftI=";
         };
@@ -72,7 +71,7 @@
                 Restart = "always";
                 RestartSec = "5s";
                 EnvironmentFile = cfg.environmentFile;
-                ExecStart = "${cfg.package}/bin/bebida-shaker";
+                ExecStart = "${cfg.package}/bin/bebida-shaker run";
               };
             };
           };

go.mod

@@ -44,7 +44,7 @@ require (
 	google.golang.org/protobuf v1.28.1 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
-	gopkg.in/yaml.v3 v3.0.1 // indirect
+	gopkg.in/yaml.v3 v3.0.1
 	k8s.io/api v0.26.0
 	k8s.io/apimachinery v0.26.0
 	k8s.io/klog/v2 v2.80.1 // indirect

utils/bebida.go

@@ -0,0 +1,49 @@
+package utils
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"strconv"
+
+	"gopkg.in/yaml.v3"
+)
+
+var prefix = "ryax.tech/"
+
+func Annotate(filePath string, deadline string, duration int, cores int, memory int) error {
+
+	fileContent, err := ioutil.ReadFile(filePath)
+	if err != nil {
+		return err
+	}
+	decoder := yaml.NewDecoder(bytes.NewBuffer(fileContent))
+	for {
+		data := make(map[interface{}]interface{})
+
+		err := decoder.Decode(&data)
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+			return err
+		}
+		if data["kind"] == "Pod" {
+			annotations := make(map[interface{}]interface{})
+			annotations[prefix+"deadline"] = deadline
+			annotations[prefix+"duration"] = strconv.Itoa(duration)
+			annotations[prefix+"resources.cores"] = strconv.Itoa(cores)
+			annotations[prefix+"resources.memory"] = strconv.Itoa(memory)
+			annotationsInYaml := data["metadata"].(map[string]interface{})
+			annotationsInYaml["annotations"] = annotations
+		}
+		strData, err := yaml.Marshal(&data)
+		if err != nil {
+			return err
+		}
+		fmt.Println(string(strData))
+		fmt.Println("---")
+	}
+	return nil
+}