diff --git a/.changelog/01fe7c45-628d-4d98-a55b-e07d304448ae.json b/.changelog/01fe7c45-628d-4d98-a55b-e07d304448ae.json new file mode 100644 index 00000000000..432d16d477e --- /dev/null +++ b/.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" + ] +} diff --git a/feature/dynamodb/expression/composite_key.go b/feature/dynamodb/expression/composite_key.go new file mode 100644 index 00000000000..630c8c13bca --- /dev/null +++ b/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 +} diff --git a/feature/dynamodb/expression/composite_key_test.go b/feature/dynamodb/expression/composite_key_test.go new file mode 100644 index 00000000000..d1626ea350f --- /dev/null +++ b/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) + } + }) + } +} diff --git a/feature/dynamodb/expression/key_condition.go b/feature/dynamodb/expression/key_condition.go index ad52671fbfc..b67e618dfb4 100644 --- a/feature/dynamodb/expression/key_condition.go +++ b/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) }