diff --git a/internal/pgscv/config.go b/internal/pgscv/config.go index c131f880..3c74d101 100644 --- a/internal/pgscv/config.go +++ b/internal/pgscv/config.go @@ -10,6 +10,7 @@ import ( "os" "path/filepath" "regexp" + "strings" "time" ) @@ -41,10 +42,10 @@ 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)) @@ -52,14 +53,15 @@ func NewConfig(configFilePath string) (*Config, error) { 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'. diff --git a/internal/pgscv/config_test.go b/internal/pgscv/config_test.go index 07c87a3f..78d0d16e 100644 --- a/internal/pgscv/config_test.go +++ b/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 diff --git a/internal/service/config.go b/internal/service/config.go new file mode 100644 index 00000000..7aaa1ece --- /dev/null +++ b/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 +} diff --git a/internal/service/config_test.go b/internal/service/config_test.go new file mode 100644 index 00000000..a59a3557 --- /dev/null +++ b/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) + } + } +} diff --git a/internal/service/service.go b/internal/service/service.go index 402d3b30..5f286dea 100644 --- a/internal/service/service.go +++ b/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 {