From 1345a0149fd89432fb239573930e4e4671f74c7a Mon Sep 17 00:00:00 2001
From: babattles <babattles@gmail.com>
Date: Sat, 3 Aug 2024 00:59:20 -0700
Subject: [PATCH] omitempty for NULL attribute values from custom marshalers

---
 feature/dynamodb/attributevalue/encode.go     |  7 +++++
 .../dynamodb/attributevalue/encode_test.go    | 29 +++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/feature/dynamodb/attributevalue/encode.go b/feature/dynamodb/attributevalue/encode.go
index 005a23c3b01..4a0aad4c658 100644
--- a/feature/dynamodb/attributevalue/encode.go
+++ b/feature/dynamodb/attributevalue/encode.go
@@ -452,6 +452,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 isNullAttributeValue(av) && fieldTag.OmitEmpty {
+			return nil, nil
 		} else if av != nil {
 			return av, nil
 		}
@@ -893,3 +895,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
index 97a385c6679..4c90022c278 100644
--- a/feature/dynamodb/attributevalue/encode_test.go
+++ b/feature/dynamodb/attributevalue/encode_test.go
@@ -420,6 +420,35 @@ 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"`
+	CustomNullPresent customNullMarshaler
+}
+
+func TestMarshalOmitEmptyCustom(t *testing.T) {
+	expect := &types.AttributeValueMemberM{
+		Value: map[string]types.AttributeValue{
+			"CustomNullPresent": &types.AttributeValueMemberNULL{Value: true},
+		},
+	}
+
+	m := testOmitEmptyCustom{}
+
+	actual, err := Marshal(m)
+	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