From c2d33fa2eb7918d7577d59318036c0414f52270c Mon Sep 17 00:00:00 2001
From: charithabandi <chbandi302@gmail.com>
Date: Tue, 7 May 2024 15:35:19 -0500
Subject: [PATCH] Load env variables into KwildConfig

---
 cmd/kwild/config/config.go | 44 ++++++++++++++++++++++++++++++++------
 cmd/kwild/root/root.go     |  2 ++
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/cmd/kwild/config/config.go b/cmd/kwild/config/config.go
index d63b1270d..42749fe9b 100644
--- a/cmd/kwild/config/config.go
+++ b/cmd/kwild/config/config.go
@@ -426,15 +426,28 @@ func LoadConfigFile(configPath string) (*KwildConfig, error) {
 
 // LoadEnvConfig loads a config from environment variables.
 func LoadEnvConfig() (*KwildConfig, error) {
-	viper.SetEnvPrefix("KWILD")
-	viper.AutomaticEnv()
-
-	var cfg KwildConfig
-	if err := viper.Unmarshal(&cfg); err != nil {
+	// Manually bind environment variables to viper keys.
+	for _, key := range viper.AllKeys() {
+		// Replace dashes with underscores in the key to match the flag name.
+		// This is required because there is inconsistency between our flag names
+		// and the struct tags. The struct tags use underscores, but the flag names
+		// use dashes. Viper uses the flag names to bind environment variables
+		// and this conversion is required to map it to the struct fields correctly.
+		bindKey := strings.ReplaceAll(key, "-", "_")
+		envKey := "KWILD_" + strings.ToUpper(strings.ReplaceAll(bindKey, ".", "_"))
+		viper.BindEnv(bindKey, envKey)
+	}
+
+	// var cfg KwildConfig, won't work because, viper won't be able to extract
+	// the heirarchical keys from the config structure as fields like cfg.app set to nil.
+	// It can only extract the first level keys [app, chain, log] in this case.
+	// To remedy that, we use DefaultEmptyConfig with all the sub fields initialized.
+	cfg := DefaultEmptyConfig()
+	if err := viper.Unmarshal(cfg); err != nil {
 		return nil, fmt.Errorf("decoding config: %v", err)
 	}
 
-	return &cfg, nil
+	return cfg, nil
 }
 
 var ErrConfigFileNotFound = fmt.Errorf("config file not found")
@@ -444,6 +457,25 @@ func fileExists(path string) bool {
 	return !os.IsNotExist(err)
 }
 
+// DefaultEmptyConfig returns a config with all fields set to their zero values.
+// This is used by viper to extract all the heirarchical keys from the config
+// structure.
+func DefaultEmptyConfig() *KwildConfig {
+	return &KwildConfig{
+		AppCfg: &AppConfig{
+			Extensions: make(map[string]map[string]string),
+		},
+		ChainCfg: &ChainConfig{
+			P2P:       &P2PConfig{},
+			RPC:       &ChainRPCConfig{},
+			Mempool:   &MempoolConfig{},
+			StateSync: &StateSyncConfig{},
+			Consensus: &ConsensusConfig{},
+		},
+		Logging: &Logging{},
+	}
+}
+
 func DefaultConfig() *KwildConfig {
 	return &KwildConfig{
 		AppCfg: &AppConfig{
diff --git a/cmd/kwild/root/root.go b/cmd/kwild/root/root.go
index 5610f4ad4..eab5cf4ba 100644
--- a/cmd/kwild/root/root.go
+++ b/cmd/kwild/root/root.go
@@ -23,6 +23,7 @@ import (
 	"github.com/kwilteam/kwil-db/internal/version"
 
 	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
 )
 
 func RootCmd() *cobra.Command {
@@ -89,6 +90,7 @@ func RootCmd() *cobra.Command {
 	flagSet := cmd.Flags()
 	flagSet.SortFlags = false
 	config.AddConfigFlags(flagSet, flagCfg)
+	viper.BindPFlags(flagSet)
 
 	flagSet.BoolVarP(&autoGen, "autogen", "a", false,
 		"auto generate private key, genesis file, and config file if not exist")