omitempty for NULL attribute values from custom marshalers (#2739)

aws · Aug 24, 2024 · 84ca95e · 84ca95e
1 parent d7a7f5a
commit 84ca95e
Show file tree

Hide file tree

Showing 5 changed files with 107 additions and 0 deletions.
diff --git a/.changelog/281f2e28209e44068c04c9d74bab8970.json b/.changelog/281f2e28209e44068c04c9d74bab8970.json
@@ -0,0 +1,9 @@
+{
+    "id": "281f2e28-209e-4406-8c04-c9d74bab8970",
+    "type": "feature",
+    "description": "Add Encoder option to obey omitempty tag for NULL attribute values.",
+    "modules": [
+        "feature/dynamodb/attributevalue",
+        "feature/dynamodbstreams/attributevalue"
+    ]
+}
diff --git a/feature/dynamodb/attributevalue/encode.go b/feature/dynamodb/attributevalue/encode.go
@@ -392,6 +392,13 @@ type EncoderOptions struct {
 	// The results of a MarshalText call will convert to string (S), results
 	// from a MarshalBinary call will convert to binary (B).
 	UseEncodingMarshalers bool
+
+	// When enabled, the encoder will omit null (NULL) attribute values
+	// returned from custom marshalers tagged with `omitempty`.
+	//
+	// NULL attribute values returned from the standard marshaling routine will
+	// always respect omitempty regardless of this setting.
+	OmitNullAttributeValues bool
 }
 
 // An Encoder provides marshaling Go value types to AttributeValues.
@@ -452,6 +459,8 @@ func (e *Encoder) encode(v reflect.Value, fieldTag tag) (types.AttributeValue, e
 	if v.Kind() != reflect.Invalid {
 		if av, err := e.tryMarshaler(v); err != nil {
 			return nil, err
+		} else if e.options.OmitNullAttributeValues && fieldTag.OmitEmpty && isNullAttributeValue(av) {
+			return nil, nil
 		} else if av != nil {
 			return av, nil
 		}
@@ -893,3 +902,8 @@ func defaultEncodeTime(t time.Time) (types.AttributeValue, error) {
 		Value: t.Format(time.RFC3339Nano),
 	}, nil
 }
+
+func isNullAttributeValue(av types.AttributeValue) bool {
+	n, ok := av.(*types.AttributeValueMemberNULL)
+	return ok && n.Value
+}
diff --git a/feature/dynamodb/attributevalue/encode_test.go b/feature/dynamodb/attributevalue/encode_test.go
@@ -420,6 +420,41 @@ func TestMarshalOmitEmpty(t *testing.T) {
 	}
 }
 
+type customNullMarshaler struct{}
+
+func (m customNullMarshaler) MarshalDynamoDBAttributeValue() (types.AttributeValue, error) {
+	return &types.AttributeValueMemberNULL{Value: true}, nil
+}
+
+type testOmitEmptyCustom struct {
+	CustomNullOmit       customNullMarshaler `dynamodbav:",omitempty"`
+	CustomNullOmitTagKey customNullMarshaler `tagkey:",omitempty"`
+	CustomNullPresent    customNullMarshaler
+	EmptySetOmit         []string `dynamodbav:",omitempty"`
+}
+
+func TestMarshalOmitEmptyCustom(t *testing.T) {
+	expect := &types.AttributeValueMemberM{
+		Value: map[string]types.AttributeValue{
+			"CustomNullPresent": &types.AttributeValueMemberNULL{Value: true},
+		},
+	}
+
+	m := testOmitEmptyCustom{}
+
+	actual, err := MarshalWithOptions(m, func(eo *EncoderOptions) {
+		eo.TagKey = "tagkey"
+		eo.OmitNullAttributeValues = true
+		eo.NullEmptySets = true
+	})
+	if err != nil {
+		t.Errorf("expect nil, got %v", err)
+	}
+	if e, a := expect, actual; !reflect.DeepEqual(e, a) {
+		t.Errorf("expect %v, got %v", e, a)
+	}
+}
+
 func TestEncodeEmbeddedPointerStruct(t *testing.T) {
 	type B struct {
 		Bint int

diff --git a/feature/dynamodbstreams/attributevalue/encode.go b/feature/dynamodbstreams/attributevalue/encode.go
@@ -392,6 +392,13 @@ type EncoderOptions struct {
 	// The results of a MarshalText call will convert to string (S), results
 	// from a MarshalBinary call will convert to binary (B).
 	UseEncodingMarshalers bool
+
+	// When enabled, the encoder will omit null (NULL) attribute values
+	// returned from custom marshalers tagged with `omitempty`.
+	//
+	// NULL attribute values returned from the standard marshaling routine will
+	// always respect omitempty regardless of this setting.
+	OmitNullAttributeValues bool
 }
 
 // An Encoder provides marshaling Go value types to AttributeValues.
@@ -452,6 +459,8 @@ func (e *Encoder) encode(v reflect.Value, fieldTag tag) (types.AttributeValue, e
 	if v.Kind() != reflect.Invalid {
 		if av, err := e.tryMarshaler(v); err != nil {
 			return nil, err
+		} else if e.options.OmitNullAttributeValues && fieldTag.OmitEmpty && isNullAttributeValue(av) {
+			return nil, nil
 		} else if av != nil {
 			return av, nil
 		}
@@ -893,3 +902,8 @@ func defaultEncodeTime(t time.Time) (types.AttributeValue, error) {
 		Value: t.Format(time.RFC3339Nano),
 	}, nil
 }
+
+func isNullAttributeValue(av types.AttributeValue) bool {
+	n, ok := av.(*types.AttributeValueMemberNULL)
+	return ok && n.Value
+}
diff --git a/feature/dynamodbstreams/attributevalue/encode_test.go b/feature/dynamodbstreams/attributevalue/encode_test.go
@@ -420,6 +420,41 @@ func TestMarshalOmitEmpty(t *testing.T) {
 	}
 }
 
+type customNullMarshaler struct{}
+
+func (m customNullMarshaler) MarshalDynamoDBStreamsAttributeValue() (types.AttributeValue, error) {
+	return &types.AttributeValueMemberNULL{Value: true}, nil
+}
+
+type testOmitEmptyCustom struct {
+	CustomNullOmit       customNullMarshaler `dynamodbav:",omitempty"`
+	CustomNullOmitTagKey customNullMarshaler `tagkey:",omitempty"`
+	CustomNullPresent    customNullMarshaler
+	EmptySetOmit         []string `dynamodbav:",omitempty"`
+}
+
+func TestMarshalOmitEmptyCustom(t *testing.T) {
+	expect := &types.AttributeValueMemberM{
+		Value: map[string]types.AttributeValue{
+			"CustomNullPresent": &types.AttributeValueMemberNULL{Value: true},
+		},
+	}
+
+	m := testOmitEmptyCustom{}
+
+	actual, err := MarshalWithOptions(m, func(eo *EncoderOptions) {
+		eo.TagKey = "tagkey"
+		eo.OmitNullAttributeValues = true
+		eo.NullEmptySets = true
+	})
+	if err != nil {
+		t.Errorf("expect nil, got %v", err)
+	}
+	if e, a := expect, actual; !reflect.DeepEqual(e, a) {
+		t.Errorf("expect %v, got %v", e, a)
+	}
+}
+
 func TestEncodeEmbeddedPointerStruct(t *testing.T) {
 	type B struct {
 		Bint int