From ff87797c8efe922768744cdec05b584ea1b15fe9 Mon Sep 17 00:00:00 2001
From: Travis Harmon <harmtrav@gmail.com>
Date: Tue, 27 Jun 2023 00:41:52 -0500
Subject: [PATCH] Make sure key includes additional attributes when necessary

---
 dynago.go | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/dynago.go b/dynago.go
index ab51529..bf8297e 100644
--- a/dynago.go
+++ b/dynago.go
@@ -160,7 +160,8 @@ func (d *Dynago) key(v Keyer) (map[string]*dynamodb.AttributeValue, error) {
 	if err != nil {
 		return nil, fmt.Errorf("d.cachedStruct: %w", err)
 	}
-	for _, primKey := range v.PrimaryKeys() {
+	primKeys := v.PrimaryKeys()
+	for _, primKey := range primKeys {
 		for i := 0; i < ty.NumField(); i++ {
 			if cache[i].attrName == primKey || slices.Contains(cache[i].attrsToCopy, primKey) {
 				av, err := cache[i].attrVal(val)
@@ -171,6 +172,21 @@ func (d *Dynago) key(v Keyer) (map[string]*dynamodb.AttributeValue, error) {
 			}
 		}
 	}
+	if d.config.AdditionalAttrs != nil {
+		d.config.AdditionalAttrs(m, val)
+		for k := range m {
+			isKey := false
+			for _, pk := range primKeys {
+				if pk == k {
+					isKey = true
+					break
+				}
+			}
+			if !isKey {
+				delete(m, k)
+			}
+		}
+	}
 	return m, nil
 }