diff --git a/plugin/env.go b/plugin/env.go index 12f39067..85ea1c89 100644 --- a/plugin/env.go +++ b/plugin/env.go @@ -1,21 +1,31 @@ package plugin import ( + "fmt" + "golang.org/x/exp/maps" "log" - "math" "os" "strconv" + "strings" ) const ( envMaxConcurrentConnection = "STEAMPIPE_MAX_CONCURRENT_CONNECTIONS" - envMaxMemoryMb = "STEAMPIPE_MAX_MEMORY_MB" envFreeMemInterval = "STEAMPIPE_FREE_MEM_INTERVAL" - defaultMaxConcurrentConnections = 25 // default to 25 concurrent connections - defaultMaxMemoryMb = math.MaxInt64 // default to no memory limit - defaultFreeMemInterval = 100 // default to freeing memory every 100 rows + defaultMaxConcurrentConnections = 25 // default to 25 concurrent connections + defaultFreeMemInterval = 100 // default to freeing memory every 100 rows + + EnvDiagnosticsLevel = "STEAMPIPE_DIAGNOSTIC_LEVEL" + EnvLegacyDiagnosticsLevel = "STEAMPIPE_DIAGNOSTICS_LEVEL" + DiagnosticsAll = "ALL" + DiagnosticsNone = "NONE" ) +var ValidDiagnosticsLevels = map[string]struct{}{ + DiagnosticsAll: {}, + DiagnosticsNone: {}, +} + func getMaxConcurrentConnections() int { maxConcurrentConnections, _ := strconv.Atoi(os.Getenv(envMaxConcurrentConnection)) if maxConcurrentConnections == 0 { @@ -25,17 +35,6 @@ func getMaxConcurrentConnections() int { return maxConcurrentConnections } -func GetMaxMemoryBytes() int64 { - maxMemoryMb, _ := strconv.Atoi(os.Getenv(envMaxMemoryMb)) - if maxMemoryMb == 0 { - log.Printf("[TRACE] No memory limit set") - maxMemoryMb = defaultMaxMemoryMb - } else { - log.Printf("[TRACE] Setting max memory %dMb", maxMemoryMb) - } - return int64(1024 * 1024 * maxMemoryMb) -} - func GetFreeMemInterval() int64 { freeMemInterval := defaultFreeMemInterval intervalEnv, ok := os.LookupEnv(envFreeMemInterval) @@ -48,3 +47,23 @@ func GetFreeMemInterval() int64 { return int64(freeMemInterval) } + +func loadDiagnosticsEnvVar() string { + // load both the legacy and current diagnostics env vars + diagnostics := strings.ToUpper(os.Getenv(EnvLegacyDiagnosticsLevel)) + if newDiagnostics, isSet := os.LookupEnv(EnvDiagnosticsLevel); isSet { + diagnostics = strings.ToUpper(newDiagnostics) + } + return diagnostics +} + +func ValidateDiagnosticsEnvVar() error { + diagnostics := loadDiagnosticsEnvVar() + if diagnostics == "" { + return nil + } + if _, isValid := ValidDiagnosticsLevels[strings.ToUpper(diagnostics)]; !isValid { + return fmt.Errorf(`invalid value of '%s' (%s), must be one of: %s`, EnvDiagnosticsLevel, diagnostics, strings.Join(maps.Keys(ValidDiagnosticsLevels), ", ")) + } + return nil +} diff --git a/plugin/row_with_metadata.go b/plugin/row_with_metadata.go index 41c5fa98..3e618bc5 100644 --- a/plugin/row_with_metadata.go +++ b/plugin/row_with_metadata.go @@ -1,21 +1,5 @@ package plugin -import ( - "os" - "strings" -) - -const ( - EnvDiagnosticsLevel = "STEAMPIPE_DIAGNOSTICS_LEVEL" - DiagnosticsAll = "ALL" - DiagnosticsNone = "NONE" -) - -var ValidDiagnosticsLevels = map[string]struct{}{ - DiagnosticsAll: {}, - DiagnosticsNone: {}, -} - type hydrateMetadata struct { Type string `json:"type"` FuncName string `json:"function_name"` @@ -38,7 +22,7 @@ func newRowCtxData(rd *rowData) *rowCtxData { Connection: d.Connection.Name, } - if strings.ToUpper(os.Getenv(EnvDiagnosticsLevel)) == DiagnosticsAll { + if loadDiagnosticsEnvVar() == DiagnosticsAll { calls := append([]*hydrateMetadata{d.fetchMetadata}, rd.hydrateMetadata...) if d.parentHydrateMetadata != nil { calls = append([]*hydrateMetadata{d.parentHydrateMetadata}, calls...)