auto commit

Author: rishav-karanjit

Examples/runtimes/java/DynamoDbEncryption/src/main/java/software/amazon/cryptography/examples/BasicPutGetExample.java

@@ -26,15 +26,18 @@
   is provided in CLI arguments.
   This table must be configured with the following
   primary key configuration:
-    - Partition key is named "partition_key" with type (S)
-    - Sort key is named "sort_key" with type (N)
+    - Partition key is named `partitionKeyName` with type (S)
+    - Sort key is named `sortKeyName` with type (N)
  */
 public class BasicPutGetExample {
 
   public static void PutItemGetItem(
-    String kmsKeyId,
-    String ddbTableName,
-    String PartitionKeyName
+    final String kmsKeyId,
+    final String ddbTableName,
+    final String partitionKeyName,
+    final String sortKeyName,
+    final String partitionKeyValue,
+    final String sortKeyValue
   ) {
     // 1. Create a Keyring. This Keyring will be responsible for protecting the data keys that protect your data.
     //    For this example, we will create a AWS KMS Keyring with the AWS KMS Key we want to use.
@@ -57,8 +60,8 @@ public static void PutItemGetItem(
     //      - SIGN_ONLY: The attribute not encrypted, but is still included in the signature
     //      - DO_NOTHING: The attribute is not encrypted and not included in the signature
     final Map<String, CryptoAction> attributeActionsOnEncrypt = new HashMap<>();
-    attributeActionsOnEncrypt.put("partition_key", CryptoAction.SIGN_ONLY); // Our partition attribute must be SIGN_ONLY
-    attributeActionsOnEncrypt.put("sort_key", CryptoAction.SIGN_ONLY); // Our sort attribute must be SIGN_ONLY
+    attributeActionsOnEncrypt.put(partitionKeyName, CryptoAction.SIGN_ONLY); // Our partition attribute must be SIGN_ONLY
+    attributeActionsOnEncrypt.put(sortKeyName, CryptoAction.SIGN_ONLY); // Our sort attribute must be SIGN_ONLY
     attributeActionsOnEncrypt.put("attribute1", CryptoAction.ENCRYPT_AND_SIGN);
     attributeActionsOnEncrypt.put("attribute2", CryptoAction.SIGN_ONLY);
     attributeActionsOnEncrypt.put(":attribute3", CryptoAction.DO_NOTHING);
@@ -99,8 +102,8 @@ public static void PutItemGetItem(
     final DynamoDbTableEncryptionConfig config = DynamoDbTableEncryptionConfig
       .builder()
       .logicalTableName(ddbTableName)
-      .partitionKeyName("partition_key")
-      .sortKeyName("sort_key")
+      .partitionKeyName(partitionKeyName)
+      .sortKeyName(sortKeyName)
       .attributeActionsOnEncrypt(attributeActionsOnEncrypt)
       .keyring(kmsKeyring)
       .allowedUnsignedAttributePrefix(unsignAttrPrefix)
@@ -146,10 +149,10 @@ public static void PutItemGetItem(
     //    client-side, according to our configuration.
     final HashMap<String, AttributeValue> item = new HashMap<>();
     item.put(
-      "partition_key",
-      AttributeValue.builder().s(PartitionKeyName).build()
+      partitionKeyName,
+      AttributeValue.builder().s(partitionKeyValue).build()
     );
-    item.put("sort_key", AttributeValue.builder().n("0").build());
+    item.put(sortKeyName, AttributeValue.builder().n(sortKeyValue).build());
     item.put(
       "attribute1",
       AttributeValue.builder().s("encrypt and sign me!").build()
@@ -173,10 +176,10 @@ public static void PutItemGetItem(
     //    back the original item.
     final HashMap<String, AttributeValue> keyToGet = new HashMap<>();
     keyToGet.put(
-      "partition_key",
-      AttributeValue.builder().s(PartitionKeyName).build()
+      partitionKeyName,
+      AttributeValue.builder().s(partitionKeyValue).build()
     );
-    keyToGet.put("sort_key", AttributeValue.builder().n("0").build());
+    keyToGet.put(sortKeyName, AttributeValue.builder().n(sortKeyValue).build());
 
     final GetItemRequest getRequest = GetItemRequest
       .builder()
@@ -199,13 +202,19 @@ public static void PutItemGetItem(
   }
 
   public static void main(final String[] args) {
-    if (args.length < 2) {
+    if (args.length < 6) {
       throw new IllegalArgumentException(
-        "To run this example, include the kmsKeyId as args[0] and ddbTableName as args[1]"
+        "To run this example, include the kmsKeyId as args[0], ddbTableName as args[1],"
+        + " partitionKeyName as args[2], sortKeyName as args[3], partitionKeyValue as args[4]"
+        + " sortKeyValue as args[5]"
       );
     }
     final String kmsKeyId = args[0];
     final String ddbTableName = args[1];
-    PutItemGetItem(kmsKeyId, ddbTableName, "BasicPutGetExample");
+    final String partitionKeyName = args[2];
+    final String sortKeyName = args[3];
+    final String partitionKeyValue = args[4];
+    final String sortKeyValue = args[5];
+    PutItemGetItem(kmsKeyId, ddbTableName, partitionKeyName, sortKeyName, partitionKeyValue, sortKeyValue);
   }
 }

Examples/runtimes/java/DynamoDbEncryption/src/test/java/software/amazon/cryptography/examples/TestBasicPutGetExample.java

@@ -7,10 +7,21 @@ public class TestBasicPutGetExample {
 
   @Test
   public void TestPutGet() {
+    final String partitionKeyValue = "BasicPutGetExample" + UUID.randomUUID();
     BasicPutGetExample.PutItemGetItem(
       TestUtils.TEST_KMS_KEY_ID,
       TestUtils.TEST_DDB_TABLE_NAME,
-      "BasicPutGetExample" + UUID.randomUUID()
+      "partition_key",
+      "sort_key",
+      partitionKeyValue,
+      "0"
+    );
+    TestUtils.cleanUpDDBItem(
+      TestUtils.TEST_DDB_TABLE_NAME,
+      "partition_key",
+      "sort_key",
+      partitionKeyValue,
+      "0"
     );
   }
 }

Examples/runtimes/java/DynamoDbEncryption/src/test/java/software/amazon/cryptography/examples/TestUtils.java

@@ -1,5 +1,10 @@
 package software.amazon.cryptography.examples;
 
+import java.util.HashMap;
+import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
+import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
+import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
+
 public class TestUtils {
 
   public static final String TEST_KEYSTORE_NAME = "KeyStoreDdbTable";
@@ -30,4 +35,40 @@ public class TestUtils {
   // Our tests require access to DDB Table with this name
   public static final String TEST_DDB_TABLE_NAME =
     "DynamoDbEncryptionInterceptorTestTable";
+
+  /**
+   * Deletes an item from a DynamoDB table.
+   *
+   * @param tableName The name of the DynamoDB table
+   * @param partitionKeyName The name of partition key
+   * @param sortKeyName The name of sort key
+   * @param partitionKeyValue The value of the partition key
+   * @param sortKeyValue The value of the sort key (can be null if table doesn't have a sort key)
+   */
+  public static void cleanUpDDBItem(
+    final String tableName,
+    final String partitionKeyName,
+    final String sortKeyName,
+    final String partitionKeyValue,
+    final String sortKeyValue
+  ) {
+    final DynamoDbClient ddb = DynamoDbClient.builder().build();
+    final HashMap<String, AttributeValue> keyToDelete = new HashMap<>();
+    keyToDelete.put(
+      partitionKeyName,
+      AttributeValue.builder().s(partitionKeyValue).build()
+    );
+    if (sortKeyValue != null) {
+      keyToDelete.put(
+        sortKeyName,
+        AttributeValue.builder().n(sortKeyValue).build()
+      );
+    }
+    final DeleteItemRequest deleteRequest = DeleteItemRequest
+      .builder()
+      .tableName(tableName)
+      .key(keyToDelete)
+      .build();
+    ddb.deleteItem(deleteRequest);
+  }
 }