From 27e5f8f0b3c815cb540f9c88e6c3b9d6b8e0811a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Piotr=20Poniedzia=C5=82ek?= <pondzix@gmail.com>
Date: Fri, 19 Jul 2024 11:23:38 +0200
Subject: [PATCH] Remove `null` fields from JQ transformation output

We don't want explicit null to be present on JQ output, so we have to
get rid of them somehow. As there is no option neither in `gojq` nor
`json.Marshal` allowing us to to that, it seems the best way is to simply
iterate over a map (we have a map between `gojq` output and `json.Marshal` input in transformation) and filter out fields being nulls.
---
 pkg/transform/jq.go      | 30 ++++++++++++++++
 pkg/transform/jq_test.go | 76 ++++++++++++++++++++++++++++++++++++++--
 2 files changed, 103 insertions(+), 3 deletions(-)

diff --git a/pkg/transform/jq.go b/pkg/transform/jq.go
index 14b708a4..0d6f729b 100644
--- a/pkg/transform/jq.go
+++ b/pkg/transform/jq.go
@@ -63,6 +63,8 @@ func (jqm *jqMapper) RunFunction() TransformationFunction {
 			return nil, nil, message, nil
 		}
 
+		removeNullFields(v)
+
 		// here v is any, so we Marshal. alternative: gojq.Marshal
 		data, err := json.Marshal(v)
 		if err != nil {
@@ -169,3 +171,31 @@ func mkJQInput(jqm *jqMapper, message *models.Message, interState interface{}) (
 
 	return spInput, nil
 }
+
+func removeNullFields(data any) {
+	switch input := data.(type) {
+	case map[string]any:
+		removeNullFromMap(input)
+	case []any:
+		removeNullFromSlice(input)
+	default:
+		return
+	}
+}
+
+func removeNullFromMap(input map[string]any) {
+	for key := range input {
+		field := input[key]
+		if field == nil {
+			delete(input, key)
+			continue
+		}
+		removeNullFields(field)
+	}
+}
+
+func removeNullFromSlice(input []any) {
+	for _, item := range input {
+		removeNullFields(item)
+	}
+}
diff --git a/pkg/transform/jq_test.go b/pkg/transform/jq_test.go
index 4493c964..62a31e53 100644
--- a/pkg/transform/jq_test.go
+++ b/pkg/transform/jq_test.go
@@ -216,16 +216,86 @@ func TestJQRunFunction_SpMode_false(t *testing.T) {
 			JQCommand: `
       { 
         explicit_null: .explicit | epoch,
-        no_such_field: .nonexistent | epoch
+        no_such_field: .nonexistent | epoch,
+        non_null: .non_null
       }`,
 			InputMsg: &models.Message{
-				Data:         []byte(`{"explicit": null}`),
+				Data:         []byte(`{"explicit": null, "non_null": "hello"}`),
 				PartitionKey: "some-key",
 			},
 			InputInterState: nil,
 			Expected: map[string]*models.Message{
 				"success": {
-					Data:         []byte(`{"explicit_null":null,"no_such_field":null}`),
+					Data:         []byte(`{"non_null":"hello"}`),
+					PartitionKey: "some-key",
+				},
+				"filtered": nil,
+				"failed":   nil,
+			},
+			ExpInterState: nil,
+			Error:         nil,
+		},
+		{
+			Scenario:  "remove_nulls_struct",
+			JQCommand: ".",
+			InputMsg: &models.Message{
+				Data: []byte(`
+        {
+          "f1": "value1",
+          "f2": 2,
+          "f3": {
+            "f5": null,
+            "f6": "value6",
+            "f7": {
+              "f8": 100,
+              "f9": null
+             }
+           },
+          "f4": null
+        }`),
+				PartitionKey: "some-key",
+			},
+			InputInterState: nil,
+			Expected: map[string]*models.Message{
+				"success": {
+					Data:         []byte(`{"f1":"value1","f2":2,"f3":{"f6":"value6","f7":{"f8":100}}}`),
+					PartitionKey: "some-key",
+				},
+				"filtered": nil,
+				"failed":   nil,
+			},
+			ExpInterState: nil,
+			Error:         nil,
+		},
+		{
+			Scenario:  "remove_nulls_arrays",
+			JQCommand: ".",
+			InputMsg: &models.Message{
+				Data: []byte(`
+          {
+            "items": [
+              {
+                "f1": "value1",
+                "f2": null,
+                "f3": [
+                  {
+                    "f4": 1,
+                    "f5": null
+                  },
+                  {
+                    "f4": null,
+                    "f5": 20
+                  }
+                ]
+              }
+            ]
+          }`),
+				PartitionKey: "some-key",
+			},
+			InputInterState: nil,
+			Expected: map[string]*models.Message{
+				"success": {
+					Data:         []byte(`{"items":[{"f1":"value1","f3":[{"f4":1},{"f5":20}]}]}`),
 					PartitionKey: "some-key",
 				},
 				"filtered": nil,