Implement env parsing from env vars when config file is not passed

cmd/event-collector/main.go

@@ -5,12 +5,14 @@ import (
 	"fmt"
 	"os"
 	"os/signal"
+	"strings"
 	"syscall"
 	"time"
 
 	"github.com/knadh/koanf/parsers/yaml"
+	"github.com/knadh/koanf/providers/env"
 	"github.com/knadh/koanf/providers/file"
-	"github.com/knadh/koanf/v2"
+	koanf "github.com/knadh/koanf/v2"
 
 	"github.com/G-Research/yunikorn-history-server/internal/config"
 	"github.com/G-Research/yunikorn-history-server/internal/repository"
@@ -26,21 +28,40 @@ var (
 	eventCounts   config.EventTypeCounts
 )
 
-func main() {
-	if len(os.Args) != 2 {
-		fmt.Fprintf(os.Stderr, "usage: %s path/to/config.yml\n", os.Args[0])
-		os.Exit(1)
+func loadConfig(cfgFile string) (*koanf.Koanf, error) {
+	k := koanf.New(".")
+
+	// Try to load from the config file if it's provided
+	if cfgFile != "" {
+		if _, err := os.Stat(cfgFile); err == nil {
+			if err := k.Load(file.Provider(cfgFile), yaml.Parser()); err == nil {
+				// Successfully loaded from file, return
+				return k, nil
+			}
+		}
 	}
 
-	cfgFile := os.Args[1]
-	if _, err := os.Stat(cfgFile); err != nil {
-		fmt.Fprintf(os.Stderr, "Error: cannot open config file %s: %v\n", cfgFile, err)
-		os.Exit(1)
+	// If there's no config file or there's an error reading it, default to env vars
+	// YHS_PARENT1_CHILD1_NAME will be converted into "parent1.child1.name"
+	if err := k.Load(env.Provider("YHS_", ".", func(s string) string {
+		return strings.Replace(strings.ToLower(
+			strings.TrimPrefix(s, "YHS_")), "_", ".", -1)
+	}), nil); err != nil {
+		return nil, fmt.Errorf("error loading environment variables: %v", err)
 	}
 
-	k := koanf.New(".")
-	if err := k.Load(file.Provider(cfgFile), yaml.Parser()); err != nil {
-		fmt.Fprintf(os.Stderr, "error loading file: %v", err)
+	return k, nil
+}
+
+func main() {
+	cfgFile := ""
+	if len(os.Args) == 2 {
+		cfgFile = os.Args[1]
+	}
+
+	k, err := loadConfig(cfgFile)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
 		os.Exit(1)
 	}
 

cmd/event-collector/main_test.go

@@ -0,0 +1,75 @@
+package main
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLoadConfigFromFile(t *testing.T) {
+	// Create a temporary configuration file
+	tmpfile, err := os.CreateTemp("", "example.*.yaml")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		err := os.Remove(tmpfile.Name())
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	// Write a test configuration to the temporary file
+	text := []byte("yunikorn:\n  protocol: http\n  host: localhost\n  port: 8080\nyhs:\n  serverAddr: localhost:8081")
+	if _, err := tmpfile.Write(text); err != nil {
+		t.Fatal(err)
+	}
+	if err := tmpfile.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	// Set some environment variables that should be ignored
+	err = os.Setenv("YHS_CONFIG_IGNORED", "IGNORED")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	k, err := loadConfig(tmpfile.Name())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assert.Equal(t, "", k.String("config.ignored"))
+	assert.Equal(t, "http", k.String("yunikorn.protocol"))
+	assert.Equal(t, "localhost", k.String("yunikorn.host"))
+	assert.Equal(t, 8080, k.Int("yunikorn.port"))
+	assert.Equal(t, "localhost:8081", k.String("yhs.serverAddr"))
+}
+
+func TestLoadConfigFromEnv(t *testing.T) {
+	// Set environment variables
+	err := os.Setenv("YHS_YUNIKORN_PROTOCOL", "http")
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = os.Setenv("YHS_YUNIKORN_HOST", "localhost")
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = os.Setenv("YHS_YUNIKORN_PORT", "8080")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Load with empty config file
+	configFile := "this_file_does_not_exist.yaml"
+	k, err := loadConfig(configFile)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assert.Equal(t, "http", k.String("yunikorn.protocol"))
+	assert.Equal(t, "localhost", k.String("yunikorn.host"))
+	assert.Equal(t, 8080, k.Int("yunikorn.port"))
+}

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/jackc/pgx/v5 v5.5.5
 	github.com/julienschmidt/httprouter v1.3.0
 	github.com/knadh/koanf/parsers/yaml v0.1.0
+	github.com/knadh/koanf/providers/env v0.1.0
 	github.com/knadh/koanf/providers/file v0.1.0
 	github.com/knadh/koanf/v2 v2.1.1
 	github.com/pkg/errors v0.9.1

go.sum

@@ -31,6 +31,8 @@ github.com/knadh/koanf/maps v0.1.1 h1:G5TjmUh2D7G2YWf5SQQqSiHRJEjaicvU0KpypqB3NI
 github.com/knadh/koanf/maps v0.1.1/go.mod h1:npD/QZY3V6ghQDdcQzl1W4ICNVTkohC8E73eI2xW4yI=
 github.com/knadh/koanf/parsers/yaml v0.1.0 h1:ZZ8/iGfRLvKSaMEECEBPM1HQslrZADk8fP1XFUxVI5w=
 github.com/knadh/koanf/parsers/yaml v0.1.0/go.mod h1:cvbUDC7AL23pImuQP0oRw/hPuccrNBS2bps8asS0CwY=
+github.com/knadh/koanf/providers/env v0.1.0 h1:LqKteXqfOWyx5Ab9VfGHmjY9BvRXi+clwyZozgVRiKg=
+github.com/knadh/koanf/providers/env v0.1.0/go.mod h1:RE8K9GbACJkeEnkl8L/Qcj8p4ZyPXZIQ191HJi44ZaQ=
 github.com/knadh/koanf/providers/file v0.1.0 h1:fs6U7nrV58d3CFAFh8VTde8TM262ObYf3ODrc//Lp+c=
 github.com/knadh/koanf/providers/file v0.1.0/go.mod h1:rjJ/nHQl64iYCtAW2QQnF0eSmDEX/YZ/eNFj5yR6BvA=
 github.com/knadh/koanf/v2 v2.1.1 h1:/R8eXqasSTsmDCsAyYj+81Wteg8AqrV9CP6gvsTsOmM=