From aa81230e5361ba4ba8efb29da765e7abad99f67f Mon Sep 17 00:00:00 2001
From: oseoin <oseoin@users.noreply.github.com>
Date: Thu, 25 Jan 2024 10:18:06 +0000
Subject: [PATCH] Add Environment Variable options to flags managed in
 exporter-toolkit (#607)

* add env var options to flags managed in exporter-toolkit
---
 README.md        |  6 +++--
 exporter.go      | 24 +++++++++++++++++++
 exporter_test.go | 61 ++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 89 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 8582efe..c204e6f 100644
--- a/README.md
+++ b/README.md
@@ -110,9 +110,11 @@ usage: nginx-prometheus-exporter [<flags>]
 
 Flags:
   -h, --[no-]help                Show context-sensitive help (also try --help-long and --help-man).
+      --[no-]web.systemd-socket  Use systemd socket activation listeners instead
+                                 of port listeners (Linux only). ($SYSTEMD_SOCKET)
       --web.listen-address=:9113 ...
-                                 Addresses on which to expose metrics and web interface. Repeatable for multiple addresses.
-      --web.config.file=""       Path to configuration file that can enable TLS or authentication. See: https://github.com/prometheus/exporter-toolkit/blob/master/docs/web-configuration.md
+                                 Addresses on which to expose metrics and web interface. Repeatable for multiple addresses. ($LISTEN_ADDRESS)
+      --web.config.file=""       Path to configuration file that can enable TLS or authentication. See: https://github.com/prometheus/exporter-toolkit/blob/master/docs/web-configuration.md ($CONFIG_FILE)
       --web.telemetry-path="/metrics"
                                  Path under which to expose metrics. ($TELEMETRY_PATH)
       --[no-]nginx.plus          Start the exporter for NGINX Plus. By default, the exporter is started for NGINX. ($NGINX_PLUS)
diff --git a/exporter.go b/exporter.go
index 8abd28d..93dd8e1 100644
--- a/exporter.go
+++ b/exporter.go
@@ -114,6 +114,9 @@ func main() {
 	flag.AddFlags(kingpin.CommandLine, promlogConfig)
 	kingpin.Version(version.Print(exporterName))
 	kingpin.HelpFlag.Short('h')
+
+	addMissingEnvironmentFlags(kingpin.CommandLine)
+
 	kingpin.Parse()
 	logger := promlog.New(promlogConfig)
 
@@ -281,3 +284,24 @@ func cloneRequest(req *http.Request) *http.Request {
 	}
 	return r
 }
+
+// addMissingEnvironmentFlags sets Envar on any flag which has
+// the "web." prefix which doesn't already have an Envar set
+func addMissingEnvironmentFlags(ka *kingpin.Application) {
+	for _, f := range ka.Model().FlagGroupModel.Flags {
+		if strings.HasPrefix(f.Name, "web.") && f.Envar == "" {
+			flag := ka.GetFlag(f.Name)
+			if flag != nil {
+				flag.Envar(convertFlagToEnvar(strings.TrimPrefix(f.Name, "web.")))
+			}
+		}
+	}
+}
+
+func convertFlagToEnvar(f string) string {
+	env := strings.ToUpper(f)
+	for _, s := range []string{"-", "."} {
+		env = strings.ReplaceAll(env, s, "_")
+	}
+	return env
+}
diff --git a/exporter_test.go b/exporter_test.go
index 9005b23..eba5928 100644
--- a/exporter_test.go
+++ b/exporter_test.go
@@ -4,6 +4,9 @@ import (
 	"reflect"
 	"testing"
 	"time"
+
+	"github.com/alecthomas/kingpin/v2"
+	"github.com/prometheus/exporter-toolkit/web/kingpinflag"
 )
 
 func TestParsePositiveDuration(t *testing.T) {
@@ -103,3 +106,61 @@ func TestParseUnixSocketAddress(t *testing.T) {
 		})
 	}
 }
+
+func TestAddMissingEnvironmentFlags(t *testing.T) {
+	expectedMatches := map[string]string{
+		"non-matching-flag":  "",
+		"web.missing-env":    "MISSING_ENV",
+		"web.has-env":        "HAS_ENV_ALREADY",
+		"web.listen-address": "LISTEN_ADDRESS",
+		"web.config.file":    "CONFIG_FILE",
+	}
+	kingpinflag.AddFlags(kingpin.CommandLine, ":9113")
+	kingpin.Flag("non-matching-flag", "").String()
+	kingpin.Flag("web.missing-env", "").String()
+	kingpin.Flag("web.has-env", "").Envar("HAS_ENV_ALREADY").String()
+	addMissingEnvironmentFlags(kingpin.CommandLine)
+
+	// using Envar() on a flag returned from GetFlag()
+	// adds an additional flag, which is processed correctly
+	// at runtime but means that we need to check for a match
+	// instead of checking the envar of each matching flag name
+	for k, v := range expectedMatches {
+		matched := false
+		for _, f := range kingpin.CommandLine.Model().FlagGroupModel.Flags {
+			if f.Name == k && f.Envar == v {
+				matched = true
+			}
+		}
+		if !matched {
+			t.Errorf("missing %s envar for %s", v, k)
+		}
+	}
+}
+
+func TestConvertFlagToEnvar(t *testing.T) {
+	cases := []struct {
+		input  string
+		output string
+	}{
+		{
+			input:  "dot.separate",
+			output: "DOT_SEPARATE",
+		},
+		{
+			input:  "underscore_separate",
+			output: "UNDERSCORE_SEPARATE",
+		},
+		{
+			input:  "mixed_separate_options",
+			output: "MIXED_SEPARATE_OPTIONS",
+		},
+	}
+
+	for _, c := range cases {
+		res := convertFlagToEnvar(c.input)
+		if res != c.output {
+			t.Errorf("expected %s to resolve to %s but got %s", c.input, c.output, res)
+		}
+	}
+}