feature/dynamodb/expression: support for composite keys

.changelog/01fe7c45-628d-4d98-a55b-e07d304448ae.json

@@ -0,0 +1,9 @@
+{
+    "id": "01fe7c45-628d-4d98-a55b-e07d304448ae",
+    "type": "feature",
+    "description": "expression: support for composite keys",
+    "collapse": false,
+    "modules": [
+        "feature/dynamodb/expression"
+    ]
+}

feature/dynamodb/expression/composite_key.go

@@ -0,0 +1,88 @@
+package expression
+
+import (
+	"fmt"
+	"strings"
+)
+
+// compositeKeyCondition builds the composite key condition expression for DynamoDB queries.
+// It counts partition and sort keys, validates the presence of at least one partition key,
+// and formats the expression accordingly.
+func compositeKeyCondition(keyConditionBuilder KeyConditionBuilder, node exprNode) (exprNode, error) {
+	pks := 0
+	sks := 0
+	for _, kcb := range keyConditionBuilder.keyConditionList {
+		switch kcb.compositeKeyMode {
+		case compositeKeyCondPartition:
+			pks++
+		case compositeKeyCondSort:
+			sks++
+		}
+	}
+
+	if pks == 0 {
+		return exprNode{}, newInvalidParameterError("compositeKeyCondition", "KeyConditionBuilder")
+	}
+
+	node.fmtExpr = fmt.Sprintf(
+		"(%s)",
+		strings.Repeat(" AND ($c)", pks+sks)[5:],
+	)
+
+	return node, nil
+}
+
+// CompositeKey initializes a KeyConditionBuilder in composite key mode.
+//
+// Note: When using CompositeKey, you must use AddPartitionKey and AddSortKey in the desired order.
+// The order in which you call these methods determines the order of keys in the resulting expression.
+// The And() method is not supported for composite keys.
+//
+// Example:
+//
+//	// Create a composite key condition with partition and sort keys
+//	partitionKey := expression.Key("TeamName").Equal(expression.Value("Wildcats"))
+//	sortKey := expression.Key("Number").Equal(expression.Value(1))
+//	keyCondition := expression.CompositeKey().AddPartitionKey(partitionKey).AddSortKey(sortKey)
+//	builder := expression.NewBuilder().WithKeyCondition(keyCondition)
+func CompositeKey() KeyConditionBuilder {
+	return KeyConditionBuilder{
+		mode: compositeKeyCond,
+	}
+}
+
+// AddPartitionKey adds a partition key condition to the composite key builder.
+//
+// Note: The order in which AddPartitionKey and AddSortKey are called determines
+// the order of keys in the resulting expression. Do not change the order after building.
+//
+// Example:
+//
+//	partitionKey := expression.Key("TeamName").Equal(expression.Value("Wildcats"))
+//	keyCondition := expression.CompositeKey().AddPartitionKey(partitionKey)
+func (kcb KeyConditionBuilder) AddPartitionKey(pk KeyConditionBuilder) KeyConditionBuilder {
+	pk.compositeKeyMode = compositeKeyCondPartition
+	kcb.keyConditionList = append(kcb.keyConditionList, pk)
+
+	return kcb
+}
+
+// AddSortKey adds a sort key condition to the composite key builder.
+//
+// Note: The order in which AddPartitionKey and AddSortKey are called determines
+// the order of keys in the resulting expression. Do not change the order after building.
+//
+// IMPORTANT: You cannot use AddSortKey alone; you must also add at least one partition key
+// using AddPartitionKey, or the query will fail.
+//
+// Example:
+//
+//	partitionKey := expression.Key("TeamName").Equal(expression.Value("Wildcats"))
+//	sortKey := expression.Key("Number").Equal(expression.Value(1))
+//	keyCondition := expression.CompositeKey().AddPartitionKey(partitionKey).AddSortKey(sortKey)
+func (kcb KeyConditionBuilder) AddSortKey(sk KeyConditionBuilder) KeyConditionBuilder {
+	sk.compositeKeyMode = compositeKeyCondSort
+	kcb.keyConditionList = append(kcb.keyConditionList, sk)
+
+	return kcb
+}

feature/dynamodb/expression/composite_key_test.go

@@ -0,0 +1,180 @@
+package expression
+
+import (
+	"reflect"
+	"strconv"
+	"testing"
+
+	"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
+)
+
+func TestCompositeKey(t *testing.T) {
+	cases := []struct {
+		input       KeyConditionBuilder
+		expected    Expression
+		expectError bool
+	}{
+		{
+			input:       CompositeKey(),
+			expectError: true,
+		},
+		{
+			input: CompositeKey().
+				AddPartitionKey(Key("pk1").Equal(Value("1"))),
+			expected: Expression{
+				expressionMap: map[expressionType]string{
+					keyCondition: "((#0 = :0))",
+				},
+				namesMap: map[string]string{
+					"#0": "pk1",
+				},
+				valuesMap: map[string]types.AttributeValue{
+					":0": &types.AttributeValueMemberS{Value: "1"},
+				},
+			},
+		},
+		{
+			input: CompositeKey().
+				AddPartitionKey(Key("pk1").Equal(Value("1"))).
+				AddSortKey(Key("sk1").Equal(Value("1"))),
+			expected: Expression{
+				expressionMap: map[expressionType]string{
+					keyCondition: "((#0 = :0) AND (#1 = :1))",
+				},
+				namesMap: map[string]string{
+					"#0": "pk1",
+					"#1": "sk1",
+				},
+				valuesMap: map[string]types.AttributeValue{
+					":0": &types.AttributeValueMemberS{Value: "1"},
+					":1": &types.AttributeValueMemberS{Value: "1"},
+				},
+			},
+		},
+		{
+			input: CompositeKey().
+				AddPartitionKey(Key("pk1").Equal(Value("1"))).
+				AddPartitionKey(Key("pk2").Equal(Value("1"))).
+				AddPartitionKey(Key("pk3").Equal(Value("1"))).
+				AddPartitionKey(Key("pk4").Equal(Value("1"))),
+			expected: Expression{
+				expressionMap: map[expressionType]string{
+					keyCondition: "((#0 = :0) AND (#1 = :1) AND (#2 = :2) AND (#3 = :3))",
+				},
+				namesMap: map[string]string{
+					"#0": "pk1",
+					"#1": "pk2",
+					"#2": "pk3",
+					"#3": "pk4",
+				},
+				valuesMap: map[string]types.AttributeValue{
+					":0": &types.AttributeValueMemberS{Value: "1"},
+					":1": &types.AttributeValueMemberS{Value: "1"},
+					":2": &types.AttributeValueMemberS{Value: "1"},
+					":3": &types.AttributeValueMemberS{Value: "1"},
+				},
+			},
+		},
+		{
+			input: CompositeKey().
+				AddPartitionKey(Key("pk1").Equal(Value("1"))).
+				AddPartitionKey(Key("pk2").Equal(Value("1"))).
+				AddPartitionKey(Key("pk3").Equal(Value("1"))).
+				AddPartitionKey(Key("pk4").Equal(Value("1"))).
+				AddSortKey(Key("sk1").Equal(Value("1"))).
+				AddSortKey(Key("sk2").Equal(Value("1"))).
+				AddSortKey(Key("sk3").Equal(Value("1"))).
+				AddSortKey(Key("sk4").Equal(Value("1"))),
+			expected: Expression{
+				expressionMap: map[expressionType]string{
+					keyCondition: "((#0 = :0) AND (#1 = :1) AND (#2 = :2) AND (#3 = :3) AND (#4 = :4) AND (#5 = :5) AND (#6 = :6) AND (#7 = :7))",
+				},
+				namesMap: map[string]string{
+					"#0": "pk1",
+					"#1": "pk2",
+					"#2": "pk3",
+					"#3": "pk4",
+					"#4": "sk1",
+					"#5": "sk2",
+					"#6": "sk3",
+					"#7": "sk4",
+				},
+				valuesMap: map[string]types.AttributeValue{
+					":0": &types.AttributeValueMemberS{Value: "1"},
+					":1": &types.AttributeValueMemberS{Value: "1"},
+					":2": &types.AttributeValueMemberS{Value: "1"},
+					":3": &types.AttributeValueMemberS{Value: "1"},
+					":4": &types.AttributeValueMemberS{Value: "1"},
+					":5": &types.AttributeValueMemberS{Value: "1"},
+					":6": &types.AttributeValueMemberS{Value: "1"},
+					":7": &types.AttributeValueMemberS{Value: "1"},
+				},
+			},
+		},
+		{
+			input: CompositeKey().
+				AddSortKey(Key("sk1").Equal(Value("1"))),
+			expectError: true,
+		},
+		{
+			input: CompositeKey().
+				AddSortKey(Key("sk1").Equal(Value("1"))).
+				AddSortKey(Key("sk2").Equal(Value("1"))).
+				AddSortKey(Key("sk3").Equal(Value("1"))).
+				AddSortKey(Key("sk4").Equal(Value("1"))),
+			expectError: true,
+		},
+		{
+			input: CompositeKey().
+				AddPartitionKey(Key("pk1").Equal(Value("1"))).
+				And(Key("sk1").Equal(Value("1"))),
+			expectError: true,
+		},
+		{
+			input: CompositeKey().
+				AddPartitionKey(Key("pk1").GreaterThanEqual(Value("1"))).
+				AddPartitionKey(Key("pk1").LessThanEqual(Value("2"))),
+			expected: Expression{
+				expressionMap: map[expressionType]string{
+					keyCondition: "((#0 >= :0) AND (#0 <= :1))",
+				},
+				namesMap: map[string]string{
+					"#0": "pk1",
+				},
+				valuesMap: map[string]types.AttributeValue{
+					":0": &types.AttributeValueMemberS{Value: "1"},
+					":1": &types.AttributeValueMemberS{Value: "2"},
+				},
+			},
+		},
+	}
+
+	for idx, c := range cases {
+		t.Run(strconv.Itoa(idx), func(t *testing.T) {
+			if !c.input.IsSet() {
+				t.Errorf("IsSet() is false")
+				t.Fail()
+			}
+
+			actual, err := NewBuilder().WithKeyCondition(c.input).Build()
+
+			if c.expectError && err == nil {
+				t.Error("expected error")
+				t.Fail()
+			}
+			if !c.expectError && err != nil {
+				t.Errorf("unexpected error: %v", err)
+				t.Fail()
+			}
+
+			if c.expectError && err != nil {
+				t.Logf("found expected error: %v", err)
+				return
+			}
+
+			if !reflect.DeepEqual(c.expected, actual) {
+				t.Errorf("unexpted diff: %v vs %v", c.expected, actual)
+			}
+		})
+	}
+}

feature/dynamodb/expression/key_condition.go

@@ -29,6 +29,13 @@ const (
 	betweenKeyCond
 	// beginsWithKeyCond represents the Begins With KeyCondition
 	beginsWithKeyCond
+
+	// compositeKeyCond represents a composite key condition (for multiple partition/sort keys)
+	compositeKeyCond
+	// compositeKeyCondPartition represents a partition key within a composite key condition
+	compositeKeyCondPartition
+	// compositeKeyCondSort represents a sort key within a composite key condition
+	compositeKeyCondSort
 )
 
 // KeyConditionBuilder represents Key Condition Expressions in DynamoDB.
@@ -38,6 +45,7 @@ type KeyConditionBuilder struct {
 	operandList      []OperandBuilder
 	keyConditionList []KeyConditionBuilder
 	mode             keyConditionMode
+	compositeKeyMode keyConditionMode
 }
 
 // KeyEqual returns a KeyConditionBuilder representing the equality clause
@@ -357,6 +365,9 @@ func KeyAnd(left, right KeyConditionBuilder) KeyConditionBuilder {
 //	// ExpressionAttributeValues representing the item attribute "Number",
 //	// the value "Wildcats", and the value 1
 //	"(TeamName = :teamName) AND (#NUMBER = :one)"
+//
+// WARNING: This method is not supported for composite keys (created with CompositeKey()).
+// Calling And() on a composite key will fail. Use AddPartitionKey() and AddSortKey() instead.
 func (kcb KeyConditionBuilder) And(right KeyConditionBuilder) KeyConditionBuilder {
 	return KeyAnd(kcb, right)
 }
@@ -486,6 +497,8 @@ func (kcb KeyConditionBuilder) buildTree() (exprNode, error) {
 		return exprNode{}, newUnsetParameterError("buildTree", "KeyConditionBuilder")
 	case invalidKeyCond:
 		return exprNode{}, fmt.Errorf("buildKeyCondition error: invalid key condition constructed")
+	case compositeKeyCond:
+		return compositeKeyCondition(kcb, ret)
 	default:
 		return exprNode{}, fmt.Errorf("buildKeyCondition error: unsupported mode: %v", kcb.mode)
 	}