[RC4] Fix for datadog logs output (#196)

* fix for datadog logs output * updating mocks * PR comment
DataDog · May 29, 2024 · 879a968 · 879a968
1 parent 3e3ff1d
commit 879a968
Show file tree

Hide file tree

Showing 6 changed files with 56 additions and 29 deletions.
diff --git a/pkg/globals/application.go b/pkg/globals/application.go
@@ -1,12 +1,22 @@
 package globals
 
+import (
+	"os"
+)
+
 const (
-	DDServiceName = "kubehound"
-	DDEnv         = "prod"
+	DDServiceName    = "kubehound"
+	DefaultDDEnv     = "dev"
+	DefaultComponent = "kubehound-ingestor"
 )
 
-func ProdEnv() bool {
-	return false
+func GetDDEnv() string {
+	env := os.Getenv("DD_ENV")
+	if env == "" {
+		return DefaultDDEnv
+	}
+
+	return env
 }
 
 const (

diff --git a/pkg/ingestor/puller/mocks/mock_puller.go b/pkg/ingestor/puller/mocks/mock_puller.go
diff --git a/pkg/kubehound/graph/edge/mocks/edge.go b/pkg/kubehound/graph/edge/mocks/edge.go
diff --git a/pkg/telemetry/log/logger.go b/pkg/telemetry/log/logger.go
@@ -2,6 +2,7 @@ package log
 
 import (
 	"context"
+	"os"
 	"strconv"
 	"sync"
 
@@ -18,14 +19,10 @@ type LoggerConfig struct {
 	DD   bool          // Whether Datadog integration is enabled.
 }
 
-const (
-	DefaultComponent = "kubehound"
-)
-
 var globalConfig = LoggerConfig{
 	Tags: logrus.Fields{
 		globals.TagService:   globals.DDServiceName,
-		globals.TagComponent: DefaultComponent,
+		globals.TagComponent: globals.DefaultComponent,
 	},
 	Mu: &sync.Mutex{},
 	DD: true,
@@ -59,9 +56,7 @@ func spanID(span tracer.Span) string {
 // Base returns the base logger for the application.
 func Base() *KubehoundLogger {
 	logger := logrus.WithFields(globalConfig.Tags)
-	if globals.ProdEnv() {
-		logger.Logger.SetFormatter(&logrus.JSONFormatter{})
-	}
+	logger.Logger.SetFormatter(GetLogrusFormatter())
 
 	return &KubehoundLogger{logger}
 }
@@ -121,3 +116,23 @@ func Trace(ctx context.Context, opts ...LoggerOption) *KubehoundLogger {
 
 	return &KubehoundLogger{logger}
 }
+
+func GetLogrusFormatter() logrus.Formatter {
+	switch logFormat := os.Getenv("KH_LOG_FORMAT"); {
+	// Datadog require the logged field to be "message" and not "msg"
+	case logFormat == "dd":
+		formatter := &logrus.JSONFormatter{
+			FieldMap: logrus.FieldMap{
+				logrus.FieldKeyMsg: "message",
+			},
+		}
+
+		return formatter
+	case logFormat == "json":
+		return &logrus.JSONFormatter{}
+	case logFormat == "text":
+		return &logrus.TextFormatter{}
+	default:
+		return &logrus.TextFormatter{}
+	}
+}
diff --git a/pkg/telemetry/profiler/profiler.go b/pkg/telemetry/profiler/profiler.go
@@ -11,7 +11,7 @@ import (
 func Initialize(cfg *config.KubehoundConfig) {
 	opts := []profiler.Option{
 		profiler.WithService(globals.DDServiceName),
-		profiler.WithEnv(globals.DDEnv),
+		profiler.WithEnv(globals.GetDDEnv()),
 		profiler.WithVersion(config.BuildVersion),
 		profiler.WithProfileTypes(
 			profiler.CPUProfile,

diff --git a/pkg/telemetry/tracer/tracer.go b/pkg/telemetry/tracer/tracer.go
@@ -13,7 +13,7 @@ import (
 func Initialize(cfg *config.KubehoundConfig) {
 	// Default options
 	opts := []tracer.StartOption{
-		tracer.WithEnv(globals.DDEnv),
+		tracer.WithEnv(globals.GetDDEnv()),
 		tracer.WithService(globals.DDServiceName),
 		tracer.WithServiceVersion(config.BuildVersion),
 		tracer.WithLogStartup(false),