Add support of environment variables.

pgSCV can be configured using config file only, which might be tricky in containerized environments (requires volumes for config). This commit allows define basic configuration using environment variables. This is not support complex configuration related to collectors settings or user-defined metrics.

internal/pgscv/config.go

@@ -10,6 +10,7 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"strings"
 	"time"
 )
 
@@ -41,25 +42,26 @@ type Config struct {
 	DatabasesRE           *regexp.Regexp           // Regular expression object compiled from Databases
 }
 
-// NewConfig creates new config based on config file or return default config of config is not exists.
+// NewConfig creates new config based on config file or return default config if config file is not specified.
 func NewConfig(configFilePath string) (*Config, error) {
 	if configFilePath == "" {
-		return &Config{Defaults: map[string]string{}}, nil
+		return newConfigFromEnv()
 	}
 
 	content, err := os.ReadFile(filepath.Clean(configFilePath))
 	if err != nil {
 		return nil, err
 	}
 
-	config := Config{Defaults: map[string]string{}}
-	err = yaml.Unmarshal(content, &config)
+	config := &Config{Defaults: map[string]string{}}
+
+	err = yaml.Unmarshal(content, config)
 	if err != nil {
 		return nil, err
 	}
 
 	log.Infoln("read configuration from ", configFilePath)
-	return &config, nil
+	return config, nil
 }
 
 // Validate checks configuration for stupid values and set defaults
@@ -214,6 +216,75 @@ func validateCollectorSettings(cs model.CollectorsSettings) error {
 	return nil
 }
 
+// newConfigFromEnv create config using environment variables.
+func newConfigFromEnv() (*Config, error) {
+	config := &Config{
+		Defaults: map[string]string{},
+	}
+
+	for _, env := range os.Environ() {
+		if !strings.HasPrefix(env, "PGSCV_") &&
+			!strings.HasPrefix(env, "POSTGRES_DSN") &&
+			!strings.HasPrefix(env, "PGBOUNCER_DSN") {
+			continue
+		}
+
+		ff := strings.SplitN(env, "=", 2)
+
+		key, value := ff[0], ff[1]
+
+		// Parse POSTGRES_DSN.
+		if strings.HasPrefix(key, "POSTGRES_DSN") {
+			id, cs, err := service.ParsePostgresDSNEnv(key, value)
+			if err != nil {
+				return nil, err
+			}
+
+			if config.ServicesConnsSettings == nil {
+				config.ServicesConnsSettings = map[string]service.ConnSetting{}
+			}
+
+			config.ServicesConnsSettings[id] = cs
+		}
+
+		// Parse PGBOUNCER_DSN.
+		if strings.HasPrefix(key, "PGBOUNCER_DSN") {
+			id, cs, err := service.ParsePgbouncerDSNEnv(key, value)
+			if err != nil {
+				return nil, err
+			}
+
+			if config.ServicesConnsSettings == nil {
+				config.ServicesConnsSettings = map[string]service.ConnSetting{}
+			}
+
+			config.ServicesConnsSettings[id] = cs
+		}
+
+		switch key {
+		case "PGSCV_LISTEN_ADDRESS":
+			config.ListenAddress = value
+		case "PGSCV_AUTOUPDATE":
+			config.AutoUpdate = value
+		case "PGSCV_NO_TRACK_MODE":
+			switch value {
+			case "y", "yes", "Yes", "YES", "t", "true", "True", "TRUE", "1", "on":
+				config.NoTrackMode = true
+			default:
+				config.NoTrackMode = false
+			}
+		case "PGSCV_SEND_METRICS_URL":
+			config.SendMetricsURL = value
+		case "PGSCV_API_KEY":
+			config.APIKey = value
+		case "PGSCV_DATABASES":
+			config.Databases = value
+		}
+	}
+
+	return config, nil
+}
+
 // toggleAutoupdate control auto-update setting.
 func toggleAutoupdate(value string) (string, error) {
 	// Empty value explicitly set to 'off'.

internal/pgscv/config_test.go

@@ -5,6 +5,7 @@ import (
 	"github.com/weaponry/pgscv/internal/filter"
 	"github.com/weaponry/pgscv/internal/model"
 	"github.com/weaponry/pgscv/internal/service"
+	"os"
 	"testing"
 )
 
@@ -383,6 +384,76 @@ func Test_validateCollectorSettings(t *testing.T) {
 	}
 }
 
+func Test_newConfigFromEnv(t *testing.T) {
+	testcases := []struct {
+		valid   bool
+		envvars map[string]string
+		want    *Config
+	}{
+		{
+			valid:   true, // No env variables
+			envvars: map[string]string{},
+			want:    &Config{Defaults: map[string]string{}},
+		},
+		{
+			valid: true, // Completely valid variables
+			envvars: map[string]string{
+				"PGSCV_LISTEN_ADDRESS":   "127.0.0.1:12345",
+				"PGSCV_AUTOUPDATE":       "1",
+				"PGSCV_NO_TRACK_MODE":    "yes",
+				"PGSCV_SEND_METRICS_URL": "127.0.0.1:54321",
+				"PGSCV_API_KEY":          "example",
+				"PGSCV_DATABASES":        "exampledb",
+				"POSTGRES_DSN":           "example_dsn",
+				"POSTGRES_DSN_EXAMPLE1":  "example_dsn",
+				"PGBOUNCER_DSN":          "example_dsn",
+				"PGBOUNCER_DSN_EXAMPLE2": "example_dsn",
+			},
+			want: &Config{
+				ListenAddress:  "127.0.0.1:12345",
+				AutoUpdate:     "1",
+				NoTrackMode:    true,
+				SendMetricsURL: "127.0.0.1:54321",
+				APIKey:         "example",
+				Databases:      "exampledb",
+				ServicesConnsSettings: map[string]service.ConnSetting{
+					"postgres":  {ServiceType: "postgres", Conninfo: "example_dsn"},
+					"EXAMPLE1":  {ServiceType: "postgres", Conninfo: "example_dsn"},
+					"pgbouncer": {ServiceType: "pgbouncer", Conninfo: "example_dsn"},
+					"EXAMPLE2":  {ServiceType: "pgbouncer", Conninfo: "example_dsn"},
+				},
+				Defaults: map[string]string{},
+			},
+		},
+		{
+			valid:   false, // Invalid postgres DSN key
+			envvars: map[string]string{"POSTGRES_DSN_": "example_dsn"},
+		},
+		{
+			valid:   false, // Invalid pgbouncer DSN key
+			envvars: map[string]string{"PGBOUNCER_DSN_": "example_dsn"},
+		},
+	}
+
+	for _, tc := range testcases {
+		for k, v := range tc.envvars {
+			assert.NoError(t, os.Setenv(k, v))
+		}
+
+		got, err := newConfigFromEnv()
+		if tc.valid {
+			assert.NoError(t, err)
+			assert.Equal(t, tc.want, got)
+		} else {
+			assert.Error(t, err)
+		}
+
+		for k := range tc.envvars {
+			assert.NoError(t, os.Unsetenv(k))
+		}
+	}
+}
+
 func Test_toggleAutoupdate(t *testing.T) {
 	testcases := []struct {
 		valid bool

internal/service/config.go

@@ -0,0 +1,61 @@
+package service
+
+import (
+	"fmt"
+	"strings"
+)
+
+// ConnSetting describes connection settings required for connecting to particular service.
+// This is primarily used for describing services defined by user in the config file (or env vars).
+type ConnSetting struct {
+	// ServiceType defines type of service for which these connection settings are used.
+	ServiceType string `yaml:"service_type"`
+	// Conninfo is the connection string in service-specific format.
+	Conninfo string `yaml:"conninfo"`
+}
+
+// ConnsSettings defines a set of all connection settings of exact services.
+type ConnsSettings map[string]ConnSetting
+
+// ParsePostgresDSNEnv is a public wrapper over parseDSNEnv.
+func ParsePostgresDSNEnv(key, value string) (string, ConnSetting, error) {
+	return parseDSNEnv("POSTGRES_DSN", key, value)
+}
+
+// ParsePgbouncerDSNEnv is a public wrapper over parseDSNEnv.
+func ParsePgbouncerDSNEnv(key, value string) (string, ConnSetting, error) {
+	return parseDSNEnv("PGBOUNCER_DSN", key, value)
+}
+
+// parseDSNEnv returns valid ConnSetting accordingly to passed prefix and environment key/value.
+func parseDSNEnv(prefix, key, value string) (string, ConnSetting, error) {
+	var stype string
+	switch prefix {
+	case "POSTGRES_DSN":
+		stype = "postgres"
+	case "PGBOUNCER_DSN":
+		stype = "pgbouncer"
+	default:
+		return "", ConnSetting{}, fmt.Errorf("invalid prefix %s", prefix)
+	}
+
+	// Prefix must be the part of key.
+	if !strings.HasPrefix(key, prefix) {
+		return "", ConnSetting{}, fmt.Errorf("invalid key %s", key)
+	}
+
+	// Nothing to parse if prefix and key are the same, just use the type as service ID.
+	if key == prefix {
+		return stype, ConnSetting{ServiceType: stype, Conninfo: value}, nil
+	}
+
+	// If prefix and key are not the same, strip prefix from key and use the rest as service ID.
+	// Use double Trim to avoid leaking 'prefix' string into ID value (see unit tests for examples).
+	id := strings.TrimPrefix(strings.TrimPrefix(key, prefix), "_")
+
+	if id == "" {
+		return "", ConnSetting{}, fmt.Errorf("invalid value '%s' is in %s", value, key)
+	}
+
+	return id, ConnSetting{ServiceType: stype, Conninfo: value}, nil
+}

internal/service/config_test.go

@@ -0,0 +1,59 @@
+package service
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func Test_ParsePostgresDSNEnv(t *testing.T) {
+	gotID, gotCS, err := ParsePostgresDSNEnv("POSTGRES_DSN", "conninfo")
+	assert.NoError(t, err)
+	assert.Equal(t, "postgres", gotID)
+	assert.Equal(t, ConnSetting{ServiceType: "postgres", Conninfo: "conninfo"}, gotCS)
+
+	_, _, err = ParsePostgresDSNEnv("INVALID", "conninfo")
+	assert.Error(t, err)
+}
+
+func Test_ParsePgbouncerDSNEnv(t *testing.T) {
+	gotID, gotCS, err := ParsePgbouncerDSNEnv("PGBOUNCER_DSN", "conninfo")
+	assert.NoError(t, err)
+	assert.Equal(t, "pgbouncer", gotID)
+	assert.Equal(t, ConnSetting{ServiceType: "pgbouncer", Conninfo: "conninfo"}, gotCS)
+
+	_, _, err = ParsePgbouncerDSNEnv("INVALID", "conninfo")
+	assert.Error(t, err)
+}
+
+func Test_parseDSNEnv(t *testing.T) {
+	testcases := []struct {
+		valid    bool
+		prefix   string
+		key      string
+		wantId   string
+		wantType string
+	}{
+		{valid: true, prefix: "POSTGRES_DSN", key: "POSTGRES_DSN", wantId: "postgres", wantType: "postgres"},
+		{valid: true, prefix: "POSTGRES_DSN", key: "POSTGRES_DSN_POSTGRES_123", wantId: "POSTGRES_123", wantType: "postgres"},
+		{valid: true, prefix: "POSTGRES_DSN", key: "POSTGRES_DSN1", wantId: "1", wantType: "postgres"},
+		{valid: true, prefix: "POSTGRES_DSN", key: "POSTGRES_DSN_POSTGRES_5432", wantId: "POSTGRES_5432", wantType: "postgres"},
+		{valid: true, prefix: "PGBOUNCER_DSN", key: "PGBOUNCER_DSN", wantId: "pgbouncer", wantType: "pgbouncer"},
+		{valid: true, prefix: "PGBOUNCER_DSN", key: "PGBOUNCER_DSN_PGBOUNCER_123", wantId: "PGBOUNCER_123", wantType: "pgbouncer"},
+		{valid: true, prefix: "PGBOUNCER_DSN", key: "PGBOUNCER_DSN1", wantId: "1", wantType: "pgbouncer"},
+		{valid: true, prefix: "PGBOUNCER_DSN", key: "PGBOUNCER_DSN_PGBOUNCER_6432", wantId: "PGBOUNCER_6432", wantType: "pgbouncer"},
+		{valid: false, prefix: "POSTGRES_DSN", key: "POSTGRES_DSN_"},
+		{valid: false, prefix: "POSTGRES_DSN", key: "INVALID"},
+		{valid: false, prefix: "INVALID", key: "INVALID"},
+	}
+
+	for _, tc := range testcases {
+		gotID, gotCS, err := parseDSNEnv(tc.prefix, tc.key, "conninfo")
+		if tc.valid {
+			assert.NoError(t, err)
+			assert.Equal(t, tc.wantId, gotID)
+			assert.Equal(t, ConnSetting{ServiceType: tc.wantType, Conninfo: "conninfo"}, gotCS)
+		} else {
+			assert.Error(t, err)
+		}
+	}
+}

internal/service/service.go

@@ -68,18 +68,6 @@ type Collector interface {
 	Collect(chan<- prometheus.Metric)
 }
 
-// ConnSetting describes connection settings required for connecting to particular service. This struct primarily
-// is used for representing services defined by user in the config file.
-type ConnSetting struct {
-	// ServiceType defines type of service for which these connection settings are used.
-	ServiceType string `yaml:"service_type"`
-	// Conninfo is the connection string in service-specific format.
-	Conninfo string `yaml:"conninfo"`
-}
-
-// ConnsSettings defines a set of all connection settings of exact services.
-type ConnsSettings map[string]ConnSetting
-
 // connectionParams is the set of parameters that may be required when constructing connection string.
 // For example, this struct describes the postmaster.pid representation https://www.postgresql.org/docs/current/storage-file-layout.html
 type connectionParams struct {