Adds operation name labels for operations to identify calling op when required (aws#2846)

cenedhryn · web-flow · commit 6a11ab8196cc · 2021-11-18T18:20:19.000-08:00
diff --git a/.changes/next-release/feature-DynamoDBEnhancedClient-f030740.json b/.changes/next-release/feature-DynamoDBEnhancedClient-f030740.json
@@ -0,0 +1,6 @@
+{
+    "category": "DynamoDB Enhanced Client", 
+    "contributor": "", 
+    "type": "feature", 
+    "description": "Adds operation type labels to each operation and propagates to the extension write context"
+}
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/DynamoDbExtensionContext.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/DynamoDbExtensionContext.java
@@ -17,6 +17,7 @@
 
 import java.util.Map;
 import software.amazon.awssdk.annotations.SdkPublicApi;
+import software.amazon.awssdk.enhanced.dynamodb.internal.operations.OperationName;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 
 /**
@@ -57,6 +58,11 @@ default TableSchema<?> tableSchema() {
      */
     @SdkPublicApi
     public interface BeforeWrite extends Context {
+
+        /**
+         * @return The context under which the operation to be modified is taking place.
+         */
+        OperationName operationName();
     }
 
     /**
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/extensions/ChainExtension.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/extensions/ChainExtension.java
@@ -104,6 +104,7 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
                                                .operationContext(context.operationContext())
                                                .tableMetadata(context.tableMetadata())
                                                .tableSchema(context.tableSchema())
+                                               .operationName(context.operationName())
                                                .build();
 
             WriteModification writeModification = extension.beforeWrite(beforeWrite);
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/extensions/DefaultDynamoDbExtensionContext.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/extensions/DefaultDynamoDbExtensionContext.java
@@ -22,6 +22,7 @@
 import software.amazon.awssdk.enhanced.dynamodb.OperationContext;
 import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
 import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
+import software.amazon.awssdk.enhanced.dynamodb.internal.operations.OperationName;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 
 /**
@@ -35,12 +36,14 @@ public final class DefaultDynamoDbExtensionContext implements DynamoDbExtensionC
     private final OperationContext operationContext;
     private final TableMetadata tableMetadata;
     private final TableSchema<?> tableSchema;
+    private final OperationName operationName;
 
     private DefaultDynamoDbExtensionContext(Builder builder) {
         this.items = builder.items;
         this.operationContext = builder.operationContext;
         this.tableMetadata = builder.tableMetadata;
         this.tableSchema = builder.tableSchema;
+        this.operationName = builder.operationName != null ? builder.operationName : OperationName.NONE;
     }
 
     public static Builder builder() {
@@ -67,6 +70,11 @@ public TableSchema<?> tableSchema() {
         return tableSchema;
     }
 
+    @Override
+    public OperationName operationName() {
+        return operationName;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) {
@@ -87,7 +95,10 @@ public boolean equals(Object o) {
         if (!Objects.equals(tableMetadata, that.tableMetadata)) {
             return false;
         }
-        return Objects.equals(tableSchema, that.tableSchema);
+        if (!Objects.equals(tableSchema, that.tableSchema)) {
+            return false;
+        }
+        return Objects.equals(operationName, that.operationName);
     }
 
     @Override
@@ -96,6 +107,7 @@ public int hashCode() {
         result = 31 * result + (operationContext != null ? operationContext.hashCode() : 0);
         result = 31 * result + (tableMetadata != null ? tableMetadata.hashCode() : 0);
         result = 31 * result + (tableSchema != null ? tableSchema.hashCode() : 0);
+        result = 31 * result + (operationName != null ? operationName.hashCode() : 0);
         return result;
     }
 
@@ -104,6 +116,7 @@ public static final class Builder {
         private OperationContext operationContext;
         private TableMetadata tableMetadata;
         private TableSchema<?> tableSchema;
+        private OperationName operationName;
 
         public Builder items(Map<String, AttributeValue> item) {
             this.items = item;
@@ -125,6 +138,11 @@ public Builder tableSchema(TableSchema<?> tableSchema) {
             return this;
         }
 
+        public Builder operationName(OperationName operationName) {
+            this.operationName = operationName;
+            return this;
+        }
+
         public DefaultDynamoDbExtensionContext build() {
             return new DefaultDynamoDbExtensionContext(this);
         }
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/BatchGetItemOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/BatchGetItemOperation.java
@@ -50,6 +50,11 @@ public static BatchGetItemOperation create(BatchGetItemEnhancedRequest request)
         return new BatchGetItemOperation(request);
     }
 
+    @Override
+    public OperationName operationName() {
+        return OperationName.BATCH_GET_ITEM;
+    }
+
     @Override
     public BatchGetItemRequest generateRequest(DynamoDbEnhancedClientExtension extension) {
         Map<String, KeysAndAttributes> requestItems = new HashMap<>();
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/BatchWriteItemOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/BatchWriteItemOperation.java
@@ -48,6 +48,11 @@ public static BatchWriteItemOperation create(BatchWriteItemEnhancedRequest reque
         return new BatchWriteItemOperation(request);
     }
 
+    @Override
+    public OperationName operationName() {
+        return OperationName.BATCH_WRITE_ITEM;
+    }
+
     @Override
     public BatchWriteItemRequest generateRequest(DynamoDbEnhancedClientExtension extension) {
         Map<String, List<WriteRequest>> allRequestItems = new HashMap<>();
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/CommonOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/CommonOperation.java
@@ -140,4 +140,9 @@ default CompletableFuture<ResultT> executeAsync(TableSchema<ItemT> tableSchema,
         CompletableFuture<ResponseT> response = asyncServiceCall(dynamoDbAsyncClient).apply(request);
         return response.thenApply(r -> transformResponse(r, tableSchema, context, extension));
     }
+
+    /**
+     * The type, or name, of the operation.
+     */
+    OperationName operationName();
 }
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/CreateTableOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/CreateTableOperation.java
@@ -53,6 +53,11 @@ public static <T> CreateTableOperation<T> create(CreateTableEnhancedRequest requ
         return new CreateTableOperation<>(request);
     }
 
+    @Override
+    public OperationName operationName() {
+        return OperationName.CREATE_TABLE;
+    }
+
     @Override
     public CreateTableRequest generateRequest(TableSchema<T> tableSchema,
                                               OperationContext operationContext,
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/DatabaseOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/DatabaseOperation.java
@@ -103,4 +103,10 @@ default CompletableFuture<ResultT> executeAsync(DynamoDbAsyncClient dynamoDbAsyn
         CompletableFuture<ResponseT> response = asyncServiceCall(dynamoDbAsyncClient).apply(request);
         return response.thenApply(r -> transformResponse(r, extension));
     }
+
+
+    /**
+     * The type, or name, of the operation.
+     */
+    OperationName operationName();
 }
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/DeleteItemOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/DeleteItemOperation.java
@@ -66,6 +66,11 @@ public static <T> DeleteItemOperation<T> create(TransactDeleteItemEnhancedReques
         return new DeleteItemOperation<>(request);
     }
 
+    @Override
+    public OperationName operationName() {
+        return OperationName.DELETE_ITEM;
+    }
+
     @Override
     public DeleteItemRequest generateRequest(TableSchema<T> tableSchema,
                                              OperationContext operationContext,
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/DeleteTableOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/DeleteTableOperation.java
@@ -36,6 +36,11 @@ public static <T> DeleteTableOperation<T> create() {
         return new DeleteTableOperation<>();
     }
 
+    @Override
+    public OperationName operationName() {
+        return OperationName.DELETE_ITEM;
+    }
+
     @Override
     public DeleteTableRequest generateRequest(TableSchema<T> tableSchema,
                                               OperationContext operationContext,
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/DescribeTableOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/DescribeTableOperation.java
@@ -36,6 +36,11 @@ public static <T> DescribeTableOperation<T> create() {
         return new DescribeTableOperation<>();
     }
 
+    @Override
+    public OperationName operationName() {
+        return OperationName.DESCRIBE_TABLE;
+    }
+
     @Override
     public DescribeTableRequest generateRequest(TableSchema<T> tableSchema,
                                                 OperationContext operationContext,
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/GetItemOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/GetItemOperation.java
@@ -57,6 +57,11 @@ public Key key() {
         return this.request.key();
     }
 
+    @Override
+    public OperationName operationName() {
+        return OperationName.GET_ITEM;
+    }
+
     @Override
     public GetItemRequest generateRequest(TableSchema<T> tableSchema,
                                           OperationContext context,
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/OperationName.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/OperationName.java
@@ -0,0 +1,50 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.enhanced.dynamodb.internal.operations;
+
+import software.amazon.awssdk.annotations.SdkInternalApi;
+
+@SdkInternalApi
+public enum OperationName {
+    NONE(null),
+    BATCH_GET_ITEM("BatchGetItem"),
+    BATCH_WRITE_ITEM("BatchWriteItem"),
+    CREATE_TABLE("CreateTable"),
+    DELETE_ITEM("DeleteItem"),
+    DELETE_TABLE("DeleteTable"),
+    DESCRIBE_TABLE("DescribeTable"),
+    GET_ITEM("GetItem"),
+    QUERY("Query"),
+    PUT_ITEM("PutItem"),
+    SCAN("Scan"),
+    TRANSACT_GET_ITEMS("TransactGetItems"),
+    TRANSACT_WRITE_ITEMS("TransactWriteItems"),
+    UPDATE_ITEM("UpdateItem");
+
+    private final String label;
+
+    OperationName() {
+        this.label = null;
+    }
+
+    OperationName(String label) {
+        this.label = label;
+    }
+
+    public String label() {
+        return label;
+    }
+}
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/PaginatedDatabaseOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/PaginatedDatabaseOperation.java
@@ -110,4 +110,9 @@ default SdkPublisher<ResultT> executeAsync(DynamoDbAsyncClient dynamoDbAsyncClie
         SdkPublisher<ResponseT> response = asyncServiceCall(dynamoDbAsyncClient).apply(request);
         return response.map(r -> transformResponse(r, extension));
     }
+
+    /**
+     * The type, or name, of the operation.
+     */
+    OperationName operationName();
 }
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/PaginatedTableOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/PaginatedTableOperation.java
@@ -82,4 +82,9 @@ default PagePublisher<ItemT> executeOnPrimaryIndexAsync(TableSchema<ItemT> table
         OperationContext context = DefaultOperationContext.create(tableName, TableMetadata.primaryIndexName());
         return executeAsync(tableSchema, context, extension, dynamoDbAsyncClient);
     }
+
+    /**
+     * The type, or name, of the operation.
+     */
+    OperationName operationName();
 }
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/PutItemOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/PutItemOperation.java
@@ -66,6 +66,11 @@ public static <T> PutItemOperation<T> create(TransactPutItemEnhancedRequest<T> r
         return new PutItemOperation<>(request);
     }
 
+    @Override
+    public OperationName operationName() {
+        return OperationName.PUT_ITEM;
+    }
+
     @Override
     public PutItemRequest generateRequest(TableSchema<T> tableSchema,
                                           OperationContext operationContext,
@@ -91,6 +96,7 @@ public PutItemRequest generateRequest(TableSchema<T> tableSchema,
                                                .operationContext(operationContext)
                                                .tableMetadata(tableMetadata)
                                                .tableSchema(tableSchema)
+                                               .operationName(operationName())
                                                .build())
                 : null;
 
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/QueryOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/QueryOperation.java
@@ -49,6 +49,11 @@ public static <T> QueryOperation<T> create(QueryEnhancedRequest request) {
         return new QueryOperation<>(request);
     }
 
+    @Override
+    public OperationName operationName() {
+        return OperationName.QUERY;
+    }
+
     @Override
     public QueryRequest generateRequest(TableSchema<T> tableSchema,
                                         OperationContext operationContext,
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/ScanOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/ScanOperation.java
@@ -50,6 +50,11 @@ public static <T> ScanOperation<T> create(ScanEnhancedRequest request) {
         return new ScanOperation<>(request);
     }
 
+    @Override
+    public OperationName operationName() {
+        return OperationName.SCAN;
+    }
+
     @Override
     public ScanRequest generateRequest(TableSchema<T> tableSchema,
                                        OperationContext operationContext,
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/TransactGetItemsOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/TransactGetItemsOperation.java
@@ -43,6 +43,11 @@ public static TransactGetItemsOperation create(TransactGetItemsEnhancedRequest r
         return new TransactGetItemsOperation(request);
     }
 
+    @Override
+    public OperationName operationName() {
+        return OperationName.TRANSACT_GET_ITEMS;
+    }
+
     @Override
     public TransactGetItemsRequest generateRequest(DynamoDbEnhancedClientExtension extension) {
         return TransactGetItemsRequest.builder()
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/TransactWriteItemsOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/TransactWriteItemsOperation.java
@@ -39,6 +39,11 @@ public static TransactWriteItemsOperation create(TransactWriteItemsEnhancedReque
         return new TransactWriteItemsOperation(request);
     }
 
+    @Override
+    public OperationName operationName() {
+        return OperationName.TRANSACT_WRITE_ITEMS;
+    }
+
     @Override
     public TransactWriteItemsRequest generateRequest(DynamoDbEnhancedClientExtension extension) {
         return TransactWriteItemsRequest.builder()
diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/UpdateItemOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/UpdateItemOperation.java
@@ -84,6 +84,11 @@ public static <T> UpdateItemOperation<T> create(TransactUpdateItemEnhancedReques
         return new UpdateItemOperation<>(request);
     }
 
+    @Override
+    public OperationName operationName() {
+        return OperationName.UPDATE_ITEM;
+    }
+
     @Override
     public UpdateItemRequest generateRequest(TableSchema<T> tableSchema,
                                              OperationContext operationContext,
@@ -103,11 +108,12 @@ public UpdateItemRequest generateRequest(TableSchema<T> tableSchema,
         WriteModification transformation =
             extension != null
             ? extension.beforeWrite(DefaultDynamoDbExtensionContext.builder()
-                                                                                     .items(itemMap)
-                                                                                     .operationContext(operationContext)
-                                                                                     .tableMetadata(tableMetadata)
-                                                                                     .tableSchema(tableSchema)
-                                                                                     .build())
+                                                                   .items(itemMap)
+                                                                   .operationContext(operationContext)
+                                                                   .tableMetadata(tableMetadata)
+                                                                   .tableSchema(tableSchema)
+                                                                   .operationName(operationName())
+                                                                   .build())
             : null;
 
         if (transformation != null && transformation.transformedItem() != null) {
diff --git a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/extensions/ChainExtensionTest.java b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/extensions/ChainExtensionTest.java
diff --git a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/PutItemWithResponseTest.java b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/PutItemWithResponseTest.java
diff --git a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/extensions/DefaultDynamoDbExtensionContextTest.java b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/extensions/DefaultDynamoDbExtensionContextTest.java
diff --git a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/IndexOperationTest.java b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/IndexOperationTest.java
diff --git a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/PutItemOperationTest.java b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/PutItemOperationTest.java
diff --git a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/TableOperationTest.java b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/TableOperationTest.java
diff --git a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/UpdateItemOperationTest.java b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/UpdateItemOperationTest.java

Original file line number	Diff line number	Diff line change
`@@ -140,4 +140,9 @@ default CompletableFuture<ResultT> executeAsync(TableSchema<ItemT> tableSchema,`
`140`	`140`	`CompletableFuture<ResponseT> response = asyncServiceCall(dynamoDbAsyncClient).apply(request);`
`141`	`141`	`return response.thenApply(r -> transformResponse(r, tableSchema, context, extension));`
`142`	`142`	`}`
	`143`	`+`
	`144`	`+ /**`
	`145`	`+ * The type, or name, of the operation.`
	`146`	`+ */`
	`147`	`+ OperationName operationName();`
`143`	`148`	`}`
Original file line number	Diff line number	Diff line change
`@@ -103,4 +103,10 @@ default CompletableFuture<ResultT> executeAsync(DynamoDbAsyncClient dynamoDbAsyn`
`103`	`103`	`CompletableFuture<ResponseT> response = asyncServiceCall(dynamoDbAsyncClient).apply(request);`
`104`	`104`	`return response.thenApply(r -> transformResponse(r, extension));`
`105`	`105`	`}`
	`106`	`+`
	`107`	`+`
	`108`	`+ /**`
	`109`	`+ * The type, or name, of the operation.`
	`110`	`+ */`
	`111`	`+ OperationName operationName();`
`106`	`112`	`}`